Add unit tests for task run recorder redelivery behavior

Co-Authored-By: Nate Nowack <[email protected]>

Author: devin-ai-integration[bot]

tests/server/services/test_task_run_recorder.py

@@ -1173,3 +1173,65 @@ async def test_subsequent_updates_move_update_timestamp(session: AsyncSession):
     assert second_update_timestamp is not None
 
     assert second_update_timestamp > first_update_timestamp
+
+
+async def test_event_retried_on_persist_failure(
+    pending_event: ReceivedEvent,
+    monkeypatch: pytest.MonkeyPatch,
+    caplog: pytest.LogCaptureFixture,
+):
+    """Test that events are retried when handle_task_run_events fails."""
+    call_count = 0
+
+    async def mock_handle_task_run_events(events, depth=0):
+        nonlocal call_count
+        call_count += 1
+        if call_count == 1:
+            raise Exception("Simulated DB failure")
+
+    monkeypatch.setattr(
+        "prefect.server.services.task_run_recorder.handle_task_run_events",
+        mock_handle_task_run_events,
+    )
+
+    async with task_run_recorder.consumer(
+        write_batch_size=1, flush_every=1, max_persist_retries=2
+    ) as handler:
+        with caplog.at_level("ERROR"):
+            await handler(message(pending_event))
+            await asyncio.sleep(1.5)
+
+    assert call_count == 2
+    assert "1 to retry" in caplog.text
+    assert "0 dropped" in caplog.text
+
+
+async def test_event_dropped_after_max_retries_exceeded(
+    pending_event: ReceivedEvent,
+    monkeypatch: pytest.MonkeyPatch,
+    caplog: pytest.LogCaptureFixture,
+):
+    """Test that events are dropped after exceeding max_persist_retries."""
+    call_count = 0
+
+    async def mock_handle_task_run_events(events, depth=0):
+        nonlocal call_count
+        call_count += 1
+        raise Exception("Simulated persistent DB failure")
+
+    monkeypatch.setattr(
+        "prefect.server.services.task_run_recorder.handle_task_run_events",
+        mock_handle_task_run_events,
+    )
+
+    async with task_run_recorder.consumer(
+        write_batch_size=1, flush_every=1, max_persist_retries=1
+    ) as handler:
+        with caplog.at_level("ERROR"):
+            await handler(message(pending_event))
+            await asyncio.sleep(1.5)
+
+    assert call_count == 2
+    assert "Dropping event" in caplog.text
+    assert "after 2 failed attempts" in caplog.text
+    assert "1 dropped" in caplog.text