Function _mock_runner

tests/algorithm/test_baseline.py:72–108

Simulate a runner consuming rollouts, adding spans, and marking them complete.

(
    *,
    store: InMemoryLightningStore,
    expected: int,
    artifacts: List[_RolloutArtifacts],
)

async def _mock_runner(
    *,
    store: InMemoryLightningStore,
    expected: int,
    artifacts: List[_RolloutArtifacts],
) -> None:
    """Simulate a runner consuming rollouts, adding spans, and marking them complete."""
    processed = 0
    while processed < expected:
        attempted = await store.dequeue_rollout()
        if attempted is None:
            await asyncio.sleep(0.001)
            continue

        attempt = attempted.attempt
        rollout_id = attempted.rollout_id
        await store.update_attempt(
            rollout_id,
            attempt.attempt_id,
            status="running",
            worker_id="runner-1",
        )

        span = _build_span(rollout_id, attempt.attempt_id, sequence_id=1, index=processed + 1)
        await store.add_span(span)
        await store.update_attempt(rollout_id, attempt.attempt_id, status="succeeded")
        await store.update_rollout(rollout_id, status="succeeded")

        artifacts.append(
            _RolloutArtifacts(
                rollout_id=rollout_id,
                attempt_id=attempt.attempt_id,
                attempt_sequence=attempt.sequence_id,
                span=span,
            )
        )
        processed += 1


@pytest.mark.asyncio