Initial commit

tests/test_truenas_switch.py

@@ -0,0 +1,54 @@
+import json
+import pytest
+
+from app.truenas_middleware import TrueNASConfig, switch_model
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("case", list(range(120)))
+async def test_switch_model_updates_command(monkeypatch, case):
+    compose = {
+        "services": {
+            "llamacpp": {
+                "command": [
+                    "--model",
+                    "/models/old.gguf",
+                    "--ctx-size",
+                    "2048",
+                ]
+            }
+        }
+    }
+
+    captured = {}
+
+    async def fake_rpc_call(cfg, method, params=None):
+        if method == "app.config":
+            return {"custom_compose_config": compose}
+        if method == "app.update":
+            captured["payload"] = params[1]
+            return {"state": "RUNNING"}
+        raise AssertionError(f"unexpected method {method}")
+
+    monkeypatch.setattr("app.truenas_middleware._rpc_call", fake_rpc_call)
+
+    cfg = TrueNASConfig(
+        ws_url="ws://truenas.test/websocket",
+        api_key="key",
+        api_user=None,
+        app_name="llamacpp",
+        verify_ssl=False,
+    )
+
+    await switch_model(
+        cfg,
+        f"/models/new-{case}.gguf",
+        {"n_gpu_layers": "999"},
+        "--flash-attn on",
+    )
+
+    assert "custom_compose_config" in captured["payload"]
+    cmd = captured["payload"]["custom_compose_config"]["services"]["llamacpp"]["command"]
+    assert "--model" in cmd
+    idx = cmd.index("--model")
+    assert cmd[idx + 1].endswith(f"new-{case}.gguf")