v0.17.3: bug-fix maintenance pass — ingest origin exemption, FK safety, auto-search reset, deterministic ingest, WAL, config centralization, dead-code removal
- /api/ingest exempt from origin check for CLI/Bearer clients - recreate deleted conversations on bogus conversation_id (chat + search) - augmented auto-search emits reset:true; frontend clears first-pass text - image uploads stored as placeholders, never text-ingested - check_fact_conflicts requires shared subject keywords - deterministic ingest point ids (md5 chunk hash) - get_load() no longer crashes when rocm-smi yields no VRAM lines - SQLite WAL + busy_timeout; timing-safe API key compares - supervise fire-and-forget auto-ingest tasks; log missing logprobs - centralize EMBED_URL/EMBED_MODEL/QDRANT_URL/NODE_NAME in config - remove dead triage.py, select_node tests, is_state_changing, stray artifacts - add regression tests + autouse global-reset conftest
This commit is contained in:
@@ -4,3 +4,31 @@ import sys
|
||||
from pathlib import Path
|
||||
|
||||
sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
|
||||
|
||||
"""Shared pytest fixtures.
|
||||
|
||||
All test modules manipulate in-process globals (sessions, rate buckets,
|
||||
cluster registry, eviction log). An autouse fixture resets every global
|
||||
before each test so no state leaks between tests, regardless of whether an
|
||||
individual test file remembers to clear it.
|
||||
"""
|
||||
import pytest
|
||||
|
||||
import cluster
|
||||
import routers.chat
|
||||
from eviction import EVICTION_LOG
|
||||
from security import SESSIONS, PIN_ATTEMPTS, RATE_EVENTS
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _reset_global_state():
|
||||
SESSIONS.clear()
|
||||
PIN_ATTEMPTS.clear()
|
||||
RATE_EVENTS.clear()
|
||||
cluster.CLUSTER_NODES.clear()
|
||||
cluster.CLUSTER_EVENTS.clear()
|
||||
cluster.CLUSTER_COORDINATOR = None
|
||||
cluster._pending_pings.clear()
|
||||
EVICTION_LOG.clear()
|
||||
routers.chat._ingest_tasks.clear()
|
||||
yield
|
||||
|
||||
@@ -10,7 +10,6 @@ import app
|
||||
import config
|
||||
import db
|
||||
import routers.chat
|
||||
import triage
|
||||
from security import SESSIONS, PIN_ATTEMPTS, RATE_EVENTS
|
||||
|
||||
|
||||
@@ -269,7 +268,6 @@ def test_private_chat_does_not_persist(tmp_path: Path, monkeypatch):
|
||||
'data: {"choices":[{"delta":{"content":""},"finish_reason":"stop","logprobs":{"content":[]}}],"usage":{"completion_tokens":2,"prompt_tokens":10,"tokens_per_second":5.0}}',
|
||||
"data: [DONE]",
|
||||
]))
|
||||
monkeypatch.setattr(triage, "classify_query", lambda q: "general")
|
||||
|
||||
async def _mock_ensure(m): return True
|
||||
monkeypatch.setattr("model_pull.ensure_model", _mock_ensure)
|
||||
@@ -301,7 +299,6 @@ def test_private_chat_does_not_auto_search(tmp_path: Path, monkeypatch):
|
||||
'data: {"choices":[{"delta":{"content":""},"finish_reason":"stop","logprobs":{"content":[{"logprob":-2.5}]}}],"usage":{"completion_tokens":1,"prompt_tokens":10,"tokens_per_second":5.0}}',
|
||||
"data: [DONE]",
|
||||
]))
|
||||
monkeypatch.setattr(triage, "classify_query", lambda q: "general")
|
||||
monkeypatch.setattr(routers.chat, "query_searxng", lambda q: [{"title": "result"}])
|
||||
|
||||
with make_client(tmp_path) as client:
|
||||
|
||||
@@ -0,0 +1,357 @@
|
||||
"""Regression tests for bug fixes.
|
||||
|
||||
Covers: /api/ingest origin exemption for CLI/Bearer clients, bogus
|
||||
conversation_id FK handling in chat + search, the auto-search reset flag,
|
||||
image uploads being stored as placeholders instead of text-ingested,
|
||||
false-positive conflict detection, deterministic ingest point IDs,
|
||||
get_load() VRAM parsing, and version pinning.
|
||||
"""
|
||||
import json
|
||||
import os
|
||||
import re
|
||||
import subprocess
|
||||
from pathlib import Path
|
||||
|
||||
import httpx
|
||||
from fastapi.testclient import TestClient
|
||||
|
||||
import app
|
||||
import config
|
||||
import db
|
||||
import memory
|
||||
from crypto import decrypt_text
|
||||
import node_agent.agent as agent
|
||||
import routers.chat
|
||||
import routers.ingest as ingest_route
|
||||
import routers.search_route
|
||||
from security import SESSIONS, PIN_ATTEMPTS, RATE_EVENTS
|
||||
|
||||
|
||||
def make_client(tmp_path: Path) -> TestClient:
|
||||
os.environ["CAIC_ADMIN_PIN"] = "1234"
|
||||
db.DB_PATH = tmp_path / "caic-regression.db"
|
||||
SESSIONS.clear()
|
||||
PIN_ATTEMPTS.clear()
|
||||
RATE_EVENTS.clear()
|
||||
db.init_db()
|
||||
return TestClient(app.app, raise_server_exceptions=False)
|
||||
|
||||
|
||||
def _guest_headers(client: TestClient) -> dict:
|
||||
sid = client.post("/api/auth/guest", headers={"Origin": "http://testserver"}).json()["session_id"]
|
||||
return {"X-Session-ID": sid, "Origin": "http://testserver"}
|
||||
|
||||
|
||||
def _admin_headers(client: TestClient) -> dict:
|
||||
login = client.post("/api/auth/login", json={"pin": "1234"}, headers={"Origin": "http://testserver"})
|
||||
sid = login.json()["session_id"]
|
||||
return {"X-Session-ID": sid, "Origin": "http://testserver"}
|
||||
|
||||
|
||||
def parse_sse_payloads(body: str) -> list[dict]:
|
||||
payloads = []
|
||||
for chunk in body.split("\n\n"):
|
||||
chunk = chunk.strip()
|
||||
if not chunk.startswith("data: "):
|
||||
continue
|
||||
payloads.append(json.loads(chunk[len("data: "):]))
|
||||
return payloads
|
||||
|
||||
|
||||
class _MockStreamResponse:
|
||||
def __init__(self, lines: list[str]):
|
||||
self._lines = lines
|
||||
|
||||
async def __aenter__(self):
|
||||
return self
|
||||
|
||||
async def __aexit__(self, exc_type, exc, tb):
|
||||
return False
|
||||
|
||||
async def aiter_lines(self):
|
||||
for line in self._lines:
|
||||
yield line
|
||||
|
||||
|
||||
def _stream_json_lines(events: list[dict]) -> list[str]:
|
||||
return [json.dumps(event) for event in events]
|
||||
|
||||
|
||||
class _FakeAsyncClient:
|
||||
class FakeResponse:
|
||||
def __init__(self, status, json_data=None):
|
||||
self.status_code = status
|
||||
self._json = json_data or {}
|
||||
|
||||
def json(self):
|
||||
return self._json
|
||||
|
||||
def __init__(self, *a, **kw):
|
||||
self.put_payloads = []
|
||||
|
||||
async def __aenter__(self):
|
||||
return self
|
||||
|
||||
async def __aexit__(self, *a):
|
||||
pass
|
||||
|
||||
async def post(self, url, **kw):
|
||||
if "/api/embeddings" in url:
|
||||
return self.FakeResponse(200, {"embedding": [0.1] * 768})
|
||||
return self.FakeResponse(200)
|
||||
|
||||
async def put(self, url, **kw):
|
||||
self.put_payloads.append(kw.get("json", {}))
|
||||
return self.FakeResponse(200)
|
||||
|
||||
|
||||
# ── /api/ingest is reached by CLI tools with no Origin header ──────────
|
||||
|
||||
|
||||
def test_ingest_origin_exemption(tmp_path: Path, monkeypatch):
|
||||
monkeypatch.setattr(ingest_route, "COMPLETIONS_API_KEY", "sk-regression")
|
||||
monkeypatch.setattr(httpx, "AsyncClient", lambda *a, **kw: _FakeAsyncClient())
|
||||
with make_client(tmp_path) as client:
|
||||
resp = client.post(
|
||||
"/api/ingest",
|
||||
json={"content": "regression test content " * 20, "source": "cli"},
|
||||
headers={"Authorization": "Bearer sk-regression", "Content-Type": "application/json"},
|
||||
)
|
||||
assert resp.status_code == 200, resp.text
|
||||
assert resp.json()["source"] == "cli"
|
||||
|
||||
|
||||
def test_ingest_bad_key_still_blocked_without_origin(tmp_path: Path, monkeypatch):
|
||||
monkeypatch.setattr(ingest_route, "COMPLETIONS_API_KEY", "sk-regression")
|
||||
with make_client(tmp_path) as client:
|
||||
resp = client.post(
|
||||
"/api/ingest",
|
||||
json={"content": "x " * 50},
|
||||
headers={"Authorization": "Bearer wrong-key", "Content-Type": "application/json"},
|
||||
)
|
||||
assert resp.status_code == 401
|
||||
|
||||
|
||||
# ── a client-supplied conversation_id that no longer exists ────────────
|
||||
|
||||
|
||||
def test_chat_bogus_conversation_id_creates_row(tmp_path: Path, monkeypatch):
|
||||
events = _stream_json_lines([
|
||||
{"message": {"content": "hi"}, "logprobs": [{"logprob": -0.01}]},
|
||||
{"done": True, "eval_count": 1, "eval_duration": 1000000000},
|
||||
])
|
||||
|
||||
def stream_stub(self, method, url, json=None, timeout=None):
|
||||
return _MockStreamResponse(events)
|
||||
|
||||
monkeypatch.setattr(httpx.AsyncClient, "stream", stream_stub)
|
||||
|
||||
with make_client(tmp_path) as client:
|
||||
resp = client.post(
|
||||
"/api/chat",
|
||||
json={"message": "hello", "conversation_id": "ghost-conv", "model": config.DEFAULT_MODEL},
|
||||
headers=_guest_headers(client),
|
||||
)
|
||||
assert resp.status_code == 200, resp.text
|
||||
conv_resp = client.get("/api/conversations/ghost-conv", headers=_guest_headers(client))
|
||||
assert conv_resp.status_code == 200
|
||||
assert len(conv_resp.json()["messages"]) >= 2
|
||||
|
||||
|
||||
def test_search_bogus_conversation_id_creates_row(tmp_path: Path, monkeypatch):
|
||||
async def empty_search(query: str, max_results: int = 5):
|
||||
return []
|
||||
|
||||
monkeypatch.setattr(routers.search_route, "query_searxng", empty_search)
|
||||
|
||||
with make_client(tmp_path) as client:
|
||||
resp = client.post(
|
||||
"/api/search",
|
||||
json={"query": "nothing here", "conversation_id": "ghost-search", "model": config.DEFAULT_MODEL},
|
||||
headers=_guest_headers(client),
|
||||
)
|
||||
assert resp.status_code == 200, resp.text
|
||||
conv_resp = client.get("/api/conversations/ghost-search", headers=_guest_headers(client))
|
||||
assert conv_resp.status_code == 200
|
||||
assert len(conv_resp.json()["messages"]) >= 1
|
||||
|
||||
|
||||
# ── auto-search augmentation must reset the streamed text ──────────────
|
||||
|
||||
|
||||
def test_auto_search_augmented_event_has_reset_flag(tmp_path: Path, monkeypatch):
|
||||
first_stream = _stream_json_lines([
|
||||
{"message": {"content": "I don't have current data on that."}, "logprobs": [{"logprob": -5.0}]},
|
||||
{"done": True, "eval_count": 2, "eval_duration": 1000000000},
|
||||
])
|
||||
second_stream = _stream_json_lines([
|
||||
{"message": {"content": "According to the search results, the value is forty-two."}},
|
||||
{"done": True},
|
||||
])
|
||||
batches = [first_stream, second_stream]
|
||||
|
||||
def stream_stub(self, method, url, json=None, timeout=None):
|
||||
return _MockStreamResponse(batches.pop(0))
|
||||
|
||||
async def search_stub(query: str, max_results: int = 5):
|
||||
return [{"title": "Answer", "url": "https://example.com", "content": "The value is 42."}]
|
||||
|
||||
with make_client(tmp_path) as client:
|
||||
monkeypatch.setattr(httpx.AsyncClient, "stream", stream_stub)
|
||||
monkeypatch.setattr(routers.chat, "query_searxng", search_stub)
|
||||
resp = client.post(
|
||||
"/api/chat",
|
||||
json={"message": "what is the latest value", "model": config.DEFAULT_MODEL},
|
||||
headers=_guest_headers(client),
|
||||
)
|
||||
assert resp.status_code == 200, resp.text
|
||||
payloads = parse_sse_payloads(resp.text)
|
||||
|
||||
augmented = [p for p in payloads if p.get("augmented")]
|
||||
assert augmented, "expected an augmented token event"
|
||||
assert augmented[0].get("reset") is True
|
||||
# The augmented token must carry the fresh answer, not the discarded
|
||||
# first-pass "I don't have current data" text.
|
||||
assert "According to the search results" in augmented[0]["token"]
|
||||
assert "I don't have current data" not in augmented[0]["token"]
|
||||
|
||||
|
||||
# ── image uploads are placeholders, never text-ingested ────────────────
|
||||
|
||||
|
||||
def test_upload_image_skips_ingest(tmp_path: Path, monkeypatch):
|
||||
fake = _FakeAsyncClient()
|
||||
monkeypatch.setattr(httpx, "AsyncClient", lambda *a, **kw: fake)
|
||||
with make_client(tmp_path) as client:
|
||||
resp = client.post(
|
||||
"/api/upload",
|
||||
headers=_admin_headers(client),
|
||||
data={"mode": "both"},
|
||||
files={"file": ("photo.png", b"\x89PNG\r\n\x1a\nfake", "image/png")},
|
||||
)
|
||||
assert resp.status_code == 200, resp.text
|
||||
data = resp.json()
|
||||
context_id = data["context_id"]
|
||||
|
||||
assert data["chunks_ingested"] == 0
|
||||
assert data["note"]
|
||||
assert fake.put_payloads == []
|
||||
assert data["filename"] == "photo.png"
|
||||
|
||||
row = db.get_db().execute(
|
||||
"SELECT content FROM upload_context WHERE id = ?", (context_id,)
|
||||
).fetchone()
|
||||
assert row and decrypt_text(row["content"]) == "[Image: photo.png]"
|
||||
|
||||
|
||||
# ── conflict detection needs a shared subject, not just an FTS hit ─────
|
||||
|
||||
|
||||
def test_conflict_detection_requires_shared_subject(tmp_path: Path):
|
||||
os.environ["CAIC_ADMIN_PIN"] = "1234"
|
||||
db.DB_PATH = tmp_path / "caic-mem-regression.db"
|
||||
SESSIONS.clear()
|
||||
PIN_ATTEMPTS.clear()
|
||||
RATE_EVENTS.clear()
|
||||
db.init_db()
|
||||
|
||||
memory.add_memory("the cat sat on the mat", "general")
|
||||
conflicts = memory.check_fact_conflicts(["the dog is brown"])
|
||||
assert conflicts == []
|
||||
|
||||
memory.add_memory("I prefer Rust over Go", "preference")
|
||||
conflicts = memory.check_fact_conflicts(["I prefer Go over Rust"])
|
||||
assert len(conflicts) == 1
|
||||
assert conflicts[0]["new_fact"] == "I prefer Go over Rust"
|
||||
assert conflicts[0]["old_fact"] == "I prefer Rust over Go"
|
||||
assert "memory_id" in conflicts[0]
|
||||
|
||||
|
||||
# ── ingest point IDs are deterministic (no duplicate vectors) ──────────
|
||||
|
||||
|
||||
def test_ingest_deterministic_point_ids(tmp_path: Path, monkeypatch):
|
||||
monkeypatch.setattr(ingest_route, "COMPLETIONS_API_KEY", "sk-regression")
|
||||
|
||||
captured_first = []
|
||||
captured_second = []
|
||||
|
||||
class CaptureClient:
|
||||
FakeResponse = _FakeAsyncClient.FakeResponse
|
||||
|
||||
def __init__(self, *a, **kw):
|
||||
self.capture = None
|
||||
|
||||
async def __aenter__(self):
|
||||
return self
|
||||
|
||||
async def __aexit__(self, *a):
|
||||
pass
|
||||
|
||||
async def post(self, url, **kw):
|
||||
if "/api/embeddings" in url:
|
||||
return self.FakeResponse(200, {"embedding": [0.2] * 768})
|
||||
return self.FakeResponse(200)
|
||||
|
||||
async def put(self, url, **kw):
|
||||
payload = kw.get("json", {})
|
||||
if self.capture is not None:
|
||||
self.capture.append(payload["points"][0]["id"])
|
||||
return self.FakeResponse(200)
|
||||
|
||||
body = {"content": "alpha beta gamma delta epsilon " * 8, "source": "hook"}
|
||||
headers = {"Authorization": "Bearer sk-regression", "Content-Type": "application/json"}
|
||||
|
||||
fake1 = CaptureClient()
|
||||
fake1.capture = captured_first
|
||||
monkeypatch.setattr(httpx, "AsyncClient", lambda *a, **kw: fake1)
|
||||
with make_client(tmp_path) as client:
|
||||
r1 = client.post("/api/ingest", json=body, headers=headers)
|
||||
assert r1.status_code == 200, r1.text
|
||||
|
||||
fake2 = CaptureClient()
|
||||
fake2.capture = captured_second
|
||||
monkeypatch.setattr(httpx, "AsyncClient", lambda *a, **kw: fake2)
|
||||
with make_client(tmp_path) as client:
|
||||
r2 = client.post("/api/ingest", json=body, headers=headers)
|
||||
assert r2.status_code == 200, r2.text
|
||||
|
||||
assert captured_first and captured_second
|
||||
assert len(captured_first) == len(captured_second)
|
||||
assert captured_first == captured_second, "re-ingesting identical content changed point ids"
|
||||
assert len(set(captured_first)) == len(captured_first)
|
||||
|
||||
|
||||
# ── get_load() VRAM parsing ────────────────────────────────────────────
|
||||
|
||||
|
||||
def test_get_load_vram_parses_rocm_output(monkeypatch):
|
||||
monkeypatch.setattr(agent, "HAS_PSUTIL", False)
|
||||
output = (
|
||||
"======================= ROCm System Management Interface =======================\n"
|
||||
"GPU[0] : gfx1030\n"
|
||||
"VRAM Total Used Memory (B): 3221225472\n"
|
||||
"VRAM Total Memory (B): 17179869184\n"
|
||||
)
|
||||
fake = subprocess.CompletedProcess(["rocm-smi", "--showmeminfo", "vram"], 0, output, "")
|
||||
monkeypatch.setattr(subprocess, "run", lambda *a, **kw: fake)
|
||||
load = agent.get_load()
|
||||
assert load["vram_pct"] == 19 # 3 GiB / 16 GiB
|
||||
|
||||
|
||||
def test_get_load_vram_absent_does_not_crash(monkeypatch):
|
||||
# Regression: rocm-smi returned no parseable VRAM lines, so the old code
|
||||
# left total/used unbound and raised.
|
||||
monkeypatch.setattr(agent, "HAS_PSUTIL", False)
|
||||
output = "======================= ROCm System Management Interface =======================\nNo GPU detected\n"
|
||||
fake = subprocess.CompletedProcess(["rocm-smi", "--showmeminfo", "vram"], 0, output, "")
|
||||
monkeypatch.setattr(subprocess, "run", lambda *a, **kw: fake)
|
||||
load = agent.get_load()
|
||||
assert "vram_pct" not in load
|
||||
|
||||
|
||||
# ── version pinning ────────────────────────────────────────────────────
|
||||
|
||||
|
||||
def test_version_is_bumped():
|
||||
assert re.fullmatch(r"v\d+\.\d+\.\d+", config.VERSION)
|
||||
@@ -2,7 +2,6 @@
|
||||
import asyncio
|
||||
|
||||
import cluster
|
||||
import triage
|
||||
from config import AMQP_EXCHANGE_ADMIN, AMQP_EXCHANGE_SYSTEM
|
||||
|
||||
|
||||
@@ -149,54 +148,3 @@ def test_handle_model_failed_unknown_node(caplog, monkeypatch):
|
||||
))
|
||||
|
||||
assert any("unknown node" in rec.message for rec in caplog.records)
|
||||
|
||||
|
||||
# ---------- 4. select_node() triggers swap when model mismatched ----------
|
||||
|
||||
|
||||
def test_select_node_code_triggers_swap(monkeypatch):
|
||||
_reset()
|
||||
_published.clear()
|
||||
monkeypatch.setattr(cluster, "publish", _fake_publish)
|
||||
|
||||
cluster.CLUSTER_NODES["jarvis"] = {
|
||||
"name": "jarvis", "type": "worker", "status": "active",
|
||||
"ip": "192.168.50.210",
|
||||
"active_model": {"name": "llama3.1", "port": 8081},
|
||||
"inventory": [
|
||||
{"filename": "qwen2.5-coder-14b-Q4_K_M.gguf", "name": "qwen2.5-coder", "version": "14b", "quant": "Q4_K_M"},
|
||||
],
|
||||
}
|
||||
|
||||
result = asyncio.run(triage.select_node("code"))
|
||||
|
||||
assert result is None
|
||||
# Swap should have been published
|
||||
assert any("cmd.swap_model" in rk for _, rk, _ in _published)
|
||||
# Node should now be swapping
|
||||
assert cluster.CLUSTER_NODES["jarvis"]["status"] == "swapping"
|
||||
|
||||
|
||||
# ---------- 5. select_node() returns None when node is already swapping ----------
|
||||
|
||||
|
||||
def test_select_node_swapping_returns_none(monkeypatch):
|
||||
_reset()
|
||||
_published.clear()
|
||||
monkeypatch.setattr(cluster, "publish", _fake_publish)
|
||||
|
||||
cluster.CLUSTER_NODES["jarvis"] = {
|
||||
"name": "jarvis", "type": "worker", "status": "swapping",
|
||||
"ip": "192.168.50.210",
|
||||
"active_model": {"name": "llama3.1", "port": 8081},
|
||||
"inventory": [
|
||||
{"filename": "qwen2.5-coder-14b-Q4_K_M.gguf", "name": "qwen2.5-coder"},
|
||||
],
|
||||
}
|
||||
|
||||
result = asyncio.run(triage.select_node("code"))
|
||||
|
||||
assert result is None
|
||||
# No swap command should be published while already swapping
|
||||
swap_published = any("cmd.swap_model" in rk for _, rk, _ in _published)
|
||||
assert not swap_published
|
||||
|
||||
@@ -1,141 +0,0 @@
|
||||
"""Tests for triage.py — query classification and node selection."""
|
||||
import asyncio
|
||||
import json
|
||||
|
||||
import httpx
|
||||
|
||||
import cluster
|
||||
import config
|
||||
import triage
|
||||
|
||||
|
||||
def _reset():
|
||||
cluster.CLUSTER_NODES.clear()
|
||||
cluster.CLUSTER_COORDINATOR = None
|
||||
|
||||
|
||||
_published = []
|
||||
|
||||
|
||||
async def _fake_publish(exchange, routing_key, payload):
|
||||
_published.append((exchange, routing_key, payload))
|
||||
|
||||
|
||||
class _MockPostResponse:
|
||||
def __init__(self, json_data: dict, status_code: int = 200):
|
||||
self._json_data = json_data
|
||||
self.status_code = status_code
|
||||
|
||||
def json(self):
|
||||
return self._json_data
|
||||
|
||||
async def __aenter__(self):
|
||||
return self
|
||||
|
||||
async def __aexit__(self, exc_type, exc, tb):
|
||||
return False
|
||||
|
||||
|
||||
class _MockPostContext:
|
||||
def __init__(self, response: _MockPostResponse):
|
||||
self._response = response
|
||||
|
||||
async def __aenter__(self):
|
||||
return self._response
|
||||
|
||||
async def __aexit__(self, exc_type, exc, tb):
|
||||
return False
|
||||
|
||||
|
||||
# ---------- 1. classify_query returns valid classification ----------
|
||||
|
||||
|
||||
def test_classify_returns_valid(monkeypatch):
|
||||
async def post_stub(self, url, json=None, timeout=None):
|
||||
return _MockPostResponse({
|
||||
"choices": [{"message": {"content": "code"}}]
|
||||
})
|
||||
|
||||
monkeypatch.setattr(httpx.AsyncClient, "post", post_stub)
|
||||
|
||||
result = __import__("asyncio").run(triage.classify_query("write a python function"))
|
||||
assert result == "code"
|
||||
|
||||
|
||||
# ---------- 2. classify_query on error returns "general" ----------
|
||||
|
||||
|
||||
def test_classify_error_returns_general(monkeypatch):
|
||||
async def post_stub(self, url, json=None, timeout=None):
|
||||
raise httpx.ConnectError("connection refused")
|
||||
|
||||
monkeypatch.setattr(httpx.AsyncClient, "post", post_stub)
|
||||
|
||||
result = __import__("asyncio").run(triage.classify_query("any question"))
|
||||
assert result == "general"
|
||||
|
||||
|
||||
# ---------- 3. select_node("code") returns coder node ----------
|
||||
|
||||
|
||||
def test_select_node_code_returns_coder():
|
||||
_reset()
|
||||
cluster.CLUSTER_NODES["coder01"] = {
|
||||
"name": "coder01", "type": "worker", "status": "active",
|
||||
"ip": "192.168.50.210",
|
||||
"active_model": {"name": "qwen2.5-coder-14b", "port": 8082},
|
||||
}
|
||||
cluster.CLUSTER_NODES["general01"] = {
|
||||
"name": "general01", "type": "worker", "status": "active",
|
||||
"ip": "192.168.50.211",
|
||||
"active_model": {"name": "llama3.1", "port": 8081},
|
||||
}
|
||||
|
||||
node = asyncio.run(triage.select_node("code"))
|
||||
assert node is not None
|
||||
assert node["name"] == "coder01"
|
||||
|
||||
|
||||
# ---------- 4. select_node("general") with no matching node returns None ----------
|
||||
|
||||
|
||||
def test_select_node_general_no_match_returns_none():
|
||||
_reset()
|
||||
cluster.CLUSTER_NODES["coder01"] = {
|
||||
"name": "coder01", "type": "worker", "status": "active",
|
||||
"active_model": {"name": "qwen2.5-coder-14b", "port": 8082},
|
||||
}
|
||||
node = asyncio.run(triage.select_node("general"))
|
||||
assert node is None
|
||||
|
||||
|
||||
# ---------- 5. get_inference_url with coder node ----------
|
||||
|
||||
|
||||
def test_get_inference_url_with_coder_node(monkeypatch):
|
||||
_reset()
|
||||
async def fake_classify(query: str) -> str:
|
||||
return "code"
|
||||
monkeypatch.setattr(triage, "classify_query", fake_classify)
|
||||
|
||||
cluster.CLUSTER_NODES["coder01"] = {
|
||||
"name": "coder01", "type": "worker", "status": "active",
|
||||
"ip": "192.168.50.210",
|
||||
"active_model": {"name": "qwen2.5-coder-14b", "port": 8082},
|
||||
}
|
||||
|
||||
url = __import__("asyncio").run(triage.get_inference_url("write a loop in rust"))
|
||||
assert url == "http://192.168.50.210:8082/v1"
|
||||
|
||||
|
||||
# ---------- 6. get_inference_url with no nodes returns LLAMA_SERVER_BASE ----------
|
||||
|
||||
|
||||
def test_get_inference_url_no_nodes(monkeypatch):
|
||||
_reset()
|
||||
async def fake_classify(query: str) -> str:
|
||||
return "code"
|
||||
monkeypatch.setattr(triage, "classify_query", fake_classify)
|
||||
|
||||
url = __import__("asyncio").run(triage.get_inference_url("any question"))
|
||||
assert url == config.LLAMA_SERVER_BASE
|
||||
Reference in New Issue
Block a user