update tests to reflect changes

Author: bjoaquinc

tests/_ai/tools/test_base.py

@@ -8,6 +8,7 @@
 
 from marimo._ai._tools.base import ToolBase, ToolContext
 from marimo._ai._tools.utils.exceptions import ToolExecutionError
+from marimo._messaging import msgspec_encoder
 
 
 @dataclass
@@ -125,3 +126,93 @@ def test_as_backend_tool() -> None:
     is_valid, msg = validator({"invalid": "field"})
     assert is_valid is False
     assert "Invalid arguments" in msg
+
+# test ToolContext methods
+
+def test_get_notebook_errors_orders_by_cell_manager():
+    """Test errors follow cell_manager order, not alphabetical."""
+    from unittest.mock import Mock
+    from marimo._messaging.cell_output import CellChannel
+    from marimo._types.ids import CellId_t, SessionId
+
+    context = ToolContext()
+    
+    # Mock error cell_op
+    error_op = Mock()
+    error_op.output = Mock()
+    error_op.output.channel = CellChannel.MARIMO_ERROR
+    error_op.output.data = [{"type": "Error", "msg": "test", "traceback": []}]
+    error_op.console = None
+    
+    # Mock session with cells c1, c2, c3
+    session = Mock()
+    session.session_view.cell_operations = {
+        CellId_t("c1"): error_op,
+        CellId_t("c2"): error_op,
+        CellId_t("c3"): error_op,
+    }
+    
+    # Cell manager returns in order: c3, c2, c1 (not alphabetical)
+    cell_data = [Mock(cell_id=CellId_t("c3")), Mock(cell_id=CellId_t("c2")), Mock(cell_id=CellId_t("c1"))]
+    session.app_file_manager.app.cell_manager.cell_data.return_value = cell_data
+    
+    context.get_session = Mock(return_value=session)
+    
+    errors = context.get_notebook_errors(SessionId("test"), include_stderr=False)
+    
+    # Should be c3, c2, c1 (not c1, c2, c3)
+    assert errors[0].cell_id == CellId_t("c3")
+    assert errors[1].cell_id == CellId_t("c2")
+    assert errors[2].cell_id == CellId_t("c1")
+
+
+def test_get_cell_errors_extracts_from_output():
+    """Test get_cell_errors extracts error details from cell output."""
+    from unittest.mock import Mock
+    from marimo._messaging.cell_output import CellChannel
+    from marimo._types.ids import CellId_t, SessionId
+    
+    context = ToolContext()
+    
+    # Mock cell_op with error
+    cell_op = Mock()
+    cell_op.output = Mock()
+    cell_op.output.channel = CellChannel.MARIMO_ERROR
+    cell_op.output.data = [
+        {"type": "ValueError", "msg": "bad value", "traceback": ["line 1"]}
+    ]
+    
+    errors = context.get_cell_errors(SessionId("test"), CellId_t("c1"), maybe_cell_op=cell_op)
+    
+    assert len(errors) == 1
+    assert errors[0].type == "ValueError"
+    assert errors[0].message == "bad value"
+    assert errors[0].traceback == ["line 1"]
+
+
+def test_get_cell_console_outputs_separates_stdout_stderr():
+    """Test get_cell_console_outputs separates stdout and stderr."""
+    from unittest.mock import Mock
+    from marimo._messaging.cell_output import CellChannel
+    
+    context = ToolContext()
+    
+    # Mock cell_op with stdout and stderr
+    stdout_output = Mock()
+    stdout_output.channel = CellChannel.STDOUT
+    stdout_output.data = "hello"
+    
+    stderr_output = Mock()
+    stderr_output.channel = CellChannel.STDERR
+    stderr_output.data = "warning"
+    
+    cell_op = Mock()
+    cell_op.console = [stdout_output, stderr_output]
+    
+    result = context.get_cell_console_outputs(cell_op)
+    
+    assert len(result.stdout) == 1
+    assert "hello" in result.stdout[0]
+    assert len(result.stderr) == 1
+    assert "warning" in result.stderr[0]
+

tests/_ai/tools/tools/test_cells.py

@@ -190,25 +190,3 @@ def test_get_visual_output_no_output():
     visual_output, mimetype = tool._get_visual_output(cell_op)  # type: ignore[arg-type]
     assert visual_output is None
     assert mimetype is None
-
-
-def test_get_console_outputs_with_stdout_stderr():
-    tool = GetCellOutputs(ToolContext())
-    console = [
-        MockConsoleOutput(CellChannel.STDOUT, "hello"),
-        MockConsoleOutput(CellChannel.STDERR, "warning"),
-    ]
-    cell_op = MockCellOp(console=console)
-
-    stdout, stderr = tool._get_console_outputs(cell_op)  # type: ignore[arg-type]
-    assert stdout == ["hello"]
-    assert stderr == ["warning"]
-
-
-def test_get_console_outputs_no_console():
-    tool = GetCellOutputs(ToolContext())
-    cell_op = MockCellOp(console=None)
-
-    stdout, stderr = tool._get_console_outputs(cell_op)  # type: ignore[arg-type]
-    assert stdout == []
-    assert stderr == []

tests/_ai/tools/tools/test_errors_tool.py

@@ -119,65 +119,6 @@ def test_get_notebook_errors_marimo_error_only(mock_context: Mock) -> None:
     assert "get_cell_runtime_data" in result.next_steps[0]
 
 
-def test_get_notebook_errors_stderr_only(mock_context: Mock) -> None:
-    """Test get_notebook_errors with STDERR only."""
-    stderr_errors = [
-        MarimoCellErrors(
-            cell_id=CellId_t("c2"),
-            errors=[
-                MarimoErrorDetail(
-                    type="STDERR",
-                    message="warning message",
-                    traceback=[],
-                )
-            ],
-        )
-    ]
-    mock_context.get_notebook_errors.return_value = stderr_errors
-
-    tool = GetNotebookErrors(ToolContext())
-    tool.context = mock_context
-
-    result = tool.handle(GetNotebookErrorsArgs(session_id=SessionId("s1")))
-
-    assert result.has_errors is False
-    assert result.total_errors == 0
-    assert result.total_cells_with_errors == 1
-    assert result.cells[0].errors[0].type == "STDERR"
-
-
-def test_get_notebook_errors_mixed_errors(mock_context: Mock) -> None:
-    """Test get_notebook_errors with both MARIMO_ERROR and STDERR."""
-    mixed_errors = [
-        MarimoCellErrors(
-            cell_id=CellId_t("c1"),
-            errors=[
-                MarimoErrorDetail(
-                    type="ValueError",
-                    message="bad value",
-                    traceback=["line 1"],
-                ),
-                MarimoErrorDetail(
-                    type="STDERR",
-                    message="warn",
-                    traceback=[],
-                ),
-            ],
-        )
-    ]
-    mock_context.get_notebook_errors.return_value = mixed_errors
-
-    tool = GetNotebookErrors(ToolContext())
-    tool.context = mock_context
-
-    result = tool.handle(GetNotebookErrorsArgs(session_id=SessionId("s1")))
-
-    assert result.has_errors is True
-    assert result.total_errors == 1
-    assert result.total_cells_with_errors == 1
-    assert len(result.cells[0].errors) == 2
-
-
 def test_get_notebook_errors_multiple_cells(mock_context: Mock) -> None:
     """Test get_notebook_errors with errors in multiple cells."""
     multiple_errors = [
@@ -190,6 +131,7 @@ def test_get_notebook_errors_multiple_cells(mock_context: Mock) -> None:
                     traceback=[],
                 )
             ],
+            stderr=[],
         ),
         MarimoCellErrors(
             cell_id=CellId_t("c2"),
@@ -198,18 +140,14 @@ def test_get_notebook_errors_multiple_cells(mock_context: Mock) -> None:
                     type="TypeError",
                     message="error in c2",
                     traceback=[],
-                )
-            ],
-        ),
-        MarimoCellErrors(
-            cell_id=CellId_t("c3"),
-            errors=[
+                ),
                 MarimoErrorDetail(
-                    type="STDERR",
-                    message="stderr in c3",
+                    type="ValueError",
+                    message="error in c2",
                     traceback=[],
                 )
             ],
+            stderr=[],
         ),
     ]
     mock_context.get_notebook_errors.return_value = multiple_errors
@@ -220,9 +158,9 @@ def test_get_notebook_errors_multiple_cells(mock_context: Mock) -> None:
     result = tool.handle(GetNotebookErrorsArgs(session_id=SessionId("s1")))
 
     assert result.has_errors is True
-    assert result.total_errors == 2
-    assert result.total_cells_with_errors == 3
-    assert len(result.cells) == 3
+    assert result.total_errors == 3
+    assert result.total_cells_with_errors == 2
+    assert len(result.cells) == 2
 
 
 def test_get_notebook_errors_respects_session_id(mock_context: Mock) -> None: