CODE_ASSISTANT/5_django_backend/tests/test_gateway.py

"""/api/ito — OpenCode → FabriX 통과 중계."""

import json

import httpx
import pytest
from django.test import AsyncClient, override_settings

from apps.gateway import views
from apps.gateway.fabrix import FabrixConfig, auth_variants, parse_models, token_value

FULL_ENV = {
    "AAF_FABRIX_BASE_URL": "https://fabrix.test/openapi/llm/",
    "AAF_FABRIX_MODEL_ID": "339",
    "AAF_FABRIX_MODELS": "339:GaussO Flash,581:GaussO Think",
    "AAF_FABRIX_CLIENT_KEY": "ck",
    "AAF_FABRIX_OPENAPI_TOKEN": "Bearer tok",
    "AAF_FABRIX_USER_EMAIL": "",
}


# ── 순수 규칙 ────────────────────────────────────────────────────
def test_parse_models_and_token():
    assert parse_models("339:GaussO Flash, 581 , ,x:") == {"339": "GaussO Flash", "581": "581", "x": "x"}
    assert token_value("Bearer abc") == "abc"
    assert token_value("abc", "bearer") == "Bearer abc"


def test_prepare_picks_model_by_header_not_body():
    cfg = FabrixConfig.from_env({**FULL_ENV, "AAF_FABRIX_MAX_TOKENS": "4096", "AAF_RELAY_STREAM_USAGE": "1"})
    assert cfg.url == "https://fabrix.test/openapi/llm/chat/completions"
    body, headers = cfg.prepare({"model": "581", "stream": True, "messages": []})
    assert body["model"] == "/mnt/models" and headers["x-llm-model-id"] == "581"
    assert body["stream_options"] == {"include_usage": True} and body["max_completion_tokens"] == 4096
    assert headers["x-openapi-token"] == "tok" and headers["x-generative-ai-client"] == "ck"
    # 목록에 없는 이름(더미)은 기본 모델
    _, headers = cfg.prepare({"model": "gpt-whatever"})
    assert headers["x-llm-model-id"] == "339"


def test_vision_model_only_when_image_attached():
    cfg = FabrixConfig.from_env({**FULL_ENV, "AAF_FABRIX_VISION_MODEL_ID": "605"})
    text = {"model": "339", "messages": [{"role": "user", "content": "hi"}]}
    img = {"model": "339", "messages": [{"role": "user", "content": [
        {"type": "text", "text": "what is this"},
        {"type": "image_url", "image_url": {"url": "data:image/png;base64,AAAA"}},
    ]}]}
    assert cfg.prepare(text)[1]["x-llm-model-id"] == "339"
    assert cfg.prepare(img)[1]["x-llm-model-id"] == "605"
    # 설정 없으면 분기 안 함
    assert FabrixConfig.from_env(FULL_ENV).prepare(img)[1]["x-llm-model-id"] == "339"


def test_missing_config_names():
    assert FabrixConfig.from_env({}).missing() == [
        "AAF_FABRIX_BASE_URL", "AAF_FABRIX_MODEL_ID", "AAF_FABRIX_CLIENT_KEY", "AAF_FABRIX_OPENAPI_TOKEN",
    ]


def test_auth_variants_order():
    headers = FabrixConfig.from_env(FULL_ENV).headers()
    keys = [k for k, _ in auth_variants(headers)]
    assert keys[0] == ("raw", "x-generative-ai-client") and len(keys) == 4
    # 마지막에 통과한 조합이 맨 앞으로
    keys2 = [k for k, _ in auth_variants(headers, ("bearer", "x-fabrix-client"))]
    assert keys2[0] == ("bearer", "x-fabrix-client") and set(keys2) == set(keys)
    hdrs = dict(auth_variants(headers))[("bearer", "x-fabrix-client")]
    assert hdrs["x-openapi-token"] == "Bearer tok" and hdrs["x-fabrix-client"] == "ck"
    assert "x-generative-ai-client" not in hdrs


# ── 뷰 (상류는 MockTransport) ────────────────────────────────────
@pytest.fixture
def upstream(monkeypatch):
    calls: list[httpx.Request] = []

    def make(handler):
        def _h(req: httpx.Request) -> httpx.Response:
            calls.append(req)
            return handler(req)

        monkeypatch.setattr(views, "TRANSPORT", httpx.MockTransport(_h))
        views._last_ok["variant"] = None
        return calls

    return make


def _sse(data: bytes, status: int = 200) -> httpx.Response:
    """상류 스트림 응답. content= 로 만들면 httpx 가 미리 읽어버려(StreamConsumed) stream= 으로."""
    return httpx.Response(status, stream=httpx.ByteStream(data))


async def _post(body: dict, auth: str | None = None):
    headers = {"Authorization": auth} if auth else {}
    return await AsyncClient().post(
        "/api/ito/chat/completions", data=json.dumps(body), content_type="application/json", headers=headers
    )


@override_settings(FABRIX_ENV=FULL_ENV)
async def test_stream_passthrough_bytes(upstream):
    calls = upstream(lambda r: _sse(b'data: {"choices":[]}\n\ndata: [DONE]\n\n'))
    resp = await _post({"model": "581", "stream": True, "messages": [{"role": "user", "content": "hi"}]})
    assert resp.status_code == 200 and resp["Content-Type"].startswith("text/event-stream")
    assert b"".join([c async for c in resp.streaming_content]) == b'data: {"choices":[]}\n\ndata: [DONE]\n\n'
    req = calls[0]
    assert str(req.url) == "https://fabrix.test/openapi/llm/chat/completions"
    assert req.headers["x-llm-model-id"] == "581" and json.loads(req.content)["model"] == "/mnt/models"


@override_settings(FABRIX_ENV=FULL_ENV)
async def test_401_retries_other_auth_shape_and_remembers(upstream):
    def handler(req: httpx.Request) -> httpx.Response:
        if req.headers.get("x-openapi-token") != "Bearer tok":
            return httpx.Response(401, content=b"OIDC-TOKEN-0")
        return httpx.Response(200, json={"choices": [{"message": {"content": "ok"}}]})

    calls = upstream(handler)
    resp = await _post({"messages": []})
    assert resp.status_code == 200 and json.loads(resp.content)["choices"][0]["message"]["content"] == "ok"
    assert calls[0].headers["x-openapi-token"] == "tok" and calls[-1].headers["x-openapi-token"] == "Bearer tok"
    # 다음 요청은 통과한 형식부터
    n = len(calls)
    await _post({"messages": []})
    assert len(calls) == n + 1 and calls[-1].headers["x-openapi-token"] == "Bearer tok"


@override_settings(FABRIX_ENV=FULL_ENV)
async def test_flaky_401_on_right_auth_is_retried_once(upstream, monkeypatch):
    """FabriX 가 맞는 조합에도 가끔 401 — 4개 다 돈 뒤 맨 앞 조합 한 번 더."""
    from apps.gateway import views

    monkeypatch.setattr(views, "RETRY_PAUSE_S", 0)
    hits = {"n": 0}

    def handler(req: httpx.Request) -> httpx.Response:
        hits["n"] += 1
        if hits["n"] <= 4:
            return httpx.Response(401, content=b"flaky")
        return httpx.Response(200, json={"choices": []})

    calls = upstream(handler)
    resp = await _post({"messages": []})
    assert resp.status_code == 200 and len(calls) == 5
    assert calls[0].headers.get("x-openapi-token") == calls[4].headers.get("x-openapi-token")


@override_settings(FABRIX_ENV=FULL_ENV)
async def test_stream_upstream_error_becomes_sse_error(upstream):
    upstream(lambda r: _sse(b"boom", 500))
    resp = await _post({"stream": True, "messages": []})
    out = b"".join([c async for c in resp.streaming_content]).decode()
    assert resp.status_code == 200 and '"type": "upstream_error"' in out and "boom" in out and "[DONE]" in out


@override_settings(FABRIX_ENV={})
async def test_missing_config_is_503_but_health_ok():
    resp = await _post({"messages": []})
    assert resp.status_code == 503 and "AAF_FABRIX_BASE_URL" in json.loads(resp.content)["error"]["message"]
    assert (await AsyncClient().get("/api/ito/healthcheck")).status_code == 200


@override_settings(FABRIX_ENV={**FULL_ENV, "AAF_GATEWAY_KEY": "secret"})
async def test_gateway_key_and_models(upstream):
    upstream(lambda r: httpx.Response(200, json={}))
    assert (await _post({"messages": []})).status_code == 401
    assert (await _post({"messages": []}, auth="Bearer secret")).status_code == 200
    resp = await AsyncClient().get("/api/ito/models", headers={"Authorization": "Bearer secret"})
    assert [m["id"] for m in json.loads(resp.content)["data"]] == ["339", "581"]