tests/test_vector_store.py

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"