55 lines
1.5 KiB
Python
55 lines
1.5 KiB
Python
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")
|