writer-work-flow/apps/api/tests/test_kimi_oauth_service.py

"""K1.3 单测：Kimi Code OAuth device-flow 服务（fake httpx，绝不联网）。

覆盖：start_device_authorization、poll_token（成功/pending/slow_down/失败）、refresh、
encrypt/decrypt round-trip、needs_refresh 启发式。
"""

from __future__ import annotations

from datetime import UTC, datetime, timedelta
from typing import Any

import pytest
from cryptography.fernet import Fernet
from ww_api.services.kimi_oauth import (
    DEVICE_AUTHORIZATION_URL,
    TOKEN_URL,
    AuthorizationPending,
    SlowDown,
    TokenSet,
    decrypt_oauth_bundle,
    encrypt_oauth_bundle,
    needs_refresh,
    poll_token,
    refresh,
    start_device_authorization,
)
from ww_shared import AppError, ErrorCode


class _FakeResponse:
    def __init__(self, status_code: int, body: dict[str, Any]) -> None:
        self.status_code = status_code
        self._body = body

    def json(self) -> Any:
        return self._body


class _FakeHttp:
    """记录每次 post 的 url + data，按 url 返回预设响应。"""

    def __init__(self, responses: list[_FakeResponse]) -> None:
        self._responses = responses
        self.calls: list[tuple[str, dict[str, str]]] = []

    async def post(self, url: str, *, data: dict[str, str]) -> _FakeResponse:
        self.calls.append((url, data))
        return self._responses[len(self.calls) - 1]


@pytest.mark.asyncio
async def test_start_device_authorization_parses_response_and_sends_scope() -> None:
    http = _FakeHttp(
        [
            _FakeResponse(
                200,
                {
                    "device_code": "dev-123",
                    "user_code": "ABCD-1234",
                    "verification_uri": "https://kimi.com/device",
                    "verification_uri_complete": "https://kimi.com/device?code=ABCD-1234",
                    "expires_in": 600,
                    "interval": 5,
                },
            )
        ]
    )

    device = await start_device_authorization(http)

    assert device.device_code == "dev-123"
    assert device.user_code == "ABCD-1234"
    assert device.interval == 5
    assert device.expires_in == 600
    # 请求打到 device authorization 端点，且**带 scope=kimi-code**（coding entitlement）。
    url, data = http.calls[0]
    assert url == DEVICE_AUTHORIZATION_URL
    assert "client_id" in data
    assert data["scope"] == "kimi-code"


@pytest.mark.asyncio
async def test_poll_token_success_returns_token_set() -> None:
    http = _FakeHttp(
        [
            _FakeResponse(
                200,
                {
                    "access_token": "acc-xyz",
                    "refresh_token": "ref-xyz",
                    "expires_in": 900,
                },
            )
        ]
    )

    token = await poll_token(http, "dev-123")

    assert token.access_token == "acc-xyz"
    assert token.refresh_token == "ref-xyz"
    assert token.expires_at > datetime.now(UTC)
    # 打到 token 端点 + device_code grant。
    url, data = http.calls[0]
    assert url == TOKEN_URL
    assert data["device_code"] == "dev-123"
    assert data["grant_type"].endswith("device_code")


@pytest.mark.asyncio
async def test_poll_token_authorization_pending_raises_pending() -> None:
    http = _FakeHttp([_FakeResponse(400, {"error": "authorization_pending"})])
    with pytest.raises(AuthorizationPending):
        await poll_token(http, "dev-123")


@pytest.mark.asyncio
async def test_poll_token_slow_down_raises_slow_down() -> None:
    http = _FakeHttp([_FakeResponse(400, {"error": "slow_down"})])
    with pytest.raises(SlowDown):
        await poll_token(http, "dev-123")


@pytest.mark.asyncio
async def test_poll_token_expired_raises_app_error() -> None:
    http = _FakeHttp([_FakeResponse(400, {"error": "expired_token"})])
    with pytest.raises(AppError) as exc:
        await poll_token(http, "dev-123")
    assert exc.value.code == ErrorCode.LLM_UNAVAILABLE


@pytest.mark.asyncio
async def test_refresh_returns_new_token_set() -> None:
    http = _FakeHttp(
        [
            _FakeResponse(
                200, {"access_token": "new-acc", "refresh_token": "new-ref", "expires_in": 900}
            )
        ]
    )

    token = await refresh(http, "ref-old")

    assert token.access_token == "new-acc"
    assert token.refresh_token == "new-ref"
    url, data = http.calls[0]
    assert url == TOKEN_URL
    assert data["grant_type"] == "refresh_token"
    assert data["refresh_token"] == "ref-old"


@pytest.mark.asyncio
async def test_refresh_failure_raises_app_error() -> None:
    http = _FakeHttp([_FakeResponse(400, {"error": "invalid_grant"})])
    with pytest.raises(AppError) as exc:
        await refresh(http, "ref-bad")
    assert exc.value.code == ErrorCode.LLM_UNAVAILABLE


def test_encrypt_decrypt_oauth_bundle_round_trip() -> None:
    key = Fernet.generate_key().decode()
    expires = datetime.now(UTC) + timedelta(seconds=900)
    token = TokenSet(access_token="acc", refresh_token="ref", expires_at=expires)

    blob = encrypt_oauth_bundle(token, key=key)
    restored = decrypt_oauth_bundle(blob, key=key)

    assert restored.access_token == "acc"
    assert restored.refresh_token == "ref"
    assert restored.expires_at == expires
    # 密文不含明文 token（加密真生效）。
    assert b"acc" not in blob
    assert b"ref" not in blob


def test_needs_refresh_true_when_near_expiry() -> None:
    now = datetime.now(UTC)
    near = TokenSet(access_token="a", refresh_token="r", expires_at=now + timedelta(seconds=60))
    assert needs_refresh(near, now=now) is True


def test_needs_refresh_false_when_fresh() -> None:
    now = datetime.now(UTC)
    fresh = TokenSet(access_token="a", refresh_token="r", expires_at=now + timedelta(seconds=900))
    assert needs_refresh(fresh, now=now) is False