|
| 1 | +import os |
| 2 | +import time |
| 3 | +import json |
| 4 | +from pathlib import Path |
| 5 | +import numpy as np |
| 6 | +import pytest |
| 7 | +from fastapi.testclient import TestClient |
| 8 | + |
| 9 | +from neuralcache.api.server import app, reranker, settings |
| 10 | +from neuralcache.narrative import NarrativeTracker |
| 11 | +from neuralcache.types import Document, RerankRequest |
| 12 | +from neuralcache.rerank import Reranker |
| 13 | +from neuralcache.config import Settings |
| 14 | + |
| 15 | + |
| 16 | +def test_narrative_purge_if_stale_json(tmp_path): |
| 17 | + """Purge uses internal _updated_ts, so set that stale and confirm JSON file removal.""" |
| 18 | + narr = NarrativeTracker(dim=4, alpha=0.5, success_gate=0.0, backend="json", path="n.json", storage_dir=str(tmp_path)) |
| 19 | + emb = np.ones((4,), dtype=np.float32) |
| 20 | + narr.update(emb, success=1.0) |
| 21 | + file_path = Path(tmp_path)/"n.json" |
| 22 | + assert file_path.exists() |
| 23 | + # Make internal timestamp stale |
| 24 | + narr._updated_ts = time.time() - 10_000 # type: ignore[attr-defined] |
| 25 | + narr.purge_if_stale(retention_seconds=3600) |
| 26 | + assert not file_path.exists(), "Narrative JSON should be deleted when stale" |
| 27 | + assert np.allclose(narr.v, np.zeros_like(narr.v)) |
| 28 | + |
| 29 | + |
| 30 | +def test_cr_empty_candidates_fallback(monkeypatch): |
| 31 | + # Force CR on but monkeypatch hierarchical_candidates to return empty |
| 32 | + s = Settings() |
| 33 | + s.cr.on = True |
| 34 | + rr = Reranker(settings=s) |
| 35 | + docs = [Document(id=str(i), text=f"text {i}") for i in range(5)] |
| 36 | + q = rr.encode_query("query") |
| 37 | + import neuralcache.cr.search as crsearch |
| 38 | + monkeypatch.setattr(crsearch, "hierarchical_candidates", lambda **kwargs: []) |
| 39 | + scored = rr.score(q, docs, query_text="query") |
| 40 | + ids = {d.id for d in scored} |
| 41 | + assert ids == {d.id for d in docs}, "Fallback should revert to all documents when CR yields none" |
| 42 | + |
| 43 | + |
| 44 | +def test_encoder_unknown_backend_warning(caplog): |
| 45 | + caplog.set_level("WARNING") |
| 46 | + from neuralcache.encoder import create_encoder |
| 47 | + enc = create_encoder("unknown-backend-xyz", dim=8) |
| 48 | + assert enc.encode("hello").shape[0] == 8 |
| 49 | + assert any("Unknown embedding backend" in rec.message for rec in caplog.records) |
| 50 | + |
| 51 | + |
| 52 | +def test_rate_limit_enforced(monkeypatch): |
| 53 | + client = TestClient(app) |
| 54 | + # Patch settings directly for test isolation |
| 55 | + orig = settings.rate_limit_per_minute |
| 56 | + settings.rate_limit_per_minute = 2 |
| 57 | + try: |
| 58 | + payload = {"query":"q", "documents":[{"id":"a","text":"A"},{"id":"b","text":"B"}], "top_k":1} |
| 59 | + # First two should pass |
| 60 | + for _ in range(2): |
| 61 | + r = client.post("/rerank", json=payload) |
| 62 | + assert r.status_code == 200 |
| 63 | + # Third should rate limit |
| 64 | + r3 = client.post("/rerank", json=payload) |
| 65 | + assert r3.status_code == 429 |
| 66 | + body = r3.json() |
| 67 | + assert body["error"]["code"] == "RATE_LIMITED" |
| 68 | + finally: |
| 69 | + settings.rate_limit_per_minute = orig |
| 70 | + |
| 71 | + |
| 72 | +def test_api_key_authentication(monkeypatch): |
| 73 | + client = TestClient(app) |
| 74 | + # Inject an API token |
| 75 | + orig_tokens = list(settings.api_tokens) |
| 76 | + settings.api_tokens = ["secret123"] |
| 77 | + try: |
| 78 | + payload = {"query":"q", "documents":[{"id":"a","text":"A"}], "top_k":1} |
| 79 | + # Missing token -> 401 |
| 80 | + r = client.post("/rerank", json=payload) |
| 81 | + assert r.status_code == 401 |
| 82 | + assert r.json()["error"]["code"] == "UNAUTHORIZED" |
| 83 | + # Provide token via x-api-key |
| 84 | + r2 = client.post("/rerank", json=payload, headers={"x-api-key":"secret123"}) |
| 85 | + assert r2.status_code == 200 |
| 86 | + finally: |
| 87 | + settings.api_tokens = orig_tokens |
| 88 | + |
| 89 | + |
| 90 | +def test_batch_rerank_gating_debug(): |
| 91 | + client = TestClient(app) |
| 92 | + # Use small document lists for clarity |
| 93 | + batch = [ |
| 94 | + {"query":"q1","documents":[{"id":"a","text":"A"},{"id":"b","text":"B"}],"top_k":1}, |
| 95 | + {"query":"q2","documents":[{"id":"c","text":"C"},{"id":"d","text":"D"}],"top_k":1} |
| 96 | + ] |
| 97 | + # Override gating to 'on' via query param use_cr left default |
| 98 | + resp = client.post("/rerank/batch?use_cr=false", json=batch) |
| 99 | + assert resp.status_code == 200 |
| 100 | + data = resp.json() |
| 101 | + assert isinstance(data, list) and len(data) == 2 |
| 102 | + # Ensure debug present and has gating keys |
| 103 | + for item in data: |
| 104 | + debug = item.get("debug") |
| 105 | + assert debug is not None |
| 106 | + gating_dbg = debug.get("gating") |
| 107 | + assert gating_dbg is not None |
| 108 | + assert set(["mode","uncertainty","use_gating","candidate_count","effective_candidate_count","total_candidates"]).issubset(gating_dbg.keys()) |
0 commit comments