lingma-openai-gateway/tests/test_auth_concurrency.py

from __future__ import annotations

import asyncio
import sys
import types
import unittest
from unittest.mock import patch

from fastapi import HTTPException
from fastapi.testclient import TestClient
from starlette.requests import Request

from app.auth import AnthropicAuthError, require_anthropic_key, require_bearer, require_metrics_access
from app.concurrency import BackpressureRejected, InFlightGuard

_playwright = types.ModuleType("playwright")
_playwright_async = types.ModuleType("playwright.async_api")


class _StubPlaywrightTimeoutError(Exception):
    pass


async def _stub_async_playwright():
    raise RuntimeError("playwright is stubbed in unit tests")


_playwright_async.TimeoutError = _StubPlaywrightTimeoutError
_playwright_async.async_playwright = _stub_async_playwright
sys.modules.setdefault("playwright", _playwright)
sys.modules.setdefault("playwright.async_api", _playwright_async)

import app.main as main


def _req(headers: dict[str, str] | None = None) -> Request:
    pairs = []
    for k, v in (headers or {}).items():
        pairs.append((k.lower().encode("latin-1"), v.encode("latin-1")))
    scope = {
        "type": "http",
        "http_version": "1.1",
        "method": "GET",
        "scheme": "http",
        "path": "/x",
        "raw_path": b"/x",
        "query_string": b"",
        "headers": pairs,
        "client": ("test", 1),
        "server": ("test", 80),
        "root_path": "",
    }
    return Request(scope)


class AuthAndConcurrencyTests(unittest.IsolatedAsyncioTestCase):
    def test_require_bearer_accepts_valid_token(self) -> None:
        request = _req({"authorization": "Bearer good"})
        require_bearer(request, ["good"])

    def test_require_bearer_rejects_invalid_token(self) -> None:
        request = _req({"authorization": "Bearer bad"})
        with self.assertRaises(HTTPException) as ctx:
            require_bearer(request, ["good"])
        self.assertEqual(ctx.exception.status_code, 401)
        self.assertEqual(ctx.exception.detail["error"]["code"], "invalid_api_key")

    def test_require_anthropic_key_accepts_x_api_key_or_bearer(self) -> None:
        request_x = _req({"x-api-key": "k1"})
        require_anthropic_key(request_x, ["k1"])

        request_b = _req({"authorization": "Bearer k2"})
        require_anthropic_key(request_b, ["k2"])

    def test_require_anthropic_key_raises_on_missing(self) -> None:
        request = _req()
        with self.assertRaises(AnthropicAuthError) as ctx:
            require_anthropic_key(request, ["k"])
        self.assertEqual(ctx.exception.status_code, 401)
        self.assertEqual(ctx.exception.error_type, "authentication_error")

    def test_require_metrics_access_503_when_no_tokens_configured(self) -> None:
        request = _req({"authorization": "Bearer any"})
        with self.assertRaises(HTTPException) as ctx:
            require_metrics_access(request, api_keys=[], metrics_token="", public=False)
        self.assertEqual(ctx.exception.status_code, 503)
        self.assertEqual(ctx.exception.detail["error"]["code"], "metrics_disabled")

    async def test_inflight_guard_unlimited_and_release_idempotent(self) -> None:
        guard = InFlightGuard(max_in_flight=0, queue_timeout_sec=0.01)
        ticket = await guard.try_acquire()
        self.assertEqual(guard.in_flight, 1)
        ticket.release()
        ticket.release()
        self.assertEqual(guard.in_flight, 0)
        self.assertEqual(guard.accepted_total, 1)

    async def test_inflight_guard_rejects_when_queue_timeout(self) -> None:
        guard = InFlightGuard(max_in_flight=1, queue_timeout_sec=0.01)
        first = await guard.try_acquire()
        with self.assertRaises(BackpressureRejected):
            await guard.try_acquire()
        self.assertEqual(guard.rejected_total, 1)
        first.release()
        self.assertEqual(guard.in_flight, 0)


class DebugRequestRecordingTests(unittest.TestCase):
    def setUp(self) -> None:
        main._DEBUG_REQUEST_LOG.clear()

    def test_redacts_sensitive_fields_and_data_urls(self) -> None:
        body = {
            "authorization": "Bearer abc",
            "x-api-key": "secret",
            "session_bundle": "very-secret",
            "images": ["data:image/png;base64,ABC"],
            "tool": {"args": "x" * 3000},
        }
        redacted = main._redact_debug_value((), body)

        self.assertEqual(redacted["authorization"], "***")
        self.assertEqual(redacted["x-api-key"], "***")
        self.assertEqual(redacted["session_bundle"], "***")
        self.assertEqual(redacted["images"][0], "[redacted-data-url]")
        self.assertIn("[truncated]", redacted["tool"]["args"])

    def test_internal_debug_requests_requires_admin_and_returns_items(self) -> None:
        with patch.object(main.settings, "api_keys", ["k1"]), patch.object(main.settings, "admin_token", "admin-1"):
            client = TestClient(main.app)
            req_payload = {
                "model": "org_auto",
                "messages": [{"role": "user", "content": "hello"}],
            }
            main._record_debug_request("openai", "/v1/chat/completions", req_payload, _req({"x-request-id": "req-1"}))

            denied = client.get("/internal/debug/requests")
            self.assertEqual(denied.status_code, 401)

            ok = client.get(
                "/internal/debug/requests?limit=1",
                headers={"Authorization": "Bearer admin-1"},
            )
            self.assertEqual(ok.status_code, 200)
            data = ok.json()
            self.assertTrue(data["ok"])
            self.assertEqual(data["count"], 1)
            self.assertEqual(data["items"][0]["protocol"], "openai")


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