From 392842ac79a0ffbc01dadd35d1d3d1569f196635 Mon Sep 17 00:00:00 2001 From: Etan Joseph Heyman Date: Wed, 1 Apr 2026 03:52:57 +0300 Subject: [PATCH 1/7] fix: update stale enrichment controller tests to match current API 5 tests were failing: 3 batch tests tested an old checkpoint/phase API that enrich_batch no longer uses, and 2 stats tests expected JSON output when the handler now returns formatted text. The batch tests also caused ModuleNotFoundError for google-genai in CI since they hit _get_gemini_client without mocking it. All tests now mock _get_gemini_client properly and assert against actual current behavior. Co-Authored-By: Claude Opus 4.6 (1M context) --- tests/test_enrichment_controller.py | 74 ++++++++++++++--------------- 1 file changed, 37 insertions(+), 37 deletions(-) diff --git a/tests/test_enrichment_controller.py b/tests/test_enrichment_controller.py index 033b2db0..b1b47f08 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 ───────────────────────────────────────────────── @@ -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 ──────────────────────────────────────────── From ab42e2c2553ff90d9279600f005ea2f5aaf90a63 Mon Sep 17 00:00:00 2001 From: Etan Joseph Heyman Date: Wed, 1 Apr 2026 03:57:43 +0300 Subject: [PATCH 2/7] fix: handle missing google-genai package gracefully in _get_gemini_client Wrap `from google import genai` in try/except ImportError so CI environments without google-genai get a clear RuntimeError instead of ModuleNotFoundError. Update test to accept either "not set" (no API key) or "not installed" (no package). Co-Authored-By: Claude Opus 4.6 (1M context) --- src/brainlayer/enrichment_controller.py | 7 ++++++- tests/test_enrichment_controller.py | 2 +- 2 files changed, 7 insertions(+), 2 deletions(-) diff --git a/src/brainlayer/enrichment_controller.py b/src/brainlayer/enrichment_controller.py index cba0782a..8d6063a9 100644 --- a/src/brainlayer/enrichment_controller.py +++ b/src/brainlayer/enrichment_controller.py @@ -81,7 +81,12 @@ 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 b1b47f08..2ac0f867 100644 --- a/tests/test_enrichment_controller.py +++ b/tests/test_enrichment_controller.py @@ -436,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() From 019c63b18e616451cea3f33e42be9b214d6fa0d4 Mon Sep 17 00:00:00 2001 From: Etan Joseph Heyman Date: Wed, 1 Apr 2026 04:00:08 +0300 Subject: [PATCH 3/7] style: fix ruff format on enrichment_controller.py Co-Authored-By: Claude Opus 4.6 (1M context) --- src/brainlayer/enrichment_controller.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/src/brainlayer/enrichment_controller.py b/src/brainlayer/enrichment_controller.py index 8d6063a9..b324cc27 100644 --- a/src/brainlayer/enrichment_controller.py +++ b/src/brainlayer/enrichment_controller.py @@ -84,9 +84,7 @@ def _get_gemini_client(): try: from google import genai except ImportError: - raise RuntimeError( - "google-genai package not installed. Install with: pip install google-genai" - ) + 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: From 5f5e446f12b64892c5197651d87fd33b8cf5ceba Mon Sep 17 00:00:00 2001 From: Etan Joseph Heyman Date: Wed, 1 Apr 2026 04:04:22 +0300 Subject: [PATCH 4/7] fix: update digest description test for current mode names The brain_digest description now uses "enrich" (not "batch") and "backfill" (not "local") to describe the enrichment modes. Co-Authored-By: Claude Opus 4.6 (1M context) --- tests/test_phase3_digest.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/test_phase3_digest.py b/tests/test_phase3_digest.py index 62c1f7b5..da86ad93 100644 --- a/tests/test_phase3_digest.py +++ b/tests/test_phase3_digest.py @@ -291,8 +291,8 @@ def test_brain_digest_description_teaches_routing(): assert "faceted tags" in desc assert "sanitizes pii" in desc assert "realtime" in desc - assert "batch" in desc - assert "local" in desc + assert "enrich" in desc + assert "backfill" in desc # --- Task 4: brain_entity MCP tool --- From 4e2bd762d469895dde5ee1fbf24e5243cef93e38 Mon Sep 17 00:00:00 2001 From: Etan Joseph Heyman Date: Wed, 1 Apr 2026 04:06:46 +0300 Subject: [PATCH 5/7] fix: update stale brain_digest mode assertions (digest/connect/enrich) The brain_digest tool description now lists modes as digest, connect, and enrich instead of the old realtime, batch, local. Update the test_brain_digest_description_teaches_routing test to match. Co-Authored-By: Claude Opus 4.6 (1M context) --- tests/test_phase3_digest.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/test_phase3_digest.py b/tests/test_phase3_digest.py index da86ad93..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 "digest" in desc + assert "connect" in desc assert "enrich" in desc - assert "backfill" in desc # --- Task 4: brain_entity MCP tool --- From fa9dfaedfa5ee6638696da4b037402a5a39298c7 Mon Sep 17 00:00:00 2001 From: Etan Joseph Heyman Date: Wed, 1 Apr 2026 04:12:18 +0300 Subject: [PATCH 6/7] =?UTF-8?q?fix:=20update=20compact=20format=20test=20?= =?UTF-8?q?=E2=80=94=20tags=20now=20included=20in=20compact=20results?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit _build_compact_result() includes tags for filtering, but the test still listed tags in the "dropped fields" assertion. Updated test to match the actual compact output contract. Co-Authored-By: Claude Opus 4.6 (1M context) --- tests/test_phase6_critical.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) 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", From 432f4a592a856c779df7b29450d5e3b1ef4e77c1 Mon Sep 17 00:00:00 2001 From: Etan Joseph Heyman Date: Wed, 1 Apr 2026 04:16:20 +0300 Subject: [PATCH 7/7] fix: update tool count test from 11 to 12 (enrich tool added) Co-Authored-By: Claude Opus 4.6 (1M context) --- tests/test_think_recall_integration.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) 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."""