"""Tests for myagents.commands.ollama and myagents.ollama_adapter.""" from __future__ import annotations import json import subprocess import threading from contextlib import contextmanager from http.client import HTTPConnection from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer import click import pytest import myagents.ollama_adapter as adapter_mod from myagents.commands import ollama from myagents.ollama_adapter import ( _Handler, is_adapter_endpoint, listen_url, normalize_system, parse_num_ctx, runtime_num_ctx, upstream_target, ) class TestNormalizeSystem: def test_moves_system_message_to_top_level(self) -> None: body = { "system": [{"type": "text", "text": "SYS1"}], "messages": [ {"role": "user", "content": [{"type": "text", "text": "hi"}]}, { "role": "system", "content": [{"type": "text", "text": "SYS2"}], }, ], } out = normalize_system(body) assert [m["role"] for m in out["messages"]] == ["user"] assert [b["text"] for b in out["system"]] == ["SYS1", "SYS2"] def test_keeps_non_system_order(self) -> None: body = { "messages": [ {"role": "user", "content": "a"}, { "role": "system", "content": [{"type": "text", "text": "s"}], }, {"role": "user", "content": "b"}, ] } out = normalize_system(body) assert [m["role"] for m in out["messages"]] == ["user", "user"] assert [m["content"] for m in out["messages"]] == ["a", "b"] def test_string_system_becomes_block(self) -> None: body = {"system": "top", "messages": [{"role": "user", "content": "a"}]} out = normalize_system(body) assert out["system"] == [{"type": "text", "text": "top"}] def test_hoists_string_content_on_system_message(self) -> None: body = { "messages": [ {"role": "user", "content": "hi"}, {"role": "system", "content": "agent-types"}, {"role": "user", "content": "ping"}, ] } out = normalize_system(body) assert [m["role"] for m in out["messages"]] == ["user", "user"] assert out["system"] == [{"type": "text", "text": "agent-types"}] def test_preserves_cache_control_on_text_blocks(self) -> None: body = { "system": [ { "type": "text", "text": "SYS", "cache_control": {"type": "ephemeral"}, } ], "messages": [ { "role": "system", "content": [ { "type": "text", "text": "more", "cache_control": {"type": "ephemeral"}, } ], } ], } out = normalize_system(body) assert out["system"][0]["cache_control"] == {"type": "ephemeral"} assert out["system"][1] == { "type": "text", "text": "more", "cache_control": {"type": "ephemeral"}, } assert out["messages"] == [] def test_drops_empty_system_turns(self) -> None: body = { "messages": [ {"role": "system", "content": ""}, {"role": "user", "content": "hi"}, ] } out = normalize_system(body) assert out["messages"] == [{"role": "user", "content": "hi"}] assert out["system"] == [] class TestListenAndUpstream: def test_listen_url(self) -> None: assert listen_url() == ( f"http://{adapter_mod.DEFAULT_HOST}:{adapter_mod.DEFAULT_PORT}" ) def test_is_adapter_endpoint_loopback_aliases(self) -> None: port = adapter_mod.DEFAULT_PORT assert is_adapter_endpoint(f"http://127.0.0.1:{port}") assert is_adapter_endpoint(f"http://127.0.0.1:{port}/") assert is_adapter_endpoint(f"http://localhost:{port}") assert not is_adapter_endpoint("http://127.0.0.1:11434") assert not is_adapter_endpoint("") assert not is_adapter_endpoint("https://api.deepseek.com/anthropic") def test_upstream_target_defaults_to_ollama_port(self) -> None: assert upstream_target("http://127.0.0.1") == ( "http", "127.0.0.1", 11434, ) assert upstream_target("127.0.0.1:11434") == ( "http", "127.0.0.1", 11434, ) assert upstream_target("https://example.com") == ( "https", "example.com", 443, ) assert upstream_target("https://example.com:8443") == ( "https", "example.com", 8443, ) class TestRuntimeNumCtx: def test_parse_num_ctx(self) -> None: assert parse_num_ctx("num_ctx 32768\n") == 32768 assert parse_num_ctx("temperature 1.0\nnum_ctx 65536") == 65536 assert parse_num_ctx(None) is None assert parse_num_ctx("") is None assert parse_num_ctx("temperature 1.0") is None def test_runtime_num_ctx_reads_show(self, monkeypatch) -> None: class _Resp: def read(self) -> bytes: return b'{"parameters":"num_ctx 65536\\n"}' def __enter__(self): return self def __exit__(self, *args: object) -> None: return None monkeypatch.setattr(adapter_mod, "urlopen", lambda *a, **k: _Resp()) assert runtime_num_ctx("qwen3.5:4b-ctx64k") == 65536 def test_runtime_num_ctx_missing_model(self) -> None: assert runtime_num_ctx("") is None def _serve(handler: type[BaseHTTPRequestHandler]) -> ThreadingHTTPServer: httpd = ThreadingHTTPServer(("127.0.0.1", 0), handler) threading.Thread(target=httpd.serve_forever, daemon=True).start() return httpd @contextmanager def _running(handler: type[BaseHTTPRequestHandler]): httpd = _serve(handler) try: yield httpd finally: httpd.shutdown() httpd.server_close() @contextmanager def _adapter_against(upstream_handler, monkeypatch): with _running(upstream_handler) as upstream: monkeypatch.setattr( adapter_mod, "OLLAMA_BASE", f"http://127.0.0.1:{upstream.server_address[1]}", ) with _running(_Handler) as proxy: yield proxy def _recorder( store: dict, *, status: int = 200, body: bytes = b'{"ok":true}', send_length: bool = True, ): class Handler(BaseHTTPRequestHandler): protocol_version = "HTTP/1.1" def _capture(self) -> None: n = int(self.headers.get("content-length") or 0) raw = self.rfile.read(n) if n else b"" store["method"] = self.command store["path"] = self.path store["body"] = raw store["headers"] = {k.lower(): v for k, v in self.headers.items()} def _reply(self, head: bool = False) -> None: self.send_response(status) self.send_header("content-type", "application/json") if send_length: self.send_header("content-length", str(len(body))) else: self.send_header("connection", "close") self.close_connection = True self.end_headers() if not head: self.wfile.write(body) def do_GET(self) -> None: self._capture() self._reply() def do_POST(self) -> None: self._capture() self._reply() def do_HEAD(self) -> None: self._capture() self._reply(head=True) def log_message(self, format: str, *args: object) -> None: pass return Handler def _request( proxy: ThreadingHTTPServer, method: str, path: str, body: bytes = b"", headers: dict[str, str] | None = None, ) -> tuple[int, bytes, str | None]: host, port = proxy.server_address conn = HTTPConnection(host, port, timeout=2) hdrs = dict(headers or {}) if body and "content-type" not in {k.lower() for k in hdrs}: hdrs["content-type"] = "application/json" conn.request(method, path, body=body or None, headers=hdrs) resp = conn.getresponse() raw = resp.read() conn.close() return resp.status, raw, resp.getheader("content-length") class TestProxy: def test_health_is_local(self, monkeypatch) -> None: store: dict = {} with _adapter_against(_recorder(store), monkeypatch) as proxy: status, raw, _length = _request(proxy, "GET", "/health") assert status == 200 assert json.loads(raw)["status"] == "ok" assert store == {} def test_get_models_is_forwarded(self, monkeypatch) -> None: store: dict = {} payload = b'{"object":"list","data":[]}' with _adapter_against( _recorder(store, body=payload), monkeypatch ) as proxy: status, raw, _length = _request(proxy, "GET", "/v1/models") assert status == 200 assert raw == payload assert store["method"] == "GET" assert store["path"] == "/v1/models" def test_post_messages_hoists_system_and_keeps_query( self, monkeypatch ) -> None: store: dict = {} payload = { "model": "x", "max_tokens": 1, "messages": [ {"role": "user", "content": "hi"}, {"role": "system", "content": "agent-types"}, {"role": "user", "content": "ping"}, ], } with _adapter_against(_recorder(store), monkeypatch) as proxy: status, _raw, _length = _request( proxy, "POST", "/v1/messages?beta=true", json.dumps(payload).encode(), ) assert status == 200 assert store["path"] == "/v1/messages?beta=true" forwarded = json.loads(store["body"]) assert [m["role"] for m in forwarded["messages"]] == ["user", "user"] assert forwarded["system"] == [ {"type": "text", "text": "agent-types"} ] def test_other_post_paths_are_forwarded_unnormalized( self, monkeypatch ) -> None: store: dict = {} payload = b'{"model":"x"}' with _adapter_against(_recorder(store), monkeypatch) as proxy: status, _raw, _length = _request( proxy, "POST", "/v1/messages/count_tokens?beta=true", payload, ) assert status == 200 assert store["path"] == "/v1/messages/count_tokens?beta=true" assert store["body"] == payload def test_forwards_client_api_key(self, monkeypatch) -> None: store: dict = {} with _adapter_against(_recorder(store), monkeypatch) as proxy: _request( proxy, "POST", "/v1/messages", b'{"messages":[]}', headers={"x-api-key": "from-client"}, ) assert store["headers"]["x-api-key"] == "from-client" def test_upstream_error_is_502(self, monkeypatch) -> None: monkeypatch.setattr(adapter_mod, "OLLAMA_BASE", "http://127.0.0.1:1") with _running(_Handler) as proxy: status, raw, _length = _request( proxy, "POST", "/v1/messages", b"{}" ) assert status == 502 assert "upstream error" in json.loads(raw)["error"]["message"] class TestForwardTerminates: def test_completes_when_upstream_sends_content_length( self, monkeypatch ) -> None: store: dict = {} body = b'{"id":"msg_1","type":"message"}' with _adapter_against( _recorder(store, body=body), monkeypatch ) as proxy: status, raw, length = _request( proxy, "POST", "/v1/messages", b'{"messages":[]}' ) assert status == 200 assert json.loads(raw)["id"] == "msg_1" assert length == str(len(raw)) def test_completes_when_upstream_omits_content_length( self, monkeypatch ) -> None: store: dict = {} body = b'data: {"type":"message_stop"}\n\n' with _adapter_against( _recorder(store, body=body, send_length=False), monkeypatch ) as proxy: status, raw, _length = _request( proxy, "POST", "/v1/messages?beta=true", b'{"messages":[]}' ) assert status == 200 assert b"message_stop" in raw def test_forwards_upstream_4xx_and_completes(self, monkeypatch) -> None: store: dict = {} body = b'{"error":{"message":"bad request"}}' with _adapter_against( _recorder(store, status=400, body=body), monkeypatch ) as proxy: status, raw, length = _request( proxy, "POST", "/v1/messages", b'{"messages":[]}' ) assert status == 400 assert json.loads(raw)["error"]["message"] == "bad request" assert length == str(len(raw)) class TestSystemdUnit: def test_unit_text_smoke(self) -> None: text = ollama._sysd_unit_text() assert "ExecStart=" in text assert "-m myagents.ollama_adapter" in text assert "WantedBy=default.target" in text def test_sysd_start_enables_and_reloads( self, monkeypatch, tmp_path ) -> None: calls: list[list[str]] = [] def fake_run(cmd, **kwargs): kwargs.pop("capture_output", None) kwargs.pop("text", None) calls.append(list(cmd)) return subprocess.CompletedProcess(cmd, 0, stdout="") monkeypatch.setattr(ollama.subprocess, "run", fake_run) monkeypatch.setattr(ollama, "UNIT_DIR", tmp_path) monkeypatch.setattr(ollama, "UNIT_PATH", tmp_path / ollama.UNIT_NAME) ollama._sysd_start() assert ["systemctl", "--user", "daemon-reload"] in calls assert [ "systemctl", "--user", "enable", "--now", ollama.UNIT_NAME, ] in calls assert any(c[0] == "loginctl" and "enable-linger" in c for c in calls) assert (tmp_path / ollama.UNIT_NAME).exists() assert ["systemctl", "--user", "restart", ollama.UNIT_NAME] not in calls def test_sysd_start_restarts_when_already_active( self, monkeypatch, tmp_path ) -> None: calls: list[list[str]] = [] def fake_run(cmd, **kwargs): kwargs.pop("capture_output", None) kwargs.pop("text", None) calls.append(list(cmd)) if cmd[:3] == ["systemctl", "--user", "is-active"]: return subprocess.CompletedProcess( cmd, 0, stdout="active\n", stderr="" ) if cmd[:3] == ["systemctl", "--user", "is-enabled"]: return subprocess.CompletedProcess( cmd, 0, stdout="enabled\n", stderr="" ) return subprocess.CompletedProcess(cmd, 0, stdout="", stderr="") monkeypatch.setattr(ollama.subprocess, "run", fake_run) monkeypatch.setattr(ollama, "UNIT_DIR", tmp_path) monkeypatch.setattr(ollama, "UNIT_PATH", tmp_path / ollama.UNIT_NAME) ollama._sysd_start() assert ["systemctl", "--user", "restart", ollama.UNIT_NAME] in calls def test_sysd_stop_disables(self, monkeypatch) -> None: calls: list[list[str]] = [] def fake_run(cmd, **kwargs): kwargs.pop("capture_output", None) kwargs.pop("text", None) calls.append(list(cmd)) return subprocess.CompletedProcess(cmd, 0) monkeypatch.setattr(ollama.subprocess, "run", fake_run) ollama._sysd_stop() assert [ "systemctl", "--user", "disable", "--now", ollama.UNIT_NAME, ] in calls assert ["systemctl", "--user", "daemon-reload"] in calls class TestMacLaunchd: def test_mac_start_bootstraps_when_not_loaded(self, monkeypatch) -> None: calls: list[list[str]] = [] monkeypatch.setattr(ollama, "_mac_loaded", lambda: False) def fake_run(cmd, **kwargs): calls.append(list(cmd)) return subprocess.CompletedProcess(cmd, 0, stdout="", stderr="") monkeypatch.setattr(ollama.subprocess, "run", fake_run) ollama._mac_start() assert calls[0][:2] == ["launchctl", "bootstrap"] assert not any(c[1] == "bootout" for c in calls) def test_mac_start_rebounds_when_loaded(self, monkeypatch) -> None: loaded = {"v": True} calls: list[list[str]] = [] monkeypatch.setattr(ollama, "_mac_loaded", lambda: loaded["v"]) def fake_run(cmd, **kwargs): if len(cmd) > 1 and cmd[1] == "bootout": loaded["v"] = False calls.append(list(cmd)) return subprocess.CompletedProcess(cmd, 0, stdout="", stderr="") monkeypatch.setattr(ollama.subprocess, "run", fake_run) ollama._mac_start() assert calls[0][:2] == ["launchctl", "bootout"] assert calls[1][:2] == ["launchctl", "bootstrap"] class TestBackendDispatch: def test_darwin_launchd(self) -> None: if ollama._platform() == "darwin": assert ollama._backend().name == "launchd" def test_linux_systemd(self, monkeypatch) -> None: monkeypatch.setattr(ollama, "_platform", lambda: "linux") assert ollama._backend().name == "systemd" def test_unsupported_platform_raises(self, monkeypatch) -> None: monkeypatch.setattr(ollama, "_platform", lambda: "win32") with pytest.raises(click.ClickException): ollama._backend()