personal-llm/tests/test_api.py at 319b33cb034fb88d6eb87e70f0a55987705a9a59 · zaidwhy/personal-llm · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
"""Gateway auth: shared token + Origin rejection (MASTER-FIX-PLAN.md Phase 3 item 12).

FakeRouter/store/vectors fixtures (conftest.py) keep this offline and keyless -
no network call, no API key, ever.
"""

from __future__ import annotations

from fastapi.testclient import TestClient

from personal_llm.engine import Engine
from personal_llm.interfaces import api
from personal_llm.voice import SpeechToText, TextToSpeech

TOKEN = "test-token"


def _fake_engine(store, vectors, router) -> Engine:
    return Engine(store=store, vectors=vectors, router=router, stt=SpeechToText(), tts=TextToSpeech())


def _client(monkeypatch, store, vectors, router):
    monkeypatch.setattr(api, "build_engine", lambda: _fake_engine(store, vectors, router))
    monkeypatch.setattr(api, "_load_or_create_gateway_token", lambda: TOKEN)
    return TestClient(api.app)


def test_tokenless_request_is_rejected(monkeypatch, store, vectors, router):
    client = _client(monkeypatch, store, vectors, router)

    resp = client.get("/stats")

    assert resp.status_code == 401


def test_wrong_token_is_rejected(monkeypatch, store, vectors, router):
    client = _client(monkeypatch, store, vectors, router)

    resp = client.get("/stats", headers={api.GATEWAY_TOKEN_HEADER: "not-the-token"})

    assert resp.status_code == 401


def test_correct_token_allows_stats(monkeypatch, store, vectors, router):
    client = _client(monkeypatch, store, vectors, router)

    resp = client.get("/stats", headers={api.GATEWAY_TOKEN_HEADER: TOKEN})

    assert resp.status_code == 200
    assert resp.json() == store.stats()


def test_origin_header_is_rejected_even_with_a_valid_token(monkeypatch, store, vectors, router):
    client = _client(monkeypatch, store, vectors, router)

    resp = client.get(
        "/stats",
        headers={api.GATEWAY_TOKEN_HEADER: TOKEN, "Origin": "https://evil.example"},
    )

    assert resp.status_code == 403


def test_ask_endpoint_requires_token(monkeypatch, store, vectors, router):
    client = _client(monkeypatch, store, vectors, router)

    resp = client.post("/ask", json={"question": "anything"})

    assert resp.status_code == 401


def test_ask_endpoint_round_trips_with_a_valid_token(monkeypatch, store, vectors, router):
    client = _client(monkeypatch, store, vectors, router)

    resp = client.post(
        "/ask",
        json={"question": "anything"},
        headers={api.GATEWAY_TOKEN_HEADER: TOKEN},
    )

    assert resp.status_code == 200
    assert "text" in resp.json()


def test_voice_ask_endpoint_requires_token(monkeypatch, store, vectors, router):
    client = _client(monkeypatch, store, vectors, router)

    resp = client.post("/voice/ask", files={"file": ("clip.wav", b"not real audio", "audio/wav")})

    assert resp.status_code == 401


def test_load_or_create_gateway_token_persists_across_calls(tmp_path, monkeypatch):
    token_path = tmp_path / "gateway_token"
    monkeypatch.setattr(api, "_gateway_token", None)
    monkeypatch.setattr(api, "get_settings", lambda: type("S", (), {"personal_llm_gateway_token_path": str(token_path)})())

    first = api._load_or_create_gateway_token()

    monkeypatch.setattr(api, "_gateway_token", None)
    second = api._load_or_create_gateway_token()

    assert first == second
    assert token_path.read_text().strip() == first