import asyncio import json import os import subprocess from pathlib import Path import httpx import psutil from fastapi.testclient import TestClient import app as app_module import config import db import hardware from security import SESSIONS, PIN_ATTEMPTS, RATE_EVENTS def make_client(tmp_path: Path) -> TestClient: os.environ["JARVISCHAT_ADMIN_PIN"] = "1234" db.DB_PATH = tmp_path / "jarvischat-hardware.db" hardware.HARDWARE_STATE_PATH = tmp_path / "hardware_state.json" SESSIONS.clear() PIN_ATTEMPTS.clear() RATE_EVENTS.clear() db.init_db() return TestClient(app_module.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"} class _MockGet: def __init__(self, status_code: int = 200, json_data: dict | None = None): self.status_code = status_code self._json_data = json_data or {} def json(self): return self._json_data def _mock_subprocess(rocm_stdout: str = "") -> object: class MockProc: returncode = 0 stdout = rocm_stdout class MockSP: TimeoutExpired = subprocess.TimeoutExpired run = staticmethod(lambda cmd, **kw: MockProc()) return MockSP() def _broken_subprocess(exception: Exception) -> object: class MockSP: TimeoutExpired = subprocess.TimeoutExpired run = staticmethod(lambda cmd, **kw: (_ for _ in ()).throw(exception)) return MockSP() def test_assess_hardware_all_services_reachable(tmp_path: Path, monkeypatch): hardware.HARDWARE_STATE_PATH = tmp_path / "hardware_state.json" monkeypatch.setattr(psutil, "virtual_memory", lambda: type("M", (), {"total": 16 * 1024 ** 3, "available": 8 * 1024 ** 3})()) monkeypatch.setattr(psutil, "cpu_count", lambda: 8) monkeypatch.setattr(hardware, "subprocess", _mock_subprocess( json.dumps({"card0": {"VRAM Total (MB)": 8192, "VRAM Free (MB)": 4096}}) )) async def mock_get(self, url, *args, **kwargs): if "v1/models" in url: return _MockGet(200, {"data": [{"id": "mistral-nemo:latest"}]}) if "6333" in url: return _MockGet(200, {"result": {"collections": [{"name": "jarvischat"}]}}) if "8888" in url: return _MockGet(200, {}) return _MockGet(200, {}) monkeypatch.setattr(httpx.AsyncClient, "get", mock_get) state = asyncio.run(hardware.assess_hardware()) assert state["ram_total_gb"] == 16.0 assert state["ram_available_gb"] == 8.0 assert state["cpu_count"] == 8 assert state["vram_total_mb"] == 8192 assert state["vram_free_mb"] == 4096 assert state["llama_reachable"] is True assert state["llama_models"] == ["mistral-nemo:latest"] assert state["qdrant_reachable"] is True assert state["qdrant_collections"] == ["jarvischat"] assert state["searxng_reachable"] is True assert tmp_path.joinpath("hardware_state.json").exists() def test_assess_hardware_rocm_smi_absent(tmp_path: Path, monkeypatch): hardware.HARDWARE_STATE_PATH = tmp_path / "hardware_state.json" monkeypatch.setattr(psutil, "virtual_memory", lambda: type("M", (), {"total": 16 * 1024 ** 3, "available": 8 * 1024 ** 3})()) monkeypatch.setattr(psutil, "cpu_count", lambda: 8) monkeypatch.setattr(hardware, "subprocess", _broken_subprocess(FileNotFoundError("no rocm-smi"))) async def mock_get(self, url, *args, **kwargs): if "v1/models" in url: return _MockGet(200, {"data": []}) if "6333" in url: return _MockGet(200, {"result": {"collections": []}}) if "8888" in url: return _MockGet(200, {}) return _MockGet(200, {}) monkeypatch.setattr(httpx.AsyncClient, "get", mock_get) state = asyncio.run(hardware.assess_hardware()) assert state["vram_total_mb"] == 0 assert state["vram_free_mb"] == 0 assert state["llama_reachable"] is True assert state["qdrant_reachable"] is True assert state["searxng_reachable"] is True def test_assess_hardware_llama_unreachable(tmp_path: Path, monkeypatch): hardware.HARDWARE_STATE_PATH = tmp_path / "hardware_state.json" monkeypatch.setattr(psutil, "virtual_memory", lambda: type("M", (), {"total": 16 * 1024 ** 3, "available": 8 * 1024 ** 3})()) monkeypatch.setattr(psutil, "cpu_count", lambda: 8) monkeypatch.setattr(hardware, "subprocess", _mock_subprocess( json.dumps({"card0": {"VRAM Total (MB)": 8192, "VRAM Free (MB)": 4096}}) )) async def mock_get(self, url, *args, **kwargs): if "v1/models" in url: raise httpx.ConnectError("refused") if "6333" in url: return _MockGet(200, {"result": {"collections": [{"name": "jarvischat"}]}}) if "8888" in url: return _MockGet(200, {}) return _MockGet(200, {}) monkeypatch.setattr(httpx.AsyncClient, "get", mock_get) state = asyncio.run(hardware.assess_hardware()) assert state["llama_reachable"] is False assert state["llama_models"] == [] assert state["qdrant_reachable"] is True assert state["searxng_reachable"] is True def test_get_hardware_endpoint(tmp_path: Path, monkeypatch): hardware.HARDWARE_STATE_PATH = tmp_path / "hardware_state.json" monkeypatch.setattr(psutil, "virtual_memory", lambda: type("M", (), {"total": 16 * 1024 ** 3, "available": 8 * 1024 ** 3})()) monkeypatch.setattr(psutil, "cpu_count", lambda: 8) monkeypatch.setattr(hardware, "subprocess", _mock_subprocess( json.dumps({"card0": {"VRAM Total (MB)": 8192, "VRAM Free (MB)": 4096}}) )) async def mock_get(self, url, *args, **kwargs): if "v1/models" in url: return _MockGet(200, {"data": [{"id": "mistral-nemo:latest"}]}) if "6333" in url: return _MockGet(200, {"result": {"collections": [{"name": "jarvischat"}]}}) if "8888" in url: return _MockGet(200, {}) return _MockGet(200, {}) monkeypatch.setattr(httpx.AsyncClient, "get", mock_get) with make_client(tmp_path) as client: resp = client.get("/api/hardware", headers=_guest_headers(client)) assert resp.status_code == 200 data = resp.json() assert data["ram_total_gb"] == 16.0 assert data["cpu_count"] == 8 assert data["vram_total_mb"] == 8192 assert data["llama_reachable"] is True assert data["qdrant_reachable"] is True assert data["searxng_reachable"] is True assert "llama_models" in data assert "qdrant_collections" in data