wehub-resource-sync
254 lines
8.6 KiB
Python
254 lines
8.6 KiB
Python
import socket
|
|
import subprocess
|
|
import sys
|
|
from unittest import mock
|
|
|
|
import pytest
|
|
from requests.adapters import HTTPAdapter
|
|
|
|
from mlflow.utils import request_utils
|
|
from mlflow.utils.request_utils import (
|
|
TCPKeepAliveHTTPAdapter,
|
|
_build_socket_options,
|
|
)
|
|
|
|
|
|
def test_request_utils_does_not_import_mlflow(tmp_path):
|
|
file_content = f"""
|
|
import importlib.util
|
|
import os
|
|
import sys
|
|
|
|
file_path = r"{request_utils.__file__}"
|
|
module_name = "mlflow.utils.request_utils"
|
|
|
|
spec = importlib.util.spec_from_file_location(module_name, file_path)
|
|
module = importlib.util.module_from_spec(spec)
|
|
sys.modules[module_name] = module
|
|
spec.loader.exec_module(module)
|
|
|
|
assert "mlflow" not in sys.modules
|
|
assert "mlflow.utils.request_utils" in sys.modules
|
|
"""
|
|
test_file = tmp_path.joinpath("test_request_utils_does_not_import_mlflow.py")
|
|
test_file.write_text(file_content)
|
|
|
|
subprocess.check_call([sys.executable, str(test_file)])
|
|
|
|
|
|
class IncompleteResponse:
|
|
def __init__(self):
|
|
self.headers = {"Content-Length": "100"}
|
|
raw = mock.MagicMock()
|
|
raw.tell.return_value = 50
|
|
self.raw = raw
|
|
|
|
def __enter__(self):
|
|
return self
|
|
|
|
def __exit__(self, *args):
|
|
pass
|
|
|
|
|
|
def test_download_chunk_incomplete_read(tmp_path):
|
|
with mock.patch.object(
|
|
request_utils, "cloud_storage_http_request", return_value=IncompleteResponse()
|
|
):
|
|
download_path = tmp_path / "chunk"
|
|
download_path.touch()
|
|
with pytest.raises(IOError, match="Incomplete read"):
|
|
request_utils.download_chunk(
|
|
range_start=0,
|
|
range_end=999,
|
|
headers={},
|
|
download_path=download_path,
|
|
http_uri="https://example.com",
|
|
)
|
|
|
|
|
|
@pytest.mark.parametrize("env_value", ["0", "false", "False", "FALSE"])
|
|
def test_redirects_disabled_if_env_var_set(monkeypatch, env_value):
|
|
monkeypatch.setenv("MLFLOW_ALLOW_HTTP_REDIRECTS", env_value)
|
|
|
|
with mock.patch("requests.Session.request") as mock_request:
|
|
mock_request.return_value.status_code = 302
|
|
mock_request.return_value.text = "mock response"
|
|
|
|
response = request_utils.cloud_storage_http_request("GET", "http://localhost:5000")
|
|
|
|
assert response.text == "mock response"
|
|
mock_request.assert_called_once_with(
|
|
"GET",
|
|
"http://localhost:5000",
|
|
allow_redirects=False,
|
|
timeout=None,
|
|
)
|
|
|
|
|
|
@pytest.mark.parametrize("env_value", ["1", "true", "True", "TRUE"])
|
|
def test_redirects_enabled_if_env_var_set(monkeypatch, env_value):
|
|
monkeypatch.setenv("MLFLOW_ALLOW_HTTP_REDIRECTS", env_value)
|
|
|
|
with mock.patch("requests.Session.request") as mock_request:
|
|
mock_request.return_value.status_code = 302
|
|
mock_request.return_value.text = "mock response"
|
|
|
|
response = request_utils.cloud_storage_http_request(
|
|
"GET",
|
|
"http://localhost:5000",
|
|
)
|
|
|
|
assert response.text == "mock response"
|
|
mock_request.assert_called_once_with(
|
|
"GET",
|
|
"http://localhost:5000",
|
|
allow_redirects=True,
|
|
timeout=None,
|
|
)
|
|
|
|
|
|
@pytest.mark.parametrize("env_value", ["0", "false", "False", "FALSE"])
|
|
def test_redirect_kwarg_overrides_env_value_false(monkeypatch, env_value):
|
|
monkeypatch.setenv("MLFLOW_ALLOW_HTTP_REDIRECTS", env_value)
|
|
|
|
with mock.patch("requests.Session.request") as mock_request:
|
|
mock_request.return_value.status_code = 302
|
|
mock_request.return_value.text = "mock response"
|
|
|
|
response = request_utils.cloud_storage_http_request(
|
|
"GET", "http://localhost:5000", allow_redirects=True
|
|
)
|
|
|
|
assert response.text == "mock response"
|
|
mock_request.assert_called_once_with(
|
|
"GET",
|
|
"http://localhost:5000",
|
|
allow_redirects=True,
|
|
timeout=None,
|
|
)
|
|
|
|
|
|
@pytest.mark.parametrize("env_value", ["1", "true", "True", "TRUE"])
|
|
def test_redirect_kwarg_overrides_env_value_true(monkeypatch, env_value):
|
|
monkeypatch.setenv("MLFLOW_ALLOW_HTTP_REDIRECTS", env_value)
|
|
|
|
with mock.patch("requests.Session.request") as mock_request:
|
|
mock_request.return_value.status_code = 302
|
|
mock_request.return_value.text = "mock response"
|
|
|
|
response = request_utils.cloud_storage_http_request(
|
|
"GET", "http://localhost:5000", allow_redirects=False
|
|
)
|
|
|
|
assert response.text == "mock response"
|
|
mock_request.assert_called_once_with(
|
|
"GET",
|
|
"http://localhost:5000",
|
|
allow_redirects=False,
|
|
timeout=None,
|
|
)
|
|
|
|
|
|
def test_redirects_enabled_by_default():
|
|
with mock.patch("requests.Session.request") as mock_request:
|
|
mock_request.return_value.status_code = 302
|
|
mock_request.return_value.text = "mock response"
|
|
|
|
response = request_utils.cloud_storage_http_request(
|
|
"GET",
|
|
"http://localhost:5000",
|
|
)
|
|
|
|
assert response.text == "mock response"
|
|
mock_request.assert_called_once_with(
|
|
"GET",
|
|
"http://localhost:5000",
|
|
allow_redirects=True,
|
|
timeout=None,
|
|
)
|
|
|
|
|
|
# --- TCP Keepalive tests ---
|
|
|
|
|
|
def test_build_socket_options_includes_keepalive():
|
|
options = _build_socket_options()
|
|
assert (socket.SOL_SOCKET, socket.SO_KEEPALIVE, 1) in options
|
|
|
|
|
|
def test_build_socket_options_platform_specific():
|
|
options = _build_socket_options()
|
|
if hasattr(socket, "TCP_KEEPIDLE"):
|
|
assert (socket.IPPROTO_TCP, socket.TCP_KEEPIDLE, 30) in options
|
|
elif hasattr(socket, "TCP_KEEPALIVE"):
|
|
assert (socket.IPPROTO_TCP, socket.TCP_KEEPALIVE, 30) in options
|
|
if hasattr(socket, "TCP_KEEPINTVL"):
|
|
assert (socket.IPPROTO_TCP, socket.TCP_KEEPINTVL, 10) in options
|
|
if hasattr(socket, "TCP_KEEPCNT"):
|
|
assert (socket.IPPROTO_TCP, socket.TCP_KEEPCNT, 3) in options
|
|
|
|
|
|
def test_build_socket_options_disabled_via_env(monkeypatch):
|
|
monkeypatch.setenv("MLFLOW_HTTP_TCP_KEEPALIVE", "false")
|
|
options = _build_socket_options()
|
|
assert (socket.SOL_SOCKET, socket.SO_KEEPALIVE, 1) not in options
|
|
|
|
|
|
def test_build_socket_options_custom_values_via_env(monkeypatch):
|
|
monkeypatch.setenv("MLFLOW_HTTP_TCP_KEEPALIVE_IDLE", "60")
|
|
monkeypatch.setenv("MLFLOW_HTTP_TCP_KEEPALIVE_INTERVAL", "20")
|
|
monkeypatch.setenv("MLFLOW_HTTP_TCP_KEEPALIVE_COUNT", "5")
|
|
options = _build_socket_options()
|
|
if hasattr(socket, "TCP_KEEPIDLE"):
|
|
assert (socket.IPPROTO_TCP, socket.TCP_KEEPIDLE, 60) in options
|
|
elif hasattr(socket, "TCP_KEEPALIVE"):
|
|
assert (socket.IPPROTO_TCP, socket.TCP_KEEPALIVE, 60) in options
|
|
if hasattr(socket, "TCP_KEEPINTVL"):
|
|
assert (socket.IPPROTO_TCP, socket.TCP_KEEPINTVL, 20) in options
|
|
if hasattr(socket, "TCP_KEEPCNT"):
|
|
assert (socket.IPPROTO_TCP, socket.TCP_KEEPCNT, 5) in options
|
|
|
|
|
|
def test_tcp_keepalive_adapter_init_poolmanager():
|
|
adapter = TCPKeepAliveHTTPAdapter()
|
|
with mock.patch.object(HTTPAdapter, "init_poolmanager") as mock_init:
|
|
adapter.init_poolmanager(1, 1)
|
|
mock_init.assert_called_once()
|
|
_, kwargs = mock_init.call_args
|
|
assert "socket_options" in kwargs
|
|
assert (socket.SOL_SOCKET, socket.SO_KEEPALIVE, 1) in kwargs["socket_options"]
|
|
|
|
|
|
def test_tcp_keepalive_adapter_proxy_manager_for():
|
|
adapter = TCPKeepAliveHTTPAdapter()
|
|
with mock.patch.object(HTTPAdapter, "proxy_manager_for") as mock_proxy:
|
|
adapter.proxy_manager_for("http://proxy:8080")
|
|
mock_proxy.assert_called_once()
|
|
_, kwargs = mock_proxy.call_args
|
|
assert "socket_options" in kwargs
|
|
assert (socket.SOL_SOCKET, socket.SO_KEEPALIVE, 1) in kwargs["socket_options"]
|
|
|
|
|
|
def test_tcp_keepalive_adapter_proxy_respects_explicit_options():
|
|
adapter = TCPKeepAliveHTTPAdapter()
|
|
custom_options = [(socket.SOL_SOCKET, socket.SO_KEEPALIVE, 0)]
|
|
with mock.patch.object(HTTPAdapter, "proxy_manager_for") as mock_proxy:
|
|
adapter.proxy_manager_for("http://proxy:8080", socket_options=custom_options)
|
|
_, kwargs = mock_proxy.call_args
|
|
assert kwargs["socket_options"] == custom_options
|
|
|
|
|
|
def test_session_uses_tcp_keepalive_adapter():
|
|
request_utils._cached_get_request_session.cache_clear()
|
|
session = request_utils._get_request_session(
|
|
max_retries=3,
|
|
backoff_factor=1,
|
|
backoff_jitter=0.5,
|
|
retry_codes=(500,),
|
|
raise_on_status=True,
|
|
respect_retry_after_header=True,
|
|
)
|
|
assert isinstance(session.get_adapter("https://example.com"), TCPKeepAliveHTTPAdapter)
|
|
assert isinstance(session.get_adapter("http://example.com"), TCPKeepAliveHTTPAdapter)
|
|
request_utils._cached_get_request_session.cache_clear()
|