Capture stdout/stderr in server test and align completion timing

Update the language server test cell code to use stdout/stderr. Inserts
a short 20ms delay before setting the completion event to ensure we
capture the stdout/stderr, which are buffered and flushed every 10ms.

Ideally, we could flush prior to the "idle" cell-op for each cell run,
but that requires more complicated upstream change so this works for
now. We will do the same thing in the extension front end.

Author: manzt

pyproject.toml

@@ -54,4 +54,5 @@ ignore = [
   "SLF001",
   "PLC",
   "PLR",
+  "FIX",
 ]

tests/test_server.py

@@ -298,6 +298,13 @@ def __():
 @pytest.mark.asyncio
 async def test_simple_marimo_run(client: LanguageClient) -> None:
     """Test that we can collect marimo operations until cell reaches idle state."""
+    code = """\
+import sys
+
+print("hello, world")
+print("error message", file=sys.stderr)
+x = 42\
+"""
 
     client.notebook_document_did_open(
         lsp.DidOpenNotebookDocumentParams(
@@ -317,7 +324,7 @@ async def test_simple_marimo_run(client: LanguageClient) -> None:
                     uri="file:///exec_test.py#cell1",
                     language_id="python",
                     version=1,
-                    text="x = 42",
+                    text=code,
                 )
             ],
         ),
@@ -327,11 +334,14 @@ async def test_simple_marimo_run(client: LanguageClient) -> None:
     completion_event = asyncio.Event()
 
     @client.feature("marimo/operation")
-    def on_marimo_operation(params: Any) -> None:  # noqa: ANN401
+    async def on_marimo_operation(params: Any) -> None:  # noqa: ANN401
         # pygls dynamically makes an `Object` named tuple which makes snapshotting hard
         # we just convert to a regular dict here for snapshotting
         messages.append(to_dict(params))
         if params.op == "completed-run":
+            # FIXME: stdin/stdout are flushed every 10ms, so wait 20ms to ensure
+            # all related events. The frontend uses the same workaround.
+            await asyncio.sleep(0.02)
             completion_event.set()
 
     await client.workspace_execute_command_async(
@@ -341,7 +351,7 @@ def on_marimo_operation(params: Any) -> None:  # noqa: ANN401
                 {
                     "notebook_uri": "file:///exec_test.py",
                     "cell_ids": ["cell1"],
-                    "codes": ["x = 42"],
+                    "codes": [code],
                 }
             ],
         )
@@ -355,7 +365,15 @@ def on_marimo_operation(params: Any) -> None:  # noqa: ANN401
                 "op": "update-cell-codes",
                 "data": {
                     "cell_ids": ["cell1"],
-                    "codes": ["x = 42"],
+                    "codes": [
+                        """\
+import sys
+
+print("hello, world")
+print("error message", file=sys.stderr)
+x = 42\
+"""
+                    ],
                     "code_is_stale": False,
                 },
             },
@@ -369,7 +387,8 @@ def on_marimo_operation(params: Any) -> None:  # noqa: ANN401
                 "op": "variables",
                 "data": {
                     "variables": [
-                        {"name": "x", "declared_by": ["cell1"], "used_by": []}
+                        {"name": "sys", "declared_by": ["cell1"], "used_by": []},
+                        {"name": "x", "declared_by": ["cell1"], "used_by": []},
                     ]
                 },
             },
@@ -410,7 +429,10 @@ def on_marimo_operation(params: Any) -> None:  # noqa: ANN401
                 "notebookUri": "file:///exec_test.py",
                 "op": "variable-values",
                 "data": {
-                    "variables": [{"name": "x", "value": "42", "datatype": "int"}]
+                    "variables": [
+                        {"name": "sys", "value": "sys", "datatype": "module"},
+                        {"name": "x", "value": "42", "datatype": "int"},
+                    ]
                 },
             },
             {
@@ -446,6 +468,48 @@ def on_marimo_operation(params: Any) -> None:  # noqa: ANN401
                     "timestamp": IsFloat(),
                 },
             },
-            {"notebookUri": "file:///exec_test.py", "op": "completed-run", "data": {}},
+            {
+                "notebookUri": "file:///exec_test.py",
+                "op": "completed-run",
+                "data": {},
+            },
+            {
+                "notebookUri": "file:///exec_test.py",
+                "op": "cell-op",
+                "data": {
+                    "cell_id": "cell1",
+                    "output": None,
+                    "console": {
+                        "channel": "stdout",
+                        "mimetype": "text/plain",
+                        "data": "hello, world\n",
+                        "timestamp": IsFloat(),
+                    },
+                    "status": None,
+                    "stale_inputs": None,
+                    "run_id": None,
+                    "serialization": None,
+                    "timestamp": IsFloat(),
+                },
+            },
+            {
+                "notebookUri": "file:///exec_test.py",
+                "op": "cell-op",
+                "data": {
+                    "cell_id": "cell1",
+                    "output": None,
+                    "console": {
+                        "channel": "stderr",
+                        "mimetype": "text/plain",
+                        "data": "error message\n",
+                        "timestamp": IsFloat(),
+                    },
+                    "status": None,
+                    "stale_inputs": None,
+                    "run_id": None,
+                    "serialization": None,
+                    "timestamp": IsFloat(),
+                },
+            },
         ]
     )