"""Fail-fast bounded concurrency for heavy API compute endpoints.""" from __future__ import annotations import json import threading from io import BytesIO import pytest from scripts import api_heavy_compute_gate as gate from scripts.api_heavy_compute_gate import ( DEFAULT_CONCURRENCY, HEAVY_COMPUTE_PATHS, HeavyComputeBusy, acquire_heavy_compute_slot, is_heavy_compute_path, release_heavy_compute_slot, reset_heavy_compute_gate, ) from scripts.jyotish_api_server import ( DEFAULT_ALLOWED_HOSTS, DEFAULT_ALLOWED_ORIGINS, JyotishAPIHandler, ) class _FakeHeaders(dict): def get(self, key, default=None): return super().get(key, default) class _FakeServer: allowed_origins = DEFAULT_ALLOWED_ORIGINS allowed_hosts = DEFAULT_ALLOWED_HOSTS class _PostCaptureHandler(JyotishAPIHandler): def __init__(self, path: str, payload: dict) -> None: raw = json.dumps(payload).encode("utf-8") self.headers = _FakeHeaders( { "Content-Length": str(len(raw)), "Content-Type": "application/json", } ) self.server = _FakeServer() self.path = path self.rfile = BytesIO(raw) self.wfile = BytesIO() self.status_code = None self.response_headers = [] self.client_address = ("test-heavy-compute", 0) def send_response(self, code, message=None): # noqa: ANN001 self.status_code = code def send_header(self, key, value): # noqa: ANN001 self.response_headers.append((key, value)) def end_headers(self): return None def payload(self) -> dict: return json.loads(self.wfile.getvalue().decode("utf-8")) class _GetCaptureHandler(JyotishAPIHandler): def __init__(self, path: str) -> None: self.headers = _FakeHeaders() self.server = _FakeServer() self.path = path self.wfile = BytesIO() self.status_code = None self.response_headers = [] self.client_address = ("test-heavy-compute", 0) def send_response(self, code, message=None): # noqa: ANN001 self.status_code = code def send_header(self, key, value): # noqa: ANN001 self.response_headers.append((key, value)) def end_headers(self): return None def payload(self) -> dict: return json.loads(self.wfile.getvalue().decode("utf-8")) @pytest.fixture def limit_one_gate(monkeypatch): monkeypatch.setenv("JYOTISH_HEAVY_COMPUTE_CONCURRENCY", "1") monkeypatch.setenv("JYOTISH_HEAVY_COMPUTE_RETRY_AFTER_SECONDS", "2") monkeypatch.setenv("JYOTISH_API_RATE_LIMIT_PER_MINUTE", "0") reset_heavy_compute_gate() yield monkeypatch.delenv("JYOTISH_HEAVY_COMPUTE_CONCURRENCY", raising=False) monkeypatch.delenv("JYOTISH_HEAVY_COMPUTE_RETRY_AFTER_SECONDS", raising=False) reset_heavy_compute_gate() def test_health_and_light_paths_are_not_gated() -> None: assert not is_heavy_compute_path("/api/health") assert not is_heavy_compute_path("/api/cities") assert not is_heavy_compute_path("/api/location/resolve") assert not is_heavy_compute_path("/api/chart/jobs/abc") assert not is_heavy_compute_path("/api/high_rigor_workflow/jobs/abc") assert not is_heavy_compute_path("/api/rectification/v5/versions") assert is_heavy_compute_path("/api/high_rigor_workflow") assert is_heavy_compute_path("/api/rectification/sensitivity_scan") assert is_heavy_compute_path("/api/vedastro_gateway/jobs/job1/run") assert "/api/chart" not in HEAVY_COMPUTE_PATHS assert DEFAULT_CONCURRENCY == 2 def test_acquire_fail_fast_then_succeeds_after_release() -> None: reset_heavy_compute_gate(limit=1, retry_after_seconds=3) first = acquire_heavy_compute_slot("/api/high_rigor_workflow") assert first is not None with pytest.raises(HeavyComputeBusy) as caught: acquire_heavy_compute_slot("/api/rectification/sensitivity_scan") assert caught.value.retry_after_seconds == 3 assert caught.value.error_code == "ERR_COMPUTE_BUSY" light = acquire_heavy_compute_slot("/api/health") assert light is None release_heavy_compute_slot(first) second = acquire_heavy_compute_slot("/api/consultation_workflow") assert second is not None release_heavy_compute_slot(second) reset_heavy_compute_gate() def test_saturated_request_returns_429_with_retry_after(limit_one_gate, monkeypatch) -> None: started = threading.Event() release_first = threading.Event() first_status = {} def _slow_compute(self, body): # noqa: ANN001 started.set() assert release_first.wait(timeout=5) return {"ok": True, "endpoint": "high_rigor_workflow"} monkeypatch.setattr(JyotishAPIHandler, "_compute_high_rigor_workflow", _slow_compute) def _run_first() -> None: handler = _PostCaptureHandler("/api/high_rigor_workflow", {}) handler.do_POST() first_status["code"] = handler.status_code first_status["payload"] = handler.payload() worker = threading.Thread(target=_run_first) worker.start() assert started.wait(timeout=5) blocked = _PostCaptureHandler("/api/high_rigor_workflow", {}) blocked.do_POST() assert blocked.status_code == 429 assert ("Retry-After", "2") in blocked.response_headers payload = blocked.payload() assert payload["success"] is False assert payload["error_code"] == "ERR_COMPUTE_BUSY" health = _GetCaptureHandler("/api/health") health.do_GET() assert health.status_code == 200 assert health.payload()["status"] == "ok" light = _PostCaptureHandler("/api/location/resolve", {"city": "beijing"}) light.do_POST() assert light.status_code == 200 assert light.payload()["status"] == "local_city_match" release_first.set() worker.join(timeout=5) assert not worker.is_alive() assert first_status["code"] == 200 recovered = _PostCaptureHandler("/api/high_rigor_workflow", {}) recovered.do_POST() assert recovered.status_code == 200 assert recovered.payload()["ok"] is True def test_quality_gate_runs_heavy_compute_gate() -> None: from scripts.run_quality_gate import CORE_PYTEST_TARGETS assert "tests/test_api_heavy_compute_gate.py" in CORE_PYTEST_TARGETS assert gate.DEFAULT_CONCURRENCY == 2