wehub-resource-sync
106 lines
3.7 KiB
Python
106 lines
3.7 KiB
Python
from unittest import mock
|
|
|
|
import pytest
|
|
|
|
from mlflow.gateway.providers.utils import (
|
|
SUPPORTED_ACCEPT_ENCODING,
|
|
_aiohttp_post,
|
|
proxy_root_url,
|
|
rename_payload_keys,
|
|
)
|
|
|
|
from tests.gateway.tools import MockAsyncResponse, mock_http_client
|
|
|
|
|
|
def test_rename_payload_keys():
|
|
payload = {"old_key1": "value1", "old_key2": "value2", "old_key3": None, "old_key4": []}
|
|
mapping = {"old_key1": "new_key1", "old_key2": "new_key2"}
|
|
expected = {"new_key1": "value1", "new_key2": "value2", "old_key3": None, "old_key4": []}
|
|
assert rename_payload_keys(payload, mapping) == expected
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("payload", "mapping", "expected"),
|
|
[
|
|
(
|
|
{"old_key1": "value1", "old_key2": None, "old_key3": "value3"},
|
|
{"old_key1": "new_key1", "old_key3": "new_key3"},
|
|
{"new_key1": "value1", "old_key2": None, "new_key3": "value3"},
|
|
),
|
|
(
|
|
{"old_key1": None, "old_key2": "value2", "old_key3": []},
|
|
{"old_key1": "new_key1", "old_key3": "new_key3"},
|
|
{"new_key1": None, "old_key2": "value2", "new_key3": []},
|
|
),
|
|
(
|
|
{"old_key1": "value1", "old_key2": "value2"},
|
|
{"old_key1": "new_key1", "old_key3": "new_key3"},
|
|
{"new_key1": "value1", "old_key2": "value2"},
|
|
),
|
|
],
|
|
)
|
|
def test_rename_payload_keys_parameterized(payload, mapping, expected):
|
|
assert rename_payload_keys(payload, mapping) == expected
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_aiohttp_post_includes_supported_accept_encoding():
|
|
mock_client = mock_http_client(MockAsyncResponse({}))
|
|
with mock.patch("aiohttp.ClientSession", return_value=mock_client) as mock_session_cls:
|
|
async with _aiohttp_post(
|
|
headers={"Authorization": "Bearer key"},
|
|
base_url="https://api.example.com",
|
|
path="/v1/chat",
|
|
payload={"model": "x"},
|
|
):
|
|
pass
|
|
mock_session_cls.assert_called_once()
|
|
call_headers = mock_session_cls.call_args.kwargs["headers"]
|
|
assert call_headers.get("Accept-Encoding") == SUPPORTED_ACCEPT_ENCODING
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_aiohttp_post_uses_timeout_from_env_var(monkeypatch):
|
|
monkeypatch.setenv("MLFLOW_GATEWAY_ROUTE_TIMEOUT_SECONDS", "2")
|
|
|
|
mock_client = mock_http_client(MockAsyncResponse({}))
|
|
with mock.patch("aiohttp.ClientSession", return_value=mock_client):
|
|
async with _aiohttp_post(
|
|
headers={"Authorization": "Bearer key"},
|
|
base_url="https://api.example.com",
|
|
path="/v1/chat",
|
|
payload={"model": "x"},
|
|
):
|
|
pass
|
|
|
|
mock_client.post.assert_called_once()
|
|
assert mock_client.post.call_args.kwargs["timeout"].total == 2
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_aiohttp_post_sets_read_bufsize_for_large_sse_lines():
|
|
mock_client = mock_http_client(MockAsyncResponse({}))
|
|
with mock.patch("aiohttp.ClientSession", return_value=mock_client) as mock_session_cls:
|
|
async with _aiohttp_post(
|
|
headers={"Authorization": "Bearer key"},
|
|
base_url="https://api.example.com",
|
|
path="/v1/chat",
|
|
payload={"model": "x"},
|
|
):
|
|
pass
|
|
|
|
assert mock_session_cls.call_args.kwargs["read_bufsize"] == 2**20
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("base_url", "expected"),
|
|
[
|
|
("https://api.anthropic.com/v1", "https://api.anthropic.com"),
|
|
("https://api.openai.com/v1", "https://api.openai.com"),
|
|
("https://api.groq.com/openai/v1", "https://api.groq.com/openai"),
|
|
("https://api.example.com/v1/", "https://api.example.com"),
|
|
],
|
|
)
|
|
def test_proxy_root_url(base_url, expected):
|
|
assert proxy_root_url(base_url) == expected
|