Initial commit

tests/test_remote_wrapper.py

@@ -0,0 +1,283 @@
+import asyncio
+import json
+import os
+import ssl
+import time
+from typing import Dict, List
+
+import pytest
+import requests
+import websockets
+
+WRAPPER_BASE = os.getenv("WRAPPER_BASE", "http://192.168.1.2:9093")
+UI_BASE = os.getenv("UI_BASE", "http://192.168.1.2:9094")
+TRUENAS_WS_URL = os.getenv("TRUENAS_WS_URL", "wss://192.168.1.2/websocket")
+TRUENAS_API_KEY = os.getenv("TRUENAS_API_KEY", "")
+TRUENAS_APP_NAME = os.getenv("TRUENAS_APP_NAME", "llamacpp")
+MODEL_REQUEST = os.getenv("MODEL_REQUEST", "")
+
+
+async def _rpc_call(method: str, params: List | None = None):
+    if not TRUENAS_API_KEY:
+        pytest.skip("TRUENAS_API_KEY not set")
+    ssl_ctx = ssl.create_default_context()
+    ssl_ctx.check_hostname = False
+    ssl_ctx.verify_mode = ssl.CERT_NONE
+    async with websockets.connect(TRUENAS_WS_URL, ssl=ssl_ctx) as ws:
+        await ws.send(json.dumps({"msg": "connect", "version": "1", "support": ["1"]}))
+        connected = json.loads(await ws.recv())
+        if connected.get("msg") != "connected":
+            raise RuntimeError("failed to connect")
+        await ws.send(json.dumps({"id": 1, "msg": "method", "method": "auth.login_with_api_key", "params": [TRUENAS_API_KEY]}))
+        auth = json.loads(await ws.recv())
+        if not auth.get("result"):
+            raise RuntimeError("auth failed")
+        await ws.send(json.dumps({"id": 2, "msg": "method", "method": method, "params": params or []}))
+        while True:
+            raw = json.loads(await ws.recv())
+            if raw.get("id") != 2:
+                continue
+            if raw.get("msg") == "error":
+                raise RuntimeError(raw.get("error"))
+            return raw.get("result")
+
+
+def _get_models() -> List[str]:
+    _wait_for_http(WRAPPER_BASE + "/health")
+    resp = requests.get(WRAPPER_BASE + "/v1/models", timeout=30)
+    resp.raise_for_status()
+    data = resp.json().get("data") or []
+    return [m.get("id") for m in data if m.get("id")]
+
+
+def _assert_chat_ok(resp_json: Dict) -> str:
+    choices = resp_json.get("choices") or []
+    assert choices, "no choices"
+    message = choices[0].get("message") or {}
+    text = message.get("content") or ""
+    assert text.strip(), "empty content"
+    return text
+
+
+def _wait_for_http(url: str, timeout_s: float = 90) -> None:
+    deadline = time.time() + timeout_s
+    last_err = None
+    while time.time() < deadline:
+        try:
+            resp = requests.get(url, timeout=5)
+            if resp.status_code == 200:
+                return
+            last_err = f"status {resp.status_code}"
+        except Exception as exc:
+            last_err = str(exc)
+        time.sleep(2)
+    raise RuntimeError(f"service not ready: {url} ({last_err})")
+
+
+def _post_with_retry(url: str, payload: Dict, timeout_s: float = 300, retries: int = 6, delay_s: float = 5.0):
+    last = None
+    for _ in range(retries):
+        try:
+            resp = requests.post(url, json=payload, timeout=timeout_s)
+            if resp.status_code == 200:
+                return resp
+            last = resp
+        except requests.exceptions.RequestException as exc:
+            last = exc
+        time.sleep(delay_s)
+    if isinstance(last, Exception):
+        raise last
+    return last
+
+
+@pytest.mark.asyncio
+async def test_active_model_and_multi_gpu_flags():
+    cfg = await _rpc_call("app.config", [TRUENAS_APP_NAME])
+    command = cfg.get("command") or []
+    assert "--model" in command
+    assert "--tensor-split" in command
+    split_idx = command.index("--tensor-split") + 1
+    split = command[split_idx]
+    assert "," in split, f"tensor-split missing commas: {split}"
+    assert "--split-mode" in command
+
+
+def test_models_listed():
+    models = _get_models()
+    assert models, "no models discovered"
+
+
+def test_chat_completions_switch_and_prompts():
+    models = _get_models()
+    assert models, "no models"
+    if MODEL_REQUEST:
+        assert MODEL_REQUEST in models, f"MODEL_REQUEST not found: {MODEL_REQUEST}"
+        model_id = MODEL_REQUEST
+    else:
+        model_id = models[0]
+    payload = {
+        "model": model_id,
+        "messages": [{"role": "user", "content": "Say OK."}],
+        "max_tokens": 12,
+        "temperature": 0,
+    }
+    for _ in range(3):
+        resp = _post_with_retry(WRAPPER_BASE + "/v1/chat/completions", payload)
+        assert resp.status_code == 200
+        _assert_chat_ok(resp.json())
+
+
+def test_tools_flat_format():
+    models = _get_models()
+    assert models, "no models"
+    if MODEL_REQUEST:
+        assert MODEL_REQUEST in models, f"MODEL_REQUEST not found: {MODEL_REQUEST}"
+        model_id = MODEL_REQUEST
+    else:
+        model_id = models[0]
+    payload = {
+        "model": model_id,
+        "messages": [{"role": "user", "content": "Say OK and do not call tools."}],
+        "tools": [
+            {
+                "type": "function",
+                "name": "format_final_json_response",
+                "description": "format output",
+                "parameters": {
+                    "type": "object",
+                    "properties": {"ok": {"type": "boolean"}},
+                    "required": ["ok"],
+                },
+            }
+        ],
+        "max_tokens": 12,
+    }
+    resp = _post_with_retry(WRAPPER_BASE + "/v1/chat/completions", payload)
+    assert resp.status_code == 200
+    _assert_chat_ok(resp.json())
+
+
+def test_functions_payload_normalized():
+    models = _get_models()
+    assert models, "no models"
+    if MODEL_REQUEST:
+        assert MODEL_REQUEST in models, f"MODEL_REQUEST not found: {MODEL_REQUEST}"
+        model_id = MODEL_REQUEST
+    else:
+        model_id = models[0]
+    payload = {
+        "model": model_id,
+        "messages": [{"role": "user", "content": "Say OK and do not call tools."}],
+        "functions": [
+            {
+                "name": "format_final_json_response",
+                "description": "format output",
+                "parameters": {
+                    "type": "object",
+                    "properties": {"ok": {"type": "boolean"}},
+                    "required": ["ok"],
+                },
+            }
+        ],
+        "max_tokens": 12,
+    }
+    resp = _post_with_retry(WRAPPER_BASE + "/v1/chat/completions", payload)
+    assert resp.status_code == 200
+    _assert_chat_ok(resp.json())
+
+
+def test_return_format_json():
+    models = _get_models()
+    assert models, "no models"
+    if MODEL_REQUEST:
+        assert MODEL_REQUEST in models, f"MODEL_REQUEST not found: {MODEL_REQUEST}"
+        model_id = MODEL_REQUEST
+    else:
+        model_id = models[0]
+    payload = {
+        "model": model_id,
+        "messages": [{"role": "user", "content": "Return JSON with key ok true."}],
+        "return_format": "json",
+        "max_tokens": 32,
+        "temperature": 0,
+    }
+    resp = _post_with_retry(WRAPPER_BASE + "/v1/chat/completions", payload)
+    assert resp.status_code == 200
+    text = _assert_chat_ok(resp.json())
+    parsed = json.loads(text)
+    assert isinstance(parsed, dict)
+
+
+def test_responses_endpoint():
+    models = _get_models()
+    assert models, "no models"
+    if MODEL_REQUEST:
+        assert MODEL_REQUEST in models, f"MODEL_REQUEST not found: {MODEL_REQUEST}"
+        model_id = MODEL_REQUEST
+    else:
+        model_id = models[0]
+    payload = {
+        "model": model_id,
+        "input": "Say OK.",
+        "max_output_tokens": 16,
+    }
+    resp = _post_with_retry(WRAPPER_BASE + "/v1/responses", payload)
+    assert resp.status_code == 200
+    output = resp.json().get("output") or []
+    assert output, "responses output empty"
+    content = output[0].get("content") or []
+    text = content[0].get("text") if content else ""
+    assert text and text.strip()
+
+
+@pytest.mark.asyncio
+async def test_model_switch_applied_to_truenas():
+    models = _get_models()
+    assert models, "no models"
+    target = MODEL_REQUEST or models[0]
+    assert target in models, f"MODEL_REQUEST not found: {target}"
+    resp = requests.post(UI_BASE + "/ui/api/switch-model", json={"model_id": target, "warmup_prompt": "warmup"}, timeout=600)
+    assert resp.status_code == 200
+    cfg = await _rpc_call("app.config", [TRUENAS_APP_NAME])
+    command = cfg.get("command") or []
+    assert "--model" in command
+    model_path = command[command.index("--model") + 1]
+    assert model_path.endswith(target)
+
+
+def test_invalid_model_rejected():
+    models = _get_models()
+    assert models, "no models"
+    payload = {
+        "model": "modelx-q8:4b",
+        "messages": [{"role": "user", "content": "Say OK."}],
+        "max_tokens": 8,
+        "temperature": 0,
+    }
+    resp = requests.post(WRAPPER_BASE + "/v1/chat/completions", json=payload, timeout=60)
+    assert resp.status_code == 404
+
+
+def test_llamacpp_logs_streaming():
+    logs = ""
+    for _ in range(5):
+        try:
+            resp = requests.get(UI_BASE + "/ui/api/llamacpp-logs", timeout=10)
+            if resp.status_code == 200:
+                logs = resp.json().get("logs") or ""
+                if logs.strip():
+                    break
+        except requests.exceptions.ReadTimeout:
+            pass
+        time.sleep(2)
+    assert logs.strip(), "no logs returned"
+
+    # Force a log line before streaming.
+    try:
+        requests.get(WRAPPER_BASE + "/proxy/llamacpp/health", timeout=5)
+    except Exception:
+        pass
+
+    # Stream endpoint may not emit immediately, so validate that the endpoint responds.
+    with requests.get(UI_BASE + "/ui/api/llamacpp-logs/stream", stream=True, timeout=(5, 5)) as resp:
+        assert resp.status_code == 200