114 lines
4.9 KiB
Python
114 lines
4.9 KiB
Python
from __future__ import annotations
|
|
|
|
import importlib.util
|
|
import json
|
|
import threading
|
|
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
|
|
from pathlib import Path
|
|
from urllib.error import HTTPError
|
|
|
|
import pytest
|
|
|
|
spec = importlib.util.spec_from_file_location("modelforge_client", Path(__file__).resolve().parents[2] / "clients/python/modelforge_client.py")
|
|
assert spec and spec.loader
|
|
sdk = importlib.util.module_from_spec(spec)
|
|
spec.loader.exec_module(sdk)
|
|
|
|
|
|
def test_python_client_scope_typed_errors_and_no_redirect_or_retry() -> None:
|
|
seen = []
|
|
|
|
class Handler(BaseHTTPRequestHandler):
|
|
def do_GET(self):
|
|
seen.append((self.path, self.headers.get("Authorization")))
|
|
self.send_response(200); self.end_headers(); self.wfile.write(b'{"connected":true}')
|
|
|
|
def do_POST(self):
|
|
seen.append((self.path, self.headers.get("Authorization")))
|
|
body = json.loads(self.rfile.read(int(self.headers["Content-Length"])))
|
|
if body.get("redirect"):
|
|
self.send_response(307); self.send_header("Location", "/must-not-follow"); self.end_headers()
|
|
elif body.get("flat_error"):
|
|
self.send_response(503)
|
|
self.send_header("Content-Type", "application/json")
|
|
self.end_headers()
|
|
self.wfile.write(
|
|
b'{"code":"LEGACY_ERROR","message":"Legacy error","correlation_id":"legacy-case"}'
|
|
)
|
|
else:
|
|
self.send_response(503); self.send_header("Content-Type", "application/json"); self.end_headers()
|
|
self.wfile.write(
|
|
json.dumps(
|
|
{
|
|
"error": {
|
|
"code": "NO_ELIGIBLE_NODE",
|
|
"message": "No node",
|
|
"correlation_id": "sdk-case",
|
|
"details": {
|
|
"input": "private benchmark input that must remain redacted"
|
|
},
|
|
}
|
|
}
|
|
).encode()
|
|
)
|
|
|
|
def log_message(self, *_args):
|
|
pass
|
|
|
|
server = ThreadingHTTPServer(("127.0.0.1", 0), Handler)
|
|
worker = threading.Thread(target=server.serve_forever, daemon=True); worker.start()
|
|
token = "mf_" + "k" * 48
|
|
try:
|
|
client = sdk.ModelForgeClient(f"http://127.0.0.1:{server.server_port}", token, "rag.embedding@1", allow_http=True, timeout=2)
|
|
assert client.health() == {"connected": True}
|
|
with pytest.raises(sdk.ModelForgeError) as failure:
|
|
client.invoke({"input": ["test"]})
|
|
assert failure.value.code == "NO_ELIGIBLE_NODE"
|
|
assert str(failure.value) == "No node"
|
|
assert failure.value.status == 503 and failure.value.correlation_id == "sdk-case"
|
|
assert "private benchmark input" not in str(failure.value)
|
|
with pytest.raises(sdk.ModelForgeError) as legacy:
|
|
client.invoke({"flat_error": True})
|
|
assert legacy.value.code == "LEGACY_ERROR"
|
|
assert str(legacy.value) == "Legacy error"
|
|
assert legacy.value.correlation_id == "legacy-case"
|
|
with pytest.raises(sdk.ModelForgeError) as redirected:
|
|
client.invoke({"redirect": True})
|
|
assert redirected.value.status == 307
|
|
assert len(seen) == 4 and all("must-not-follow" not in path for path, _ in seen)
|
|
assert all(header == f"Bearer {token}" for _, header in seen)
|
|
assert token not in repr(client) + repr(failure.value)
|
|
finally:
|
|
server.shutdown(); server.server_close(); worker.join(timeout=3)
|
|
|
|
|
|
@pytest.mark.parametrize("origin", ["http://lab.invalid", "https://user:password@lab.invalid", "https://lab.invalid/path", "https://lab.invalid/?q=key"])
|
|
def test_python_client_rejects_unsafe_origin(origin: str) -> None:
|
|
with pytest.raises(ValueError):
|
|
sdk.ModelForgeClient(origin, "mf_" + "k" * 48, "rag.embedding@1")
|
|
|
|
|
|
def test_python_client_preserves_typed_error_when_error_body_times_out() -> None:
|
|
class StalledBody:
|
|
closed = False
|
|
|
|
def read(self, _limit):
|
|
raise TimeoutError("private transport details")
|
|
|
|
def close(self):
|
|
pass
|
|
|
|
class Opener:
|
|
def open(self, _request, *, timeout):
|
|
raise HTTPError("https://lab.invalid", 503, "Unavailable",
|
|
{"X-Correlation-ID": "stalled-error"}, StalledBody())
|
|
|
|
client = sdk.ModelForgeClient("https://lab.invalid", "mf_" + "k" * 48, "rag.embedding@1")
|
|
client._opener = Opener()
|
|
with pytest.raises(sdk.ModelForgeError) as failure:
|
|
client.health()
|
|
assert failure.value.code == "TRANSPORT_ERROR"
|
|
assert failure.value.status == 503
|
|
assert failure.value.correlation_id == "stalled-error"
|
|
assert "private transport details" not in str(failure.value)
|