tests/unit/test_api_auth_enforcement.py

"""Unit tests for ZeroGPU API bearer-token enforcement (P0 security).

Verifies that every GPU-burning REST endpoint refuses unauthenticated calls
when FLOW_API_KEY is set, and that the open-local mode (no FLOW_API_KEY)
still answers for read-only endpoints.

These tests statically parse app.py instead of importing it, because the
Space runtime deps (gradio, fastapi) are not installed in the local venv.
"""
import os
import re
import sys
import unittest

APP_PATH = os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "..", "app.py")

GPU_BURNING_ROUTES = [
    ("POST", "/v1/chat/completions"),
    ("POST", "/v1/audio/speech"),
    ("POST", "/v1/audio/jobs"),
    # ("/v1/images/generations" removed by policy: image generation out of Space scope.)
    ("POST", "/v1/embeddings"),
    ("POST", "/v1/warmup"),
]


def _read_source() -> str:
    with open(APP_PATH, "r", encoding="utf-8") as f:
        return f.read()


def _route_body(source: str, method: str, route: str) -> str:
    """Return the function body of a @fastapi_app.<method>("route") handler."""
    pattern = (
        r'@fastapi_app\.' + method.lower() + r'\("' + re.escape(route) + r'"\)'
        r'\s*async def \w+\(.*?\):(.*?)(?=@fastapi_app\.|\Z)'
    )
    match = re.search(pattern, source, re.DOTALL)
    assert match is not None, f"Route {method} {route} not found"
    return match.group(1)


class FlowApiKeyAuthEnforcement(unittest.TestCase):
    """With FLOW_API_KEY set, all GPU-burning endpoints must return 401 without a valid bearer."""

    def setUp(self):
        self._old_key = os.environ.get("FLOW_API_KEY")
        os.environ["FLOW_API_KEY"] = "test-secret-key-123"

    def tearDown(self):
        if self._old_key is not None:
            os.environ["FLOW_API_KEY"] = self._old_key
        else:
            os.environ.pop("FLOW_API_KEY", None)

    def test_authorize_function_exists_and_enforces_401(self):
        source = _read_source()
        self.assertIn("def _authorize_api_request", source)
        # It must raise HTTPException 401 on missing/mismatched bearer when key set
        body_match = re.search(
            r'def _authorize_api_request\(request: Request\) -> None:(.*?)\n\n',
            source,
            re.DOTALL,
        )
        self.assertIsNotNone(body_match)
        body = body_match.group(1)
        self.assertIn('os.environ.get("FLOW_API_KEY")', body)
        self.assertIn("401", body)
        self.assertIn("compare_digest", body)

    def test_all_gpu_burning_routes_call_auth_guard(self):
        """Static check: every GPU-burning route function calls _authorize_api_request."""
        source = _read_source()
        for method, route in GPU_BURNING_ROUTES:
            body = _route_body(source, method, route)
            self.assertIn(
                "_authorize_api_request(request)",
                body,
                f"Route {method} {route} does NOT call _authorize_api_request",
            )

    def test_audio_jobs_download_status_are_read_only(self):
        """Job status/download GETs are read-only; they must not leak the whole queue list."""
        source = _read_source()
        # status + download endpoints exist and are GET (no GPU burn)
        self.assertIn('@fastapi_app.get("/v1/audio/jobs/{job_id}")', source)
        self.assertIn('@fastapi_app.get("/v1/audio/jobs/{job_id}/download")', source)


class OpenLocalMode(unittest.TestCase):
    """Without FLOW_API_KEY the guard must be a documented no-op (open local mode)."""

    def setUp(self):
        self._old_key = os.environ.pop("FLOW_API_KEY", None)

    def tearDown(self):
        if self._old_key is not None:
            os.environ["FLOW_API_KEY"] = self._old_key

    def test_open_local_mode_is_explicit(self):
        source = _read_source()
        body_match = re.search(
            r'def _authorize_api_request\(request: Request\) -> None:(.*?)\n\n',
            source,
            re.DOTALL,
        )
        self.assertIsNotNone(body_match)
        self.assertIn("if not expected:", body_match.group(1))
        self.assertIn("return", body_match.group(1))


if __name__ == "__main__":
    unittest.main()

def test_local_callers_prioritize_flow_api_key(monkeypatch):
    """Verify that model_dispatcher and persona_engine prefer FLOW_API_KEY over HF_TOKEN."""
    monkeypatch.setenv("FLOW_API_KEY", "flow_secret_key_abc")
    monkeypatch.setenv("HF_TOKEN", "hf_token_fallback")

    # 1. model_dispatcher
    from engine.model_dispatcher import _space_auth_token
    assert _space_auth_token() == "flow_secret_key_abc"
    assert _space_auth_token(hf_token="explicit_hf") == "flow_secret_key_abc"

    # When FLOW_API_KEY is not set, falls back to hf_token
    monkeypatch.delenv("FLOW_API_KEY")
    assert _space_auth_token(hf_token=os.environ.get("HF_TOKEN")) == "hf_token_fallback"
    assert _space_auth_token(hf_token="explicit_hf") == "explicit_hf"

def test_gradio_ui_launch_is_optin_locked():
    """Verify the Gradio UI is password-locked only when GRADIO_UI_PASSWORD is set."""
    src = open(APP_PATH).read()
    # The launch block must gate auth on the GRADIO_UI_PASSWORD env var
    m = re.search(r'if __name__ == "__main__":(.*)', src, re.S)
    assert m, "launch block not found"
    block = m.group(1)
    assert 'os.environ.get("GRADIO_UI_PASSWORD"' in block, (
        "launch block does not gate UI auth on GRADIO_UI_PASSWORD"
    )
    assert 'auth=(_ui_username, _ui_password)' in block, (
        "gradio auth tuple not wired into launch"
    )
    # Default (no password) branch must remain open-local unchanged
    assert "demo.launch(prevent_thread_lock=True, ssr_mode=False)" in block