Spaces:
Sleeping
Sleeping
| """Outcome-classification tests for ChromaDBClient.query_observed (Task 3, Step 4). | |
| These prove the core reason retrieval failures used to be invisible: a | |
| genuinely-empty successful search and a failed search that fell back to ``[]`` | |
| now carry distinct statuses, while ``query()`` stays byte-for-byte compatible. | |
| """ | |
| from __future__ import annotations | |
| import os | |
| os.environ.setdefault("CEREBRAS_API_KEY", "test-key") | |
| import chromadb | |
| import pytest | |
| class _FakeEmbedder: | |
| def embed(self, texts: list[str]) -> list[list[float]]: | |
| return [[0.1, 0.0] for _ in texts] | |
| def db(monkeypatch): | |
| # See test_rag_behavior_unchanged.isolated_db for why this cache-clear is | |
| # required: chromadb.EphemeralClient() shares one process-wide backing | |
| # store, so collection names can otherwise leak data across tests/files. | |
| from chromadb.api.client import SharedSystemClient | |
| SharedSystemClient.clear_system_cache() | |
| monkeypatch.setattr( | |
| chromadb, "PersistentClient", lambda *a, **k: chromadb.EphemeralClient() | |
| ) | |
| from app.rag.chromadb_client import ChromaDBClient | |
| return ChromaDBClient(embedder=_FakeEmbedder()) | |
| def _seed(db, project_id="proj-1"): | |
| db.upsert( | |
| "library", | |
| ["alpha doc", "beta doc"], | |
| [ | |
| {"source": "paper-a", "document_id": "doc-a", "project_id": project_id}, | |
| {"source": "paper-b", "document_id": "doc-b", "project_id": project_id}, | |
| ], | |
| ["a", "b"], | |
| ) | |
| def test_successful_retrieval_outcome(db): | |
| _seed(db) | |
| q = db.embedder.embed(["x"])[0] | |
| outcome = db.query_observed("library", q, n_results=2, where={"project_id": "proj-1"}) | |
| assert outcome.status == "success" | |
| assert outcome.retrieval_error is False | |
| assert outcome.empty_result is False | |
| assert len(outcome.candidates) == 2 | |
| def test_legitimate_empty_is_success_empty_not_error(db): | |
| _seed(db) | |
| q = db.embedder.embed(["x"])[0] | |
| outcome = db.query_observed( | |
| "library", q, n_results=2, where={"project_id": "no-such-project"} | |
| ) | |
| assert outcome.status == "success_empty" | |
| assert outcome.empty_result is True | |
| # An empty *successful* search is NOT an error. | |
| assert outcome.retrieval_error is False | |
| assert outcome.candidates == [] | |
| def test_missing_collection_is_error_fallback(db): | |
| q = db.embedder.embed(["x"])[0] | |
| outcome = db.query_observed("never-created", q, n_results=2) | |
| assert outcome.status == "error_fallback" | |
| assert outcome.retrieval_error is True | |
| assert outcome.error_type == "collection_unavailable" | |
| assert outcome.candidates == [] | |
| def test_vector_search_failure_is_error_fallback(db, monkeypatch): | |
| _seed(db) | |
| q = db.embedder.embed(["x"])[0] | |
| real_get = db._client.get_collection | |
| class _BoomCollection: | |
| def count(self): | |
| return 2 | |
| def query(self, **kwargs): | |
| raise RuntimeError("index corrupted") | |
| monkeypatch.setattr(db._client, "get_collection", lambda name: _BoomCollection()) | |
| outcome = db.query_observed("library", q, n_results=2) | |
| assert outcome.status == "error_fallback" | |
| assert outcome.retrieval_error is True | |
| assert outcome.error_type == "vector_search_failed" | |
| assert outcome.candidates == [] | |
| def test_query_delegates_and_stays_list_compatible(db): | |
| """The public query() must remain a plain list[dict] for every caller.""" | |
| _seed(db) | |
| q = db.embedder.embed(["x"])[0] | |
| result = db.query("library", q, n_results=2, where={"project_id": "proj-1"}) | |
| assert isinstance(result, list) | |
| assert all(isinstance(r, dict) and "text" in r for r in result) | |
| # Missing collection still yields a plain [], never raises. | |
| assert db.query("never-created", q, n_results=2) == [] | |