Additional tests

haiku_rag_slim/haiku/rag/agents/rlm/docker_sandbox.py

@@ -24,7 +24,7 @@ class SandboxResult:
     success: bool
 
 
-class DockerSandbox:
+class DockerSandbox:  # pragma: no cover
     """Execute code in a persistent Docker container.
 
     Use as an async context manager to manage container lifecycle:

tests/test_app.py

@@ -669,3 +669,57 @@ def test_migrate_closes_store_on_exception(tmp_path):
             app.migrate()
 
         mock_store.close.assert_called_once()
+
+
+@pytest.mark.asyncio
+async def test_rlm(app: HaikuRAGApp, monkeypatch):
+    """Test rlm method calls client.rlm and prints results."""
+    from haiku.rag.agents.rlm.models import RLMResult
+
+    mock_result = RLMResult(
+        answer="The total is 42.",
+        program="result = sum(values)\nprint(result)",
+    )
+
+    mock_client = AsyncMock()
+    mock_client.rlm = AsyncMock(return_value=mock_result)
+    mock_client.__aenter__.return_value = mock_client
+
+    mock_print = MagicMock()
+    monkeypatch.setattr(app.console, "print", mock_print)
+
+    with patch("haiku.rag.app.HaikuRAG", return_value=mock_client):
+        await app.rlm("What is the total?")
+
+    mock_client.rlm.assert_called_once_with(
+        "What is the total?", documents=None, filter=None
+    )
+    calls = [str(c) for c in mock_print.call_args_list]
+    assert any("Question" in c for c in calls)
+    assert any("Program" in c for c in calls)
+    assert any("Answer" in c for c in calls)
+
+
+@pytest.mark.asyncio
+async def test_rlm_with_document_and_filter(app: HaikuRAGApp, monkeypatch):
+    """Test rlm method passes document and filter to client."""
+    from haiku.rag.agents.rlm.models import RLMResult
+
+    mock_result = RLMResult(
+        answer="Answer with filter",
+        program="print('filtered')",
+    )
+
+    mock_client = AsyncMock()
+    mock_client.rlm = AsyncMock(return_value=mock_result)
+    mock_client.__aenter__.return_value = mock_client
+
+    mock_print = MagicMock()
+    monkeypatch.setattr(app.console, "print", mock_print)
+
+    with patch("haiku.rag.app.HaikuRAG", return_value=mock_client):
+        await app.rlm("What is it?", document="doc-123", filter="uri LIKE '%test%'")
+
+    mock_client.rlm.assert_called_once_with(
+        "What is it?", documents=["doc-123"], filter="uri LIKE '%test%'"
+    )

tests/test_mcp.py

@@ -307,3 +307,71 @@ async def test_mcp_research_question():
             assert result.title == "Research Title"
             assert result.executive_summary == "Summary"
             mock_graph.run.assert_called_once()
+
+
+@pytest.mark.asyncio
+async def test_mcp_rlm_question():
+    """Test rlm_question tool is properly wired."""
+    with tempfile.TemporaryDirectory() as temp_dir:
+        db_path = Path(temp_dir) / "test.lancedb"
+        mcp = create_mcp_server(db_path)
+
+        from haiku.rag.agents.rlm.models import RLMResult
+
+        mock_result = RLMResult(
+            answer="The total is 42.",
+            program="result = sum(values)\nprint(result)",
+        )
+
+        with patch("haiku.rag.mcp.HaikuRAG") as mock_rag_class:
+            mock_rag = AsyncMock()
+            mock_rag.rlm = AsyncMock(return_value=mock_result)
+            mock_rag_class.return_value.__aenter__ = AsyncMock(return_value=mock_rag)
+            mock_rag_class.return_value.__aexit__ = AsyncMock(return_value=None)
+
+            tools = await mcp.get_tools()
+            rlm_tool = next(t for t in tools.values() if t.name == "rlm_question")
+
+            result = await rlm_tool.fn(question="What is the total?")
+
+            assert result == "The total is 42."
+            mock_rag.rlm.assert_called_once_with(
+                "What is the total?", documents=None, filter=None
+            )
+
+
+@pytest.mark.asyncio
+async def test_mcp_rlm_question_with_document_and_filter():
+    """Test rlm_question tool with document and filter parameters."""
+    with tempfile.TemporaryDirectory() as temp_dir:
+        db_path = Path(temp_dir) / "test.lancedb"
+        mcp = create_mcp_server(db_path)
+
+        from haiku.rag.agents.rlm.models import RLMResult
+
+        mock_result = RLMResult(
+            answer="Filtered answer",
+            program="print('filtered')",
+        )
+
+        with patch("haiku.rag.mcp.HaikuRAG") as mock_rag_class:
+            mock_rag = AsyncMock()
+            mock_rag.rlm = AsyncMock(return_value=mock_result)
+            mock_rag_class.return_value.__aenter__ = AsyncMock(return_value=mock_rag)
+            mock_rag_class.return_value.__aexit__ = AsyncMock(return_value=None)
+
+            tools = await mcp.get_tools()
+            rlm_tool = next(t for t in tools.values() if t.name == "rlm_question")
+
+            result = await rlm_tool.fn(
+                question="Analyze this",
+                document="doc-123",
+                filter="uri LIKE '%test%'",
+            )
+
+            assert result == "Filtered answer"
+            mock_rag.rlm.assert_called_once_with(
+                "Analyze this",
+                documents=["doc-123"],
+                filter="uri LIKE '%test%'",
+            )