234 lines
7.6 KiB
Python
234 lines
7.6 KiB
Python
from __future__ import annotations
|
|
|
|
import base64
|
|
from io import BytesIO
|
|
import json
|
|
import sys
|
|
import types
|
|
|
|
import numpy as np
|
|
import pytest
|
|
from PIL import Image
|
|
|
|
import client
|
|
from helpers.fakes import FakeTensor, make_chat_completion, ns
|
|
|
|
|
|
def _decode_data_url_png(data_url: str) -> Image.Image:
|
|
assert data_url.startswith("data:image/png;base64,")
|
|
encoded = data_url.split(",", maxsplit=1)[1]
|
|
raw = base64.b64decode(encoded)
|
|
return Image.open(BytesIO(raw))
|
|
|
|
|
|
def test_normalize_server_url_rejects_invalid_values() -> None:
|
|
with pytest.raises(ValueError, match="must not be empty"):
|
|
client.normalize_server_url("")
|
|
with pytest.raises(ValueError, match="http:// or https://"):
|
|
client.normalize_server_url("ftp://127.0.0.1:1234")
|
|
with pytest.raises(ValueError, match="include host"):
|
|
client.normalize_server_url("http://")
|
|
|
|
|
|
def test_to_openai_base_url_appends_v1() -> None:
|
|
assert client.to_openai_base_url("http://127.0.0.1:1234") == "http://127.0.0.1:1234/v1"
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"value,expected",
|
|
[
|
|
(None, "lm-studio"),
|
|
("", "lm-studio"),
|
|
(" ", "lm-studio"),
|
|
("-", "lm-studio"),
|
|
(" token-1 ", "token-1"),
|
|
],
|
|
)
|
|
def test_normalize_api_key(value: str | None, expected: str) -> None:
|
|
assert client.normalize_api_key(value) == expected
|
|
|
|
|
|
def test_create_openai_client_passes_expected_args(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
captured: dict[str, object] = {}
|
|
|
|
class FakeOpenAI:
|
|
def __init__(self, **kwargs):
|
|
captured.update(kwargs)
|
|
|
|
fake_module = types.ModuleType("openai")
|
|
fake_module.OpenAI = FakeOpenAI
|
|
monkeypatch.setitem(sys.modules, "openai", fake_module)
|
|
|
|
created = client.create_openai_client(
|
|
server_url="http://localhost:1234",
|
|
api_key=" api ",
|
|
timeout_seconds=30,
|
|
max_retries=4,
|
|
)
|
|
|
|
assert isinstance(created, FakeOpenAI)
|
|
assert captured == {
|
|
"api_key": "api",
|
|
"base_url": "http://localhost:1234/v1",
|
|
"timeout": 30,
|
|
"max_retries": 4,
|
|
}
|
|
|
|
|
|
def test_list_models_deduplicates_and_sorts() -> None:
|
|
model_data = [
|
|
ns(id="model-b"),
|
|
ns(id="model-a"),
|
|
ns(id="model-b"),
|
|
ns(id=""),
|
|
ns(id=None),
|
|
ns(other="ignored"),
|
|
]
|
|
fake_client = ns(models=ns(list=lambda: ns(data=model_data)))
|
|
assert client.list_models(fake_client) == ["model-a", "model-b"]
|
|
|
|
|
|
def test_get_server_models_delegates(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
fake_client = object()
|
|
monkeypatch.setattr(client, "create_openai_client", lambda **_: fake_client)
|
|
monkeypatch.setattr(client, "list_models", lambda _c: ["x", "y"])
|
|
|
|
assert (
|
|
client.get_server_models(
|
|
server_url="http://localhost:1234",
|
|
api_key="-",
|
|
timeout_seconds=5,
|
|
)
|
|
== ["x", "y"]
|
|
)
|
|
|
|
|
|
def test_coerce_seed_numpy_bool_and_bytes_error() -> None:
|
|
assert client.coerce_seed(np.int64(12)) == 12
|
|
assert client.coerce_seed(True) == 1
|
|
assert client.coerce_seed(False) == 0
|
|
|
|
with pytest.raises(ValueError, match="utf-8"):
|
|
client.coerce_seed(b"\xff")
|
|
|
|
|
|
@pytest.mark.parametrize("value", [float("inf"), float("-inf")])
|
|
def test_coerce_seed_rejects_non_finite_float(value: float) -> None:
|
|
with pytest.raises(ValueError, match="finite"):
|
|
client.coerce_seed(value)
|
|
|
|
|
|
def test_resolve_request_seed_handles_ranges(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
monkeypatch.setattr(client.secrets, "randbelow", lambda _: 42)
|
|
assert client.resolve_request_seed(None) == 43
|
|
assert client.resolve_request_seed(-1) == 43
|
|
assert client.resolve_request_seed(11) == 11
|
|
|
|
with pytest.raises(ValueError, match=">= 0"):
|
|
client.resolve_request_seed(-2)
|
|
|
|
|
|
def test_comfy_image_to_base64_png_url_single_and_batch() -> None:
|
|
single = FakeTensor(np.array([[[1.0, 0.0, 0.0]]], dtype=np.float32))
|
|
url_single = client.comfy_image_to_base64_png_url(single)
|
|
image_single = _decode_data_url_png(url_single)
|
|
assert image_single.size == (1, 1)
|
|
assert image_single.getpixel((0, 0))[:3] == (255, 0, 0)
|
|
|
|
batch = FakeTensor(
|
|
np.array(
|
|
[
|
|
[[[1.0, 0.0, 0.0]]], # red
|
|
[[[0.0, 0.0, 1.0]]], # blue
|
|
],
|
|
dtype=np.float32,
|
|
)
|
|
)
|
|
url_batch = client.comfy_image_to_base64_png_url(batch)
|
|
image_batch = _decode_data_url_png(url_batch)
|
|
assert image_batch.getpixel((0, 0))[:3] == (255, 0, 0)
|
|
|
|
|
|
def test_build_responses_input_text_handles_none() -> None:
|
|
payload = client.build_responses_input_text(None)
|
|
assert payload == [{"role": "user", "content": [{"type": "input_text", "text": ""}]}]
|
|
|
|
|
|
def test_extract_chat_completion_text_variants() -> None:
|
|
assert client.extract_chat_completion_text(ns(choices=[])) == ""
|
|
assert client.extract_chat_completion_text(ns(choices=[ns(message=None)])) == ""
|
|
assert client.extract_chat_completion_text(make_chat_completion("plain")) == "plain"
|
|
|
|
content = [
|
|
"line1",
|
|
ns(text="line2"),
|
|
{"text": "line3"},
|
|
ns(other="ignored"),
|
|
{"other": "ignored"},
|
|
]
|
|
assert client.extract_chat_completion_text(make_chat_completion(content)) == "line1\nline2\nline3"
|
|
|
|
|
|
def test_extract_responses_text_variants() -> None:
|
|
assert client.extract_responses_text(ns(output_text="from-output-text")) == "from-output-text"
|
|
|
|
response_from_dict = ns(
|
|
output=[{"content": [{"type": "output_text", "text": "from-dict"}]}]
|
|
)
|
|
assert client.extract_responses_text(response_from_dict) == "from-dict"
|
|
|
|
response_from_object = ns(output=[ns(content=[ns(text="from-object-content")]), ns(text="from-object")])
|
|
assert client.extract_responses_text(response_from_object) == "from-object-content\nfrom-object"
|
|
|
|
assert client.extract_responses_text(ns(output_text="", output=[])) == ""
|
|
|
|
|
|
def test_dump_openai_response_branches() -> None:
|
|
class WithModelDump:
|
|
def model_dump(self):
|
|
return {"a": 1}
|
|
|
|
class WithToDict:
|
|
def to_dict(self):
|
|
return {"b": 2}
|
|
|
|
class Plain:
|
|
def __str__(self) -> str:
|
|
return "plain-response"
|
|
|
|
assert json.loads(client.dump_openai_response(WithModelDump())) == {"a": 1}
|
|
assert json.loads(client.dump_openai_response(WithToDict())) == {"b": 2}
|
|
assert json.loads(client.dump_openai_response({"c": 3})) == {"c": 3}
|
|
assert client.dump_openai_response(Plain()) == "plain-response"
|
|
|
|
|
|
def test_strip_think_content_case_insensitive_and_multiple() -> None:
|
|
text = "A\n<THINK>hidden-1</THINK>\nB\n<think>hidden-2</think>\nC"
|
|
assert client.strip_think_content(text) == "A\n\nB\n\nC"
|
|
assert client.strip_think_content("Visible\n<think>unterminated") == "Visible"
|
|
assert client.strip_think_content(None) == ""
|
|
|
|
|
|
def test_build_chat_messages_with_and_without_system_prompt() -> None:
|
|
with_system = client.build_chat_messages("System", "User")
|
|
assert with_system[0] == {"role": "system", "content": "System"}
|
|
assert with_system[1] == {"role": "user", "content": "User"}
|
|
|
|
no_system = client.build_chat_messages(" ", "User")
|
|
assert no_system == [{"role": "user", "content": "User"}]
|
|
|
|
|
|
def test_build_chat_messages_with_image_shape() -> None:
|
|
messages = client.build_chat_messages_with_image(
|
|
system_prompt="sys",
|
|
user_prompt="describe",
|
|
image_data_url="data:image/png;base64,abc",
|
|
)
|
|
assert messages[0] == {"role": "system", "content": "sys"}
|
|
assert messages[1]["role"] == "user"
|
|
assert messages[1]["content"][0] == {"type": "text", "text": "describe"}
|
|
assert messages[1]["content"][1] == {
|
|
"type": "image_url",
|
|
"image_url": {"url": "data:image/png;base64,abc"},
|
|
}
|