112 lines
4.5 KiB
Python
112 lines
4.5 KiB
Python
import unittest
|
|
|
|
from fastapi import FastAPI
|
|
from fastapi.testclient import TestClient
|
|
|
|
from llamacpp_ha.config import BackendConfig, ProxyConfig
|
|
from llamacpp_ha.monitor import ProxyStats, build_router
|
|
from llamacpp_ha.queue import RequestQueue
|
|
from llamacpp_ha.registry import BackendRegistry, BackendState
|
|
from llamacpp_ha.session_store import SessionStore
|
|
from llamacpp_ha.slot_tracker import SlotTracker
|
|
|
|
|
|
def _make_monitor_app(states=None):
|
|
cfg = ProxyConfig(backends=[BackendConfig(url=s.url) for s in (states or [])])
|
|
registry = BackendRegistry(cfg)
|
|
for state in (states or []):
|
|
registry._states[state.url] = state
|
|
registry._rebuild_index()
|
|
|
|
slot_tracker = SlotTracker()
|
|
request_queue = RequestQueue()
|
|
session_store = SessionStore()
|
|
stats = ProxyStats()
|
|
|
|
app = FastAPI()
|
|
router = build_router(
|
|
registry=registry,
|
|
slot_tracker=slot_tracker,
|
|
request_queue=request_queue,
|
|
session_store=session_store,
|
|
stats=stats,
|
|
)
|
|
app.include_router(router)
|
|
return app, registry, slot_tracker, request_queue, session_store, stats
|
|
|
|
|
|
class TestMonitorEndpoints(unittest.TestCase):
|
|
def test_monitor_page_returns_html(self):
|
|
app, *_ = _make_monitor_app()
|
|
resp = TestClient(app).get("/monitor")
|
|
self.assertEqual(resp.status_code, 200)
|
|
self.assertIn("text/html", resp.headers["content-type"])
|
|
self.assertIn("llamacpp-ha", resp.text)
|
|
|
|
def test_monitor_page_no_external_deps(self):
|
|
app, *_ = _make_monitor_app()
|
|
text = TestClient(app).get("/monitor").text
|
|
self.assertNotIn("cdn.", text)
|
|
self.assertNotIn("googleapis.com", text)
|
|
self.assertNotIn("unpkg.com", text)
|
|
|
|
def test_monitor_data_structure(self):
|
|
state = BackendState(config=BackendConfig(url="http://b1"), live=True, models=["m1", "m2"])
|
|
app, _, slot_tracker, *_ = _make_monitor_app([state])
|
|
slot_tracker.set_capacity("http://b1", 4)
|
|
data = TestClient(app).get("/monitor/data").json()
|
|
for key in ("uptime", "total_requests", "queue_depth", "session_count",
|
|
"live_backend_count", "backends", "queue",
|
|
"model_stats", "backend_stats",
|
|
"session_hits", "session_misses", "session_hit_rate"):
|
|
self.assertIn(key, data)
|
|
|
|
def test_monitor_data_backend_fields(self):
|
|
state = BackendState(config=BackendConfig(url="http://b1"), live=True, models=["m1"])
|
|
app, _, slot_tracker, *_ = _make_monitor_app([state])
|
|
slot_tracker.set_capacity("http://b1", 3)
|
|
backend = TestClient(app).get("/monitor/data").json()["backends"][0]
|
|
self.assertEqual(backend["url"], "http://b1")
|
|
self.assertTrue(backend["live"])
|
|
self.assertEqual(backend["models"], ["m1"])
|
|
self.assertEqual(backend["slots_total"], 3)
|
|
self.assertIn("slots_acquired", backend)
|
|
|
|
def test_monitor_data_dead_backend(self):
|
|
state = BackendState(config=BackendConfig(url="http://b2-dead"), live=False, models=[])
|
|
app, *_ = _make_monitor_app([state])
|
|
data = TestClient(app).get("/monitor/data").json()
|
|
backends = {b["url"]: b for b in data["backends"]}
|
|
self.assertFalse(backends["http://b2-dead"]["live"])
|
|
self.assertEqual(data["live_backend_count"], 0)
|
|
|
|
def test_monitor_data_empty_state(self):
|
|
app, *_ = _make_monitor_app([])
|
|
data = TestClient(app).get("/monitor/data").json()
|
|
self.assertEqual(data["backends"], [])
|
|
self.assertEqual(data["queue"], [])
|
|
self.assertEqual(data["session_count"], 0)
|
|
|
|
def test_monitor_total_requests_reflects_stats(self):
|
|
app, *_, stats = _make_monitor_app()
|
|
stats.increment_requests()
|
|
stats.increment_requests()
|
|
data = TestClient(app).get("/monitor/data").json()
|
|
self.assertEqual(data["total_requests"], 2)
|
|
|
|
def test_monitor_uptime_is_string(self):
|
|
app, *_ = _make_monitor_app()
|
|
data = TestClient(app).get("/monitor/data").json()
|
|
self.assertRegex(data["uptime"], r"^\d{2}:\d{2}:\d{2}$")
|
|
|
|
def test_monitor_last_poll_age_never_polled(self):
|
|
"""Backend that has never been polled should show null last_poll_age."""
|
|
state = BackendState(config=BackendConfig(url="http://b1"), live=False, models=[])
|
|
app, *_ = _make_monitor_app([state])
|
|
data = TestClient(app).get("/monitor/data").json()
|
|
self.assertIsNone(data["backends"][0]["last_poll_age"])
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|