Add worker integration test with st.write_stream() to check if the coroutine is awaited

packages/kernel/src/worker-runtime.test.ts

@@ -81,6 +81,7 @@ const TEST_SOURCES: {
   files: Record<string, string>;
   entrypoint: string;
   requirements?: string[];
+  additionalAppTestCode?: string;
 }[] = [
   {
     entrypoint: "data.column_config.py",
@@ -118,6 +119,20 @@ const TEST_SOURCES: {
       ),
     },
   },
+  {
+    // Checks if the async customization of st.write_stream() and runtime AST manipulation works.
+    entrypoint: "text.write_stream.py",
+    files: {
+      "text.write_stream.py": path.resolve(
+        __dirname,
+        "../../sharing-editor/public/samples/011_component_gallery/pages/text.write_stream.py",
+      ),
+    },
+    additionalAppTestCode: `
+at.button[0].click()
+await at.run(timeout=10)
+`,
+  },
 ];
 
 suite("Worker intergration test running an app", async () => {
@@ -146,21 +161,38 @@ suite("Worker intergration test running an app", async () => {
           requirements: testSource.requirements,
         });
 
+        pyodide.globals.set(
+          "__additionalAppTestCode__",
+          testSource.additionalAppTestCode,
+        );
+
         // The code above setting up the worker env is good enough to check if the worker is set up correctly,
         // but it doesn't check the error occurred inside the Streamlit app running in the worker.
         // So, we use the code below to test if the Streamlit app runs without any error.
         await pyodide.runPythonAsync(`
+import ast
+import asyncio
+import warnings
 from streamlit.testing.v1 import AppTest
 
 at = AppTest.from_file("${testSource.entrypoint}")
 
-await at.run()
+with warnings.catch_warnings(record=True) as w:  # Test warning messages. Ref: https://docs.python.org/3/library/warnings.html#testing-warnings
+    warnings.simplefilter("always")
+    warnings.filterwarnings("ignore", message=r"\\n?Pyarrow will become a required dependency of pandas in the next major release of pandas")  # Ignore PyArrow version warning from Pandas (https://github.com/pandas-dev/pandas/blob/v2.2.0/pandas/__init__.py#L221-L231)
+
+    await at.run()
+
+    if __additionalAppTestCode__:
+        bytecode = compile(__additionalAppTestCode__, "<string>", "exec", flags=ast.PyCF_ALLOW_TOP_LEVEL_AWAIT)
+        await eval(bytecode)
 
 assert not at.exception, f"Exception occurred: {at.exception}"
+assert len(w) == 0, f"Warning occurred: {w[0].message if w else None}"
         `);
       },
       {
-        timeout: 30 * 1000,
+        timeout: 60 * 1000,
       },
     );
   }