Jyotisha/tests/test_api_heavy_compute_gate.py

"""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