update prompt tests to check for essential data in prompts only

Author: bjoaquinc

tests/_mcp/server/prompts/test_errors_prompts.py

@@ -0,0 +1,67 @@
+import pytest
+
+pytest.importorskip("mcp", reason="MCP requires Python 3.10+")
+
+from unittest.mock import Mock
+
+from marimo._ai._tools.types import (
+    MarimoCellErrors,
+    MarimoErrorDetail,
+    MarimoNotebookInfo,
+)
+from marimo._mcp.server._prompts.prompts.errors import ErrorsSummary
+from marimo._types.ids import CellId_t, SessionId
+
+
+def test_errors_summary_includes_essential_data():
+    """Test that essential data (notebook info, cell IDs, error messages) are included."""
+    context = Mock()
+    context.get_active_sessions_internal.return_value = [
+        MarimoNotebookInfo(
+            name="notebook.py",
+            path="/path/to/notebook.py",
+            session_id=SessionId("session_1"),
+        )
+    ]
+    context.get_notebook_errors.return_value = [
+        MarimoCellErrors(
+            cell_id=CellId_t("cell_1"),
+            errors=[
+                MarimoErrorDetail(
+                    type="ValueError",
+                    message="invalid value",
+                    traceback=["line 1"],
+                )
+            ],
+        ),
+        MarimoCellErrors(
+            cell_id=CellId_t("cell_2"),
+            errors=[
+                MarimoErrorDetail(
+                    type="TypeError",
+                    message="wrong type",
+                    traceback=["line 2"],
+                )
+            ],
+        ),
+    ]
+
+    prompt = ErrorsSummary(context=context)
+    messages = prompt.handle()
+    text = "\n".join(
+        msg.content.text  # type: ignore[attr-defined]
+        for msg in messages
+        if hasattr(msg.content, "text")
+    )
+
+    # Essential data must be present
+    assert "notebook.py" in text
+    assert "session_1" in text
+    assert "/path/to/notebook.py" in text
+    assert "cell_1" in text
+    assert "cell_2" in text
+    assert "ValueError" in text
+    assert "invalid value" in text
+    assert "TypeError" in text
+    assert "wrong type" in text
+

tests/_mcp/server/prompts/test_notebooks_prompts.py

@@ -4,73 +4,37 @@
 
 from unittest.mock import Mock
 
+from marimo._ai._tools.types import MarimoNotebookInfo
 from marimo._mcp.server._prompts.prompts.notebooks import ActiveNotebooks
+from marimo._types.ids import SessionId
 
 
-def test_active_notebooks_metadata():
-    """Test that name and description are properly set."""
-    prompt = ActiveNotebooks(context=Mock())
-    assert prompt.name == "active_notebooks"
-    assert (
-        prompt.description
-        == "Get current active notebooks and their session IDs and file paths."
-    )
-
-
-def test_active_notebooks_no_sessions():
-    """Test output when no sessions are active."""
+def test_active_notebooks_includes_session_ids_and_paths():
+    """Test that essential data (session IDs and file paths) are included."""
     context = Mock()
-    context.get_active_sessions_internal.return_value = []
-
-    prompt = ActiveNotebooks(context=context)
-    messages = prompt.handle()
-
-    assert len(messages) == 1
-    assert messages[0].role == "user"
-    assert messages[0].content.type == "text"
-    assert (
-        "No active marimo notebook sessions found" in messages[0].content.text
-    )
-
-
-def test_active_notebooks_with_sessions():
-    """Test output with active sessions."""
-    context = Mock()
-
-    # Mock active session objects
-    active_session1 = Mock()
-    active_session1.session_id = "session_1"
-    active_session1.path = "/path/to/notebook.py"
-
-    active_session2 = Mock()
-    active_session2.session_id = "session_2"
-    active_session2.path = None
-
     context.get_active_sessions_internal.return_value = [
-        active_session1,
-        active_session2,
+        MarimoNotebookInfo(
+            name="notebook.py",
+            path="/path/to/notebook.py",
+            session_id=SessionId("session_1"),
+        ),
+        MarimoNotebookInfo(
+            name="other.py",
+            path="/other/path.py",
+            session_id=SessionId("session_2"),
+        ),
     ]
 
     prompt = ActiveNotebooks(context=context)
     messages = prompt.handle()
-
-    assert len(messages) == 3  # 2 sessions + 1 action message
-
-    # Check first message (with file path)
-    assert messages[0].role == "user"
-    assert messages[0].content.type == "text"
-    assert "session_1" in messages[0].content.text
-    assert "/path/to/notebook.py" in messages[0].content.text
-
-    # Check second message (without file path)
-    assert messages[1].role == "user"
-    assert messages[1].content.type == "text"
-    assert "session_2" in messages[1].content.text
-
-    # Check action message
-    assert messages[2].role == "user"
-    assert messages[2].content.type == "text"
-    assert (
-        "Use these session_ids when calling marimo MCP tools"
-        in messages[2].content.text
+    text = "\n".join(
+        msg.content.text  # type: ignore[attr-defined]
+        for msg in messages
+        if hasattr(msg.content, "text")
     )
+
+    # Essential data must be present
+    assert "session_1" in text
+    assert "session_2" in text
+    assert "/path/to/notebook.py" in text
+    assert "/other/path.py" in text