Class _FakeWorkResource

tests/lib/environments/test_worker.py:90–134

class _FakeWorkResource:
    def __init__(self, *, heartbeat_state: str = "stopping", heartbeat_error: Exception | None = None) -> None:
        self._heartbeat_state = heartbeat_state
        self._heartbeat_error = heartbeat_error
        # When set, every heartbeat waits for this event before answering, so a
        # test decides *when* in the run the heartbeat's outcome lands.
        self.heartbeat_gate: asyncio.Event | None = None
        self.heartbeat_calls: list[dict[str, Any]] = []
        self.stop_calls: list[dict[str, Any]] = []

    async def heartbeat(
        self,
        work_id: str,
        *,
        environment_id: str,
        expected_last_heartbeat: str,  # noqa: ARG002
        extra_headers: Any = None,
    ) -> Any:
        self.heartbeat_calls.append(
            {"work_id": work_id, "environment_id": environment_id, "extra_headers": extra_headers}
        )
        if self.heartbeat_gate is not None:
            await self.heartbeat_gate.wait()
        if self._heartbeat_error is not None:
            raise self._heartbeat_error
        return SimpleNamespace(last_heartbeat="hb-1", ttl_seconds=60, state=self._heartbeat_state, lease_extended=True)

    async def stop(
        self,
        work_id: str,
        *,
        environment_id: str,
        force: bool = False,
        extra_headers: Any = None,
        betas: Any = None,
    ) -> None:
        self.stop_calls.append(
            {
                "work_id": work_id,
                "environment_id": environment_id,
                "force": force,
                "extra_headers": extra_headers,
                "betas": betas,
            }
        )


class _FakeSessions: