diff --git a/.gitignore b/.gitignore index afa1ad1..701c6dd 100644 --- a/.gitignore +++ b/.gitignore @@ -2,3 +2,4 @@ __pycache__/ .ruff_cache/ .venv/ +.pytest_cache/ diff --git a/Dockerfile b/Dockerfile index 700b8ab..e08d4ba 100644 --- a/Dockerfile +++ b/Dockerfile @@ -28,6 +28,9 @@ VOLUME ["/data"] EXPOSE 80 +HEALTHCHECK --start-period=5s --start-interval=1s CMD \ + test "$(curl -fsS "http://127.0.0.1:$PORT/health")" = "ok" + STOPSIGNAL SIGINT CMD ["/entrypoint.sh"] diff --git a/src/healthcheck.py b/src/healthcheck.py new file mode 100644 index 0000000..b38e59a --- /dev/null +++ b/src/healthcheck.py @@ -0,0 +1,18 @@ +import os +import urllib.request + + +def main() -> int: + port = os.environ.get("PORT", "80") + url = f"http://127.0.0.1:{port}/health" + + try: + with urllib.request.urlopen(url, timeout=2) as resp: + body = resp.read().decode("utf-8").strip() + return 0 if resp.status == 200 and body == "ok" else 1 + except Exception: + return 1 + + +if __name__ == "__main__": + exit(main()) diff --git a/src/providers/base.py b/src/providers/base.py index 478ab53..5f9cdc6 100644 --- a/src/providers/base.py +++ b/src/providers/base.py @@ -71,7 +71,7 @@ class Provider(ABC): """Usage percent when provider may switch active account/token.""" return None - def should_prepare_standby(self, usage_percent: int) -> bool: + def should_prepare_standby(self) -> bool: """Whether standby preparation should be triggered for current usage.""" return False diff --git a/src/providers/chatgpt/provider.py b/src/providers/chatgpt/provider.py index 603c247..b536611 100644 --- a/src/providers/chatgpt/provider.py +++ b/src/providers/chatgpt/provider.py @@ -122,14 +122,14 @@ class ChatGPTProvider(Provider): return True return await self._create_next_account_under_lock() - def should_prepare_standby(self, usage_percent: int) -> bool: - return usage_percent >= self.prepare_threshold and bool(load_next_tokens()) + def should_prepare_standby(self) -> bool: + return bool(load_next_tokens()) async def ensure_standby_account( self, usage_percent: int, ) -> None: - if self.should_prepare_standby(usage_percent): + if usage_percent >= self.prepare_threshold: await self.ensure_next_account() async def maybe_switch_active_account(self, usage_percent: int) -> bool: diff --git a/src/server.py b/src/server.py index be2e905..171b5c8 100644 --- a/src/server.py +++ b/src/server.py @@ -47,7 +47,7 @@ def should_trigger_standby_prepare(provider_name: str, usage_percent: int) -> bo provider = PROVIDERS.get(provider_name) if not provider: return False - return provider.should_prepare_standby(usage_percent) + return provider.should_prepare_standby() async def ensure_provider_token_ready(provider_name: str): @@ -86,7 +86,7 @@ async def ensure_standby_task(provider_name: str, usage_percent: int, reason: st if not provider: return - if not provider.should_prepare_standby(usage_percent): + if not provider.should_prepare_standby(): return try: @@ -188,6 +188,11 @@ async def token_handler(request: web.Request) -> web.Response: ) +async def health_handler(request: web.Request) -> web.Response: + del request + return web.Response(text="ok") + + async def on_startup(app: web.Application): del app for provider_name in PROVIDERS: @@ -208,6 +213,7 @@ def create_app() -> web.Application: app = web.Application(middlewares=[request_log_middleware]) app.on_startup.append(on_startup) app.on_cleanup.append(on_cleanup) + app.router.add_get("/health", health_handler) app.router.add_get("/{provider}/token", token_handler) app.router.add_get("/token", token_handler) return app diff --git a/tests/test_server_unit.py b/tests/test_server_unit.py index 68b3030..8ebefb8 100644 --- a/tests/test_server_unit.py +++ b/tests/test_server_unit.py @@ -34,8 +34,8 @@ class FakeProvider(Provider): def prepare_threshold(self) -> int: return self._prepare_threshold - def should_prepare_standby(self, usage_percent: int) -> bool: - return usage_percent >= self.prepare_threshold + def should_prepare_standby(self) -> bool: + return False @property def name(self) -> str: @@ -113,6 +113,12 @@ def test_token_handler_unknown_provider(monkeypatch): assert resp.status == 404 +def test_health_handler_ok(): + resp = asyncio.run(server.health_handler(object())) + assert resp.status == 200 + assert resp.text == "ok" + + def test_token_handler_success(monkeypatch): provider = FakeProvider() monkeypatch.setattr(server, "PROVIDERS", {"fake": provider})