test_backend_supervisor.py 2.0 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869
  1. from collections.abc import Sequence
  2. from zbt.commands.serve import run_backend_supervisor
  3. class FakeProcess:
  4. def __init__(self, return_code: int | None) -> None:
  5. self._return_code = return_code
  6. self.terminated = False
  7. self.killed = False
  8. def poll(self) -> int | None:
  9. return self._return_code
  10. def terminate(self) -> None:
  11. self.terminated = True
  12. self._return_code = 0
  13. def wait(self, timeout: float | None = None) -> int:
  14. del timeout
  15. return 0 if self._return_code is None else self._return_code
  16. def kill(self) -> None:
  17. self.killed = True
  18. self._return_code = -9
  19. def test_unified_backend_starts_api_and_worker_and_stops_peer_on_exit() -> None:
  20. api = FakeProcess(return_code=0)
  21. worker = FakeProcess(return_code=None)
  22. launched_roles: list[str] = []
  23. def launch(role: str, command: Sequence[str]) -> FakeProcess:
  24. assert command
  25. launched_roles.append(role)
  26. return api if role == "api" else worker
  27. result = run_backend_supervisor(
  28. launch=launch,
  29. api_command=("python", "-m", "uvicorn", "zbt.main:app"),
  30. worker_command=("python", "-m", "zbt.commands.run_refund_worker"),
  31. output=lambda _: None,
  32. poll_interval_seconds=0,
  33. )
  34. assert result == 0
  35. assert launched_roles == ["api", "refund-worker"]
  36. assert worker.terminated is True
  37. assert worker.killed is False
  38. def test_unified_backend_fails_when_refund_worker_exits_unexpectedly() -> None:
  39. api = FakeProcess(return_code=None)
  40. worker = FakeProcess(return_code=0)
  41. def launch(role: str, command: Sequence[str]) -> FakeProcess:
  42. del command
  43. return api if role == "api" else worker
  44. result = run_backend_supervisor(
  45. launch=launch,
  46. api_command=("api",),
  47. worker_command=("worker",),
  48. output=lambda _: None,
  49. poll_interval_seconds=0,
  50. )
  51. assert result == 1
  52. assert api.terminated is True