diff --git a/src/brainlayer/enrichment_controller.py b/src/brainlayer/enrichment_controller.py index cba0782a..b324cc27 100644 --- a/src/brainlayer/enrichment_controller.py +++ b/src/brainlayer/enrichment_controller.py @@ -81,7 +81,10 @@ def get_unsubmitted_export_files(*args, **kwargs): def _get_gemini_client(): """Create Gemini client. Uses regional endpoint when GOOGLE_CLOUD_REGION is set.""" - from google import genai + try: + from google import genai + except ImportError: + raise RuntimeError("google-genai package not installed. Install with: pip install google-genai") api_key = os.environ.get("GOOGLE_API_KEY") or os.environ.get("GOOGLE_GENERATIVE_AI_API_KEY") if not api_key: diff --git a/tests/test_enrichment_controller.py b/tests/test_enrichment_controller.py index 033b2db0..2ac0f867 100644 --- a/tests/test_enrichment_controller.py +++ b/tests/test_enrichment_controller.py @@ -7,7 +7,6 @@ Target: 35+ tests per A-R2 acceptance criteria. """ -import json from types import SimpleNamespace from unittest.mock import MagicMock @@ -251,19 +250,22 @@ def test_enrich_local_does_not_call_build_external_prompt(monkeypatch): external_prompt.assert_not_called() -def test_enrich_batch_uses_checkpoint_db_for_resume(monkeypatch): +def test_enrich_batch_returns_early_for_no_candidates(monkeypatch): from brainlayer import enrichment_controller as controller store = MagicMock() - ensure_mock = MagicMock() - monkeypatch.setattr(controller, "ensure_checkpoint_table", ensure_mock) - monkeypatch.setattr(controller, "get_pending_jobs", MagicMock(return_value=[])) - monkeypatch.setattr(controller, "get_unsubmitted_export_files", MagicMock(return_value=[])) + store.get_enrichment_candidates.return_value = [] + + # _get_gemini_client should never be reached when there are no candidates + monkeypatch.setattr( + controller, "_get_gemini_client", lambda: (_ for _ in ()).throw(AssertionError("should not be called")) + ) - result = controller.enrich_batch(store, phase="run", limit=100) + result = controller.enrich_batch(store, limit=100) - ensure_mock.assert_called_once_with(store) + store.get_enrichment_candidates.assert_called_once_with(limit=100, chunk_ids=None) assert result.mode == "batch" + assert result.enriched == 0 # ── Content-hash dedup tests ───────────────────────────────────────────────── @@ -434,7 +436,7 @@ def test_gemini_client_requires_api_key(monkeypatch): monkeypatch.delenv("GOOGLE_API_KEY", raising=False) monkeypatch.delenv("GOOGLE_GENERATIVE_AI_API_KEY", raising=False) - with pytest.raises(RuntimeError, match="not set"): + with pytest.raises(RuntimeError, match="not set|not installed"): _get_gemini_client() @@ -762,8 +764,10 @@ async def test_brain_enrich_handler_stats_mode(monkeypatch): result = await _brain_enrich(stats=True) assert result.isError is not True - data = json.loads(result.content[0].text) - assert "total_chunks" in data + text = result.content[0].text + # _enrich_stats returns formatted text with box-drawing chars, not JSON + assert "Total:" in text + assert "Enriched:" in text @pytest.mark.asyncio @@ -777,51 +781,47 @@ async def test_enrich_stats_returns_correct_structure(): store._read_cursor.return_value = cursor result = await _enrich_stats(store) - data = json.loads(result.content[0].text) + text = result.content[0].text - assert data["total_chunks"] == 1000 - assert data["enriched"] == 600 - assert data["unenriched_eligible"] == 350 - assert data["skipped_too_short"] == 50 - assert data["enriched_pct"] == 60.0 - assert data["enriched_last_24h"] == 20 + # _enrich_stats returns formatted text lines, not JSON + assert "Total: 1,000" in text + assert "Enriched: 600" in text + assert "(60.0%)" in text + assert "Remaining: 350" in text + assert "Skipped: 50" in text + assert "Last 24h: 20" in text # ── Batch mode tests ───────────────────────────────────────────────────────── -def test_enrich_batch_poll_phase_only_checks_pending(monkeypatch): +def test_enrich_batch_processes_candidates_with_gemini(monkeypatch): from brainlayer import enrichment_controller as controller store = MagicMock() - monkeypatch.setattr(controller, "ensure_checkpoint_table", MagicMock()) - pending_mock = MagicMock(return_value=[{"id": "job1"}]) - monkeypatch.setattr(controller, "get_pending_jobs", pending_mock) - export_mock = MagicMock(return_value=[]) - monkeypatch.setattr(controller, "get_unsubmitted_export_files", export_mock) + store.get_enrichment_candidates.return_value = [_candidate("c1"), _candidate("c2")] + _patch_realtime_deps(monkeypatch, controller, store) - result = controller.enrich_batch(store, phase="poll") + result = controller.enrich_batch(store, limit=10) - pending_mock.assert_called_once() - export_mock.assert_not_called() - assert result.attempted == 1 + assert result.mode == "batch" + assert result.attempted == 2 + assert result.enriched == 2 -def test_enrich_batch_submit_phase_only_checks_exports(monkeypatch): +def test_enrich_batch_graceful_when_no_gemini_key(monkeypatch): from brainlayer import enrichment_controller as controller store = MagicMock() - monkeypatch.setattr(controller, "ensure_checkpoint_table", MagicMock()) - pending_mock = MagicMock(return_value=[]) - monkeypatch.setattr(controller, "get_pending_jobs", pending_mock) - export_mock = MagicMock(return_value=["f1.jsonl", "f2.jsonl"]) - monkeypatch.setattr(controller, "get_unsubmitted_export_files", export_mock) + store.get_enrichment_candidates.return_value = [_candidate()] - result = controller.enrich_batch(store, phase="submit") + monkeypatch.setattr(controller, "_get_gemini_client", lambda: (_ for _ in ()).throw(RuntimeError("no key"))) - pending_mock.assert_not_called() - export_mock.assert_called_once() - assert result.attempted == 2 + result = controller.enrich_batch(store, limit=5) + + assert result.mode == "batch" + assert result.enriched == 0 + assert any("No Gemini client" in e for e in result.errors) # ── Realtime chunk_ids filter test ──────────────────────────────────────────── diff --git a/tests/test_phase3_digest.py b/tests/test_phase3_digest.py index 62c1f7b5..29f71611 100644 --- a/tests/test_phase3_digest.py +++ b/tests/test_phase3_digest.py @@ -290,9 +290,9 @@ def test_brain_digest_description_teaches_routing(): assert "on schedule for backfill" in desc assert "faceted tags" in desc assert "sanitizes pii" in desc - assert "realtime" in desc - assert "batch" in desc - assert "local" in desc + assert "digest" in desc + assert "connect" in desc + assert "enrich" in desc # --- Task 4: brain_entity MCP tool --- diff --git a/tests/test_phase6_critical.py b/tests/test_phase6_critical.py index 04937cfc..c232e96f 100644 --- a/tests/test_phase6_critical.py +++ b/tests/test_phase6_critical.py @@ -394,10 +394,12 @@ def test_compact_format_output_size(self): # importance is now included in compact format for relevance visibility assert "importance" in compact + # tags is now included in compact format for filtering + assert "tags" in compact + # Verbose fields dropped for dropped in ( "content_type", - "tags", "intent", "source_file", "session_summary", diff --git a/tests/test_think_recall_integration.py b/tests/test_think_recall_integration.py index 1a24890a..647cb72e 100644 --- a/tests/test_think_recall_integration.py +++ b/tests/test_think_recall_integration.py @@ -246,13 +246,13 @@ class TestMCPToolCount: """Verify MCP server has correct tool count.""" def test_tool_count(self): - """MCP server should have 11 tools: search, store, recall, digest, entity, get_person, update, expand, tags, supersede, archive.""" + """MCP server should have 12 tools: search, store, recall, digest, entity, get_person, update, expand, tags, supersede, archive, enrich.""" import asyncio from brainlayer.mcp import list_tools tools = asyncio.run(list_tools()) - assert len(tools) == 11 + assert len(tools) == 12 def test_consolidated_tools_registered(self): """brain_search, brain_store, brain_recall are registered."""