test step writes

Author: adrianlyjak

examples/dbos_durability/.gitignore

@@ -1,2 +1,2 @@
-.last_run.json
+.last_run_id
 *.sqlite3

examples/dbos_durability/.last_run_id

@@ -732,3 +732,201 @@ def test_concurrent_workers_stress(test_db_path: Path, iteration: int) -> None:
         f"Iteration {iteration} failed.\nstdout: {result.stdout}\nstderr: {result.stderr}"
     )
     assert_no_determinism_errors(result)
+
+
+# =============================================================================
+# Test 7: Streaming stress test - many write_event_to_stream and send_event calls
+# This tests determinism when a step fires many background tasks for streaming
+# and internal event sending, then gets interrupted and replayed.
+# =============================================================================
+
+STREAMING_STRESS_WORKFLOW_CODE = textwrap.dedent("""
+    import asyncio
+
+    class ProgressEvent(Event):
+        progress: int = Field(default=0)
+
+    class WorkItem(Event):
+        item_id: int = Field(default=0)
+
+    class WorkDone(Event):
+        item_id: int = Field(default=0)
+        total_processed: int = Field(default=0)
+
+    class FanOutComplete(Event):
+        pass
+
+    class StreamingStressWorkflow(Workflow):
+        @step
+        async def fan_out(self, ctx: Context, ev: StartEvent) -> FanOutComplete:
+            # Fire many stream writes and internal events concurrently
+            # This creates many background tasks that call DBOS operations
+            for i in range(15):
+                ctx.write_event_to_stream(ProgressEvent(progress=i))
+                ctx.send_event(WorkItem(item_id=i))
+            print(f"STEP:fan_out:dispatched_15_items", flush=True)
+            # Write completion signal to stream for interrupt tests
+            ctx.write_event_to_stream(ProgressEvent(progress=999))
+            return FanOutComplete()
+
+        @step(num_workers=4)
+        async def process_work(self, ctx: Context, ev: WorkItem) -> WorkDone:
+            # Each worker also writes to stream, creating more concurrent DBOS ops
+            await asyncio.sleep(0.01)  # Small delay to increase interleaving
+            ctx.write_event_to_stream(ProgressEvent(progress=100 + ev.item_id))
+            print(f"STEP:process_work:{ev.item_id}:complete", flush=True)
+            return WorkDone(item_id=ev.item_id)
+
+        @step
+        async def after_fanout(self, ctx: Context, ev: FanOutComplete) -> None:
+            # Consume FanOutComplete, don't trigger anything
+            print("STEP:after_fanout:complete", flush=True)
+            return None
+
+        @step
+        async def collect(self, ctx: Context, ev: WorkDone) -> StopEvent:
+            # First WorkDone ends the workflow
+            print(f"STEP:collect:{ev.item_id}:complete", flush=True)
+            return StopEvent(result={"first_done": ev.item_id})
+""")
+
+
+def test_streaming_stress_determinism(test_db_path: Path) -> None:
+    """Test determinism with many concurrent stream writes and send_event calls."""
+    run_id = "test-streaming-stress-001"
+    db_url = f"sqlite+pysqlite:///{test_db_path}?check_same_thread=false"
+
+    run_main = textwrap.dedent(f'''
+        wf = StreamingStressWorkflow(runtime=runtime)
+        runtime.launch()
+        try:
+            ctx = Context(wf)
+            handler = ctx._workflow_run(wf, StartEvent(), run_id="{run_id}")
+            result = await handler
+            print(f"RESULT:{{result}}", flush=True)
+            print("SUCCESS", flush=True)
+        except Exception as e:
+            print(f"ERROR:{{type(e).__name__}}:{{e}}", flush=True)
+            raise
+        finally:
+            runtime.destroy()
+    ''')
+
+    print("\n=== Running streaming stress workflow ===")
+    result = run_workflow_script(
+        make_script(
+            STREAMING_STRESS_WORKFLOW_CODE, run_main, db_url, "test-streaming-stress"
+        )
+    )
+    print(f"stdout: {result.stdout}")
+    print(f"stderr: {result.stderr}")
+
+    assert "SUCCESS" in result.stdout, (
+        f"Should complete successfully.\nstdout: {result.stdout}\nstderr: {result.stderr}"
+    )
+    assert_no_determinism_errors(result)
+
+
+def test_streaming_interrupt_resume(test_db_path: Path) -> None:
+    """Test interrupt/resume with many concurrent stream writes in flight."""
+    run_id = "test-streaming-interrupt-001"
+    db_url = f"sqlite+pysqlite:///{test_db_path}?check_same_thread=false"
+
+    # Start workflow and interrupt after seeing the ProgressEvent with progress=999
+    # (the one returned by fan_out). This ensures all the send_event and
+    # write_event_to_stream calls have been made but workers are still running.
+    start_main = textwrap.dedent(f'''
+        wf = StreamingStressWorkflow(runtime=runtime)
+        runtime.launch()
+        ctx = Context(wf)
+        handler = ctx._workflow_run(wf, StartEvent(), run_id="{run_id}")
+        async for event in handler.stream_events():
+            print(f"EVENT:{{type(event).__name__}}:{{getattr(event, 'progress', getattr(event, 'item_id', ''))}}", flush=True)
+            # Interrupt after fan_out completes (it emits progress=999)
+            if isinstance(event, ProgressEvent) and event.progress == 999:
+                print("INTERRUPTING_AFTER_FANOUT", flush=True)
+                import os
+                os._exit(0)
+    ''')
+
+    # Resume and complete
+    resume_main = textwrap.dedent(f'''
+        wf = StreamingStressWorkflow(runtime=runtime)
+        runtime.launch()
+        try:
+            ctx = Context(wf)
+            handler = ctx._workflow_run(wf, StartEvent(), run_id="{run_id}")
+            result = await handler
+            print(f"RESULT:{{result}}", flush=True)
+            print("SUCCESS", flush=True)
+        except Exception as e:
+            print(f"ERROR:{{type(e).__name__}}:{{e}}", flush=True)
+            raise
+        finally:
+            runtime.destroy()
+    ''')
+
+    print("\n=== Starting streaming workflow (will interrupt after fan_out) ===")
+    result1 = run_workflow_script(
+        make_script(
+            STREAMING_STRESS_WORKFLOW_CODE, start_main, db_url, "test-streaming-int"
+        )
+    )
+    print(f"stdout: {result1.stdout}")
+    print(f"stderr: {result1.stderr}")
+
+    assert "STEP:fan_out:dispatched_15_items" in result1.stdout, (
+        "Fan out should complete"
+    )
+    assert "INTERRUPTING_AFTER_FANOUT" in result1.stdout, "Should have interrupted"
+
+    print("\n=== Resuming streaming workflow ===")
+    result2 = run_workflow_script(
+        make_script(
+            STREAMING_STRESS_WORKFLOW_CODE, resume_main, db_url, "test-streaming-int"
+        )
+    )
+    print(f"stdout: {result2.stdout}")
+    print(f"stderr: {result2.stderr}")
+
+    assert_no_determinism_errors(result2)
+    assert "SUCCESS" in result2.stdout or result2.returncode == 0, (
+        f"Resume should succeed.\nstdout: {result2.stdout}\nstderr: {result2.stderr}"
+    )
+
+
+@pytest.mark.parametrize("iteration", range(5))
+def test_streaming_stress_repeated(test_db_path: Path, iteration: int) -> None:
+    """Stress test streaming - run 5 times to catch timing issues."""
+    run_id = f"test-streaming-repeated-{iteration}"
+    db_url = f"sqlite+pysqlite:///{test_db_path}?check_same_thread=false"
+
+    run_main = textwrap.dedent(f'''
+        wf = StreamingStressWorkflow(runtime=runtime)
+        runtime.launch()
+        try:
+            ctx = Context(wf)
+            handler = ctx._workflow_run(wf, StartEvent(), run_id="{run_id}")
+            result = await handler
+            print(f"RESULT:{{result}}", flush=True)
+            print("SUCCESS", flush=True)
+        except Exception as e:
+            print(f"ERROR:{{type(e).__name__}}:{{e}}", flush=True)
+            raise
+        finally:
+            runtime.destroy()
+    ''')
+
+    result = run_workflow_script(
+        make_script(
+            STREAMING_STRESS_WORKFLOW_CODE,
+            run_main,
+            db_url,
+            f"streaming-stress-{iteration}",
+        )
+    )
+
+    assert "SUCCESS" in result.stdout, (
+        f"Iteration {iteration} failed.\nstdout: {result.stdout}\nstderr: {result.stderr}"
+    )
+    assert_no_determinism_errors(result)