File size: 5,811 Bytes
35676b4 7880373 35676b4 7880373 35676b4 | 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 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 | import pytest
from storage import vector_store
@pytest.fixture(autouse=True)
def tmp_chroma(tmp_path, monkeypatch):
monkeypatch.setattr(vector_store, "CHROMA_PATH", tmp_path / "chroma")
def test_add_and_search_returns_match():
vector_store.add_chunks(
collection_name="filings",
chunks=["Revenue increased 5% year over year driven by iPhone sales."],
metadatas=[{"ticker": "AAPL", "source": "10-Q", "filing_date": "2024-11-01", "section": "MD&A"}],
ids=["AAPL-10Q-mda-0"],
)
results = vector_store.search("filings", "revenue growth", "AAPL", n_results=1)
assert len(results) == 1
assert "Revenue" in results[0]["text"]
assert results[0]["metadata"]["ticker"] == "AAPL"
assert results[0]["evidence_ref"]["evidence_id"].startswith("ev_")
assert len(results[0]["evidence_ref"]["content_hash"]) == 64
def test_search_unknown_collection_returns_empty():
results = vector_store.search("filings", "revenue", "MSFT", n_results=3)
assert results == []
def test_search_does_not_disguise_storage_failure_as_empty(monkeypatch):
class BrokenClient:
def get_collection(self, **_kwargs):
raise OSError("database disk image is malformed")
monkeypatch.setattr(vector_store, "_client", lambda: BrokenClient())
with pytest.raises(RuntimeError, match="unavailable or corrupt"):
vector_store.search("filings", "revenue", "MSFT", n_results=3)
def test_delete_does_not_hide_storage_failure(monkeypatch):
class BrokenClient:
def get_collection(self, **_kwargs):
raise OSError("database disk image is malformed")
monkeypatch.setattr(vector_store, "_client", lambda: BrokenClient())
with pytest.raises(RuntimeError, match="for deletion"):
vector_store.delete_by_ticker("filings", "MSFT")
def test_search_ignores_legacy_chunk_without_stored_provenance(monkeypatch):
class LegacyCollection:
def query(self, **_kwargs):
return {
"documents": [["Legacy filing text about revenue."]],
"metadatas": [[{
"ticker": "AAPL", "source": "10-Q",
"filing_date": "2024-11-01", "section": "MD&A",
}]],
}
class LegacyClient:
def get_collection(self, **_kwargs):
return LegacyCollection()
monkeypatch.setattr(vector_store, "_client", lambda: LegacyClient())
assert vector_store.search("filings", "revenue", "AAPL", n_results=3) == []
def test_search_ignores_chunk_with_tampered_provenance(monkeypatch):
class TamperedCollection:
def query(self, **_kwargs):
return {
"documents": [["Revenue increased 5%."]],
"metadatas": [[{
"ticker": "AAPL", "source": "10-Q",
"filing_date": "2024-11-01", "section": "MD&A",
"document_id": "sec:AAPL:test", "chunk_id": "chunk-1",
"content_hash": "0" * 64, "evidence_id": "ev_tampered",
}]],
}
class TamperedClient:
def get_collection(self, **_kwargs):
return TamperedCollection()
monkeypatch.setattr(vector_store, "_client", lambda: TamperedClient())
assert vector_store.search("filings", "revenue", "AAPL", n_results=3) == []
def test_search_ignores_source_that_does_not_belong_to_collection(monkeypatch):
class WrongSourceCollection:
def query(self, **_kwargs):
return {
"documents": [["A metrics record injected into filing search."]],
"metadatas": [[{
"ticker": "AAPL", "source": "metrics",
"document_id": "metrics:AAPL:test", "chunk_id": "chunk-1",
"content_hash": "0" * 64, "evidence_id": "ev_wrong_source",
}]],
}
class WrongSourceClient:
def get_collection(self, **_kwargs):
return WrongSourceCollection()
monkeypatch.setattr(vector_store, "_client", lambda: WrongSourceClient())
assert vector_store.search("filings", "revenue", "AAPL", n_results=3) == []
def test_upsert_replaces_by_id():
vector_store.add_chunks(
collection_name="filings",
chunks=["Old text."],
metadatas=[{"ticker": "AAPL", "source": "10-Q", "filing_date": "2024-01-01", "section": "MD&A"}],
ids=["AAPL-chunk-0"],
)
vector_store.add_chunks(
collection_name="filings",
chunks=["New text about revenue growth."],
metadatas=[{"ticker": "AAPL", "source": "10-Q", "filing_date": "2024-11-01", "section": "MD&A"}],
ids=["AAPL-chunk-0"],
)
results = vector_store.search("filings", "revenue", "AAPL", n_results=1)
assert "New text" in results[0]["text"]
def test_search_with_min_filing_date_filter():
vector_store.add_chunks(
collection_name="filings",
chunks=[
"Old filing from 2023 discussing revenue.",
"New filing from 2025 discussing revenue."
],
metadatas=[
{"ticker": "AAPL", "source": "10-Q", "filing_date": "2023-01-15", "section": "MD&A"},
{"ticker": "AAPL", "source": "10-Q", "filing_date": "2025-01-15", "section": "MD&A"}
],
ids=["AAPL-2023-0", "AAPL-2025-0"],
)
# Search without date filter should return both
results = vector_store.search("filings", "revenue", "AAPL", n_results=5)
assert len(results) == 2
# Search with min_filing_date should only return the newer chunk
results = vector_store.search("filings", "revenue", "AAPL", n_results=5, min_filing_date="2024-01-01")
assert len(results) == 1
assert "2025" in results[0]["text"]
assert results[0]["metadata"]["filing_date"] == "2025-01-15"
|