Files
writer-work-flow/tests/test_k1_kimi_oauth_e2e.py

517 lines
22 KiB
Python
Raw Permalink Blame History

This file contains ambiguous Unicode characters

This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.

"""K1 端到端Kimi Code OAuth device-flowdevice/token/refresh **全 mock**)→ 真 pg。
证明 K1 闭环DoDdevice-flow 连接 → 加密 token 持久化(DB 真源,无明文) → 网关建出
KimiCodeAdapter(coding base + 伪造头 + bearer) → 建网关时按需刷新更新存量 token → 断开清除)。
**绝不联网**device authorization / token / refresh 全经注入的 fake httpx 客户端,
auth.kimi.com / api.kimi.com 一次都不碰。
用例 1 · 连接 device flow → token 加密持久化DB 真源)
`POST .../oauth/start` → 202 `{job_id, user_code, verification_uri, ...}`。fake
`device_authorization` 返 device_codefake `token` 端点先 `authorization_pending`
两次、再成功 `{access_token, refresh_token, expires_in}`。后台 `run_job` 轮询循环跑在
**真 e2e_sm session** 上ASGITransport await background taskoverride
`get_session_factory`→e2e_sm、monkeypatch 模块级 `_default_http_client`→fake、
`asyncio.sleep`→no-op。`GET /jobs/{id}` → doneresult `{connected:true,
provider:"kimi-code"}` 且**不含任何 token 串**。DB 真源:`provider_credentials` 行
`auth_type="oauth"`、`oauth_enc` 非空、`api_key_enc` 为 None解密 `oauth_enc` 回环
access/refresh token证明加密持久化、非明文。`GET .../oauth/status` →
`{connected:true, expires_at}` 无 token。
用例 2 · 网关建出 Kimi Code 适配器(正确 token + 伪造头 + coding base
存 oauth 凭据 + `tier_routing` 行把某档位指向 `kimi-code:kimi-for-coding` →
`build_gateway_for_tier(...,"writer")` → 断言适配器是 `KimiCodeAdapter`,其
`AsyncOpenAI` 客户端 `base_url=https://api.kimi.com/coding/v1`、带伪造 `User-Agent`
+ 全 `X-Msh-*` 默认头、`api_key`(=bearer) = 存量 access token。
用例 3 · 建网关时刷新refresh-on-build
存一个 `expires_at` 临近过期的 oauth 包fake token 端点对 `grant_type=refresh_token`
返**新** token 组。`_build_provider_adapter(store,"kimi-code")`(经真 e2e_sm session
→ 断言 refresh 被调、适配器用**新** access token、`oauth_enc`(DB 真源)已更新为新包。
用例 4 · 断开disconnect
`POST .../oauth/disconnect` → 凭据行被清除DB 真源:之后无 oauth 凭据)。
零成本 & 确定性fake httpxdevice/token/refresh+ 真 `SqlCredentialStore`/真 pg。无 pg → skip。
坑(见 memory/gotchas
- 后台 work 自建 http 走**模块级** `routers.kimi_oauth._default_http_client()`(非 dep——
须 monkeypatch 它指到同一 scripted fake端点 device-auth call[0] + work poll call[1..]
顺序共享);端点 device-auth 走 `http` dep现为 `_http_client` 生成器依赖CR-H2
再 override `_http_client` dep。
- 真 pg 持久化:**不** monkeypatch `SqlCredentialStore`(让真 store 落库override
`get_session_factory`→`e2e_sm`(同测试 engine/loop使后台 work 用真 session 写库。
- monkeypatch `routers.kimi_oauth.asyncio.sleep`→no-op不真等 interval 秒)。
- `_resolve_kimi_code_token` 刷新走 `project_deps.httpx.AsyncClient` + `project_deps.kimi_refresh`——
用字符串 target monkeypatchmypy strict 不静态校验字符串属性)。
- token 绝不进 job 结果/状态响应/start 响应——显式负向断言。
- get_sessionmaker engine 绑首个 loop → 每 DB 测清缓存重建、结束 dispose同 M3/M4/M5
"""
from __future__ import annotations
import os
from collections.abc import AsyncIterator
from datetime import UTC, datetime, timedelta
from typing import Any
import httpx
import pytest
from asgi_lifespan import LifespanManager
from cryptography.fernet import Fernet
from sqlalchemy import delete, select
from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker
from ww_db import get_sessionmaker
from ww_db.models import Job, ProviderCredential, TierRouting
from ww_llm_gateway.adapters.kimi_code import (
KIMI_CODE_BASE_URL,
KIMI_CODE_PROVIDER,
KIMI_CODE_USER_AGENT,
KimiCodeAdapter,
)
_ACCESS_TOKEN = "kimi-acc-e2e-not-real-12345"
_REFRESH_TOKEN = "kimi-ref-e2e-not-real-67890"
_NEW_ACCESS_TOKEN = "kimi-acc-e2e-refreshed-abcde"
_NEW_REFRESH_TOKEN = "kimi-ref-e2e-refreshed-fghij"
_OAUTH_BASE = "/settings/providers/kimi-code/oauth"
_START_URL = f"{_OAUTH_BASE}/start"
_DISCONNECT_URL = f"{_OAUTH_BASE}/disconnect"
_STATUS_URL = f"{_OAUTH_BASE}/status"
_REQUIRED_HEADER_KEYS = (
"User-Agent",
"X-Msh-Platform",
"X-Msh-Version",
"X-Msh-Device-Name",
"X-Msh-Device-Model",
"X-Msh-Os-Version",
"X-Msh-Device-Id",
)
class _FakeResponse:
"""`HttpResponse` Protocol 最小实现status_code + json())。"""
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 _ScriptedHttp:
"""按调用序号返回预设响应device auth → poll attempts记录每次调用。aclose 无操作。
端点的 device authorization 调用是 call[0],后台 work 的 token 轮询是 call[1..]。
同一实例供两者顺序共享(须 monkeypatch 模块级 `_default_http_client` 指到它)。
"""
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:
idx = len(self.calls)
self.calls.append((url, data))
return self._responses[idx]
async def aclose(self) -> None:
return None
def _device_auth_response() -> _FakeResponse:
return _FakeResponse(
200,
{
"device_code": "dev-e2e-1",
"user_code": "ABCD-1234",
"verification_uri": "https://auth.kimi.com/device",
"verification_uri_complete": "https://auth.kimi.com/device?code=ABCD-1234",
"expires_in": 600,
"interval": 5,
},
)
def _pending_response() -> _FakeResponse:
return _FakeResponse(400, {"error": "authorization_pending"})
def _token_success_response() -> _FakeResponse:
return _FakeResponse(
200,
{
"access_token": _ACCESS_TOKEN,
"refresh_token": _REFRESH_TOKEN,
"expires_in": 900,
},
)
def _enc_key() -> str:
"""设置 `CREDENTIAL_ENC_KEY` 为有效 Fernet 密钥 + 清 settings 缓存,返回该密钥。"""
key = Fernet.generate_key().decode()
os.environ["CREDENTIAL_ENC_KEY"] = key
from ww_config import get_settings
get_settings.cache_clear()
return key
@pytest.fixture
async def e2e_sm() -> AsyncIterator[async_sessionmaker[AsyncSession]]:
"""真实 DB session 工厂;无 pg 时跳过(每测试清缓存重建 engine、结束 dispose"""
get_sessionmaker.cache_clear()
maker = get_sessionmaker()
try:
async with maker() as probe:
await probe.execute(select(1))
except Exception:
pytest.skip("postgres not reachable")
yield maker
await maker.kw["bind"].dispose()
get_sessionmaker.cache_clear()
async def _cleanup_credentials(e2e_sm: async_sessionmaker[AsyncSession]) -> None:
"""清理本测试族写入的 kimi-code 凭据 + kimi_oauth job 行(单用户 stub owner"""
from ww_api.services.credentials import STUB_OWNER_ID
async with e2e_sm() as cleanup:
await cleanup.execute(
delete(ProviderCredential).where(
ProviderCredential.owner_id == STUB_OWNER_ID,
ProviderCredential.provider == KIMI_CODE_PROVIDER,
)
)
await cleanup.execute(delete(Job).where(Job.kind == "kimi_oauth"))
await cleanup.execute(delete(TierRouting).where(TierRouting.project_id.is_(None)))
await cleanup.commit()
async def _seed_oauth_credential(e2e_sm: async_sessionmaker[AsyncSession], blob: bytes) -> None:
"""直接经真 store 在真 pg 写一行 oauth 凭据(用例 2/3 前置)。"""
from ww_api.services.credentials import STUB_OWNER_ID, SqlCredentialStore
async with e2e_sm() as session:
store = SqlCredentialStore(session)
await store.upsert_oauth_credential(STUB_OWNER_ID, KIMI_CODE_PROVIDER, blob)
# store 只 flush测试前置须显式提交否则独立 session建网关/校验)看不到(新契约)。
await session.commit()
async def test_k1_connect_device_flow_persists_encrypted_token(
e2e_sm: async_sessionmaker[AsyncSession],
monkeypatch: pytest.MonkeyPatch,
) -> None:
"""用例 1device flow 连接 → 后台轮询(pending×2→成功) → 加密 token 真落 pg → status 已连接。
全程零联网fake httpxtoken 绝不进 job 结果/状态响应/start 响应(负向断言)。
"""
import ww_api.routers.kimi_oauth as kimi_mod
from ww_api.main import create_app
from ww_api.services.credentials import (
AUTH_TYPE_OAUTH,
STUB_OWNER_ID,
SqlCredentialStore,
)
from ww_api.services.kimi_oauth import decrypt_oauth_bundle
from ww_api.services.project_deps import get_session_factory
enc_key = _enc_key()
# device auth(call0) → poll pending(call1) → poll pending(call2) → poll success(call3)。
http = _ScriptedHttp(
[
_device_auth_response(),
_pending_response(),
_pending_response(),
_token_success_response(),
]
)
# 后台 work 不真等 interval 秒。
async def _no_sleep(_seconds: float) -> None:
return None
monkeypatch.setattr("ww_api.routers.kimi_oauth.asyncio.sleep", _no_sleep)
# 后台 work 自建 http 走模块级 `_default_http_client()`(非 dep→ 指到同一 scripted fake。
monkeypatch.setattr(kimi_mod, "_default_http_client", lambda: http)
app = create_app()
# 端点 device-auth 走请求级生成器依赖 `_http_client`CR-H2→ override 指到同一 fake
# (保证与后台轮询顺序共享;否则真生成器会联网)。
app.dependency_overrides[kimi_mod._http_client] = lambda: http
# 后台 work 用 run_job 自建独立 session → 真 e2e_sm同测试 engine/loop真 store 落库。
app.dependency_overrides[get_session_factory] = lambda: e2e_sm
transport = httpx.ASGITransport(app=app)
try:
await _cleanup_credentials(e2e_sm)
async with LifespanManager(app):
async with httpx.AsyncClient(transport=transport, base_url="http://test") as client:
# 1) 发起 device flow → 202 + user_code/verification_uriASGITransport 等后台)。
start_resp = await client.post(_START_URL)
assert start_resp.status_code == 202
start_body = start_resp.json()
assert start_body["user_code"] == "ABCD-1234"
assert start_body["verification_uri"] == "https://auth.kimi.com/device"
assert start_body["interval"] == 5
job_id = start_body["job_id"]
# start 响应绝不含 token。
assert _ACCESS_TOKEN not in start_resp.text
assert _REFRESH_TOKEN not in start_resp.text
# 2) 轮询 job → doneresult 只 {connected, provider},无 token。
job_resp = await client.get(f"/jobs/{job_id}")
assert job_resp.status_code == 200
job = job_resp.json()
assert job["status"] == "done"
assert job["kind"] == "kimi_oauth"
assert job["result"] == {"connected": True, "provider": KIMI_CODE_PROVIDER}
# job 整体(含 result绝不泄露 token。
assert _ACCESS_TOKEN not in job_resp.text
assert _REFRESH_TOKEN not in job_resp.text
# 3) status → 已连接 + expires_at无 token 本体。
status_resp = await client.get(_STATUS_URL)
assert status_resp.status_code == 200
status_body = status_resp.json()
assert status_body["connected"] is True
assert status_body["expires_at"] is not None
assert "access_token" not in status_body
assert "refresh_token" not in status_body
assert _ACCESS_TOKEN not in status_resp.text
assert _REFRESH_TOKEN not in status_resp.text
# 后台经历两次 pending 才成功(证明 pending 路径真走device + 3 token = 4 调用。
assert len(http.calls) == 4
# DB 真源断言oauth 凭据行加密持久化(非明文),解密回环 token。
async with e2e_sm() as verify:
row = (
await verify.execute(
select(ProviderCredential).where(
ProviderCredential.owner_id == STUB_OWNER_ID,
ProviderCredential.provider == KIMI_CODE_PROVIDER,
)
)
).scalar_one()
assert row.auth_type == AUTH_TYPE_OAUTH
assert row.oauth_enc is not None
assert row.api_key_enc is None
# 密文不是明文token 不出现在原始字节里)。
assert _ACCESS_TOKEN.encode() not in row.oauth_enc
assert _REFRESH_TOKEN.encode() not in row.oauth_enc
# 解密回环:证明真的是这套 token 被加密持久化。
token = decrypt_oauth_bundle(row.oauth_enc, key=enc_key)
assert token.access_token == _ACCESS_TOKEN
assert token.refresh_token == _REFRESH_TOKEN
# 经真 store 读 status 也是已连接store 层一致)。
store = SqlCredentialStore(verify)
cred = await store.get_credential(STUB_OWNER_ID, KIMI_CODE_PROVIDER)
assert cred is not None
assert cred.auth_type == AUTH_TYPE_OAUTH
finally:
await _cleanup_credentials(e2e_sm)
async def test_k1_gateway_builds_kimi_code_adapter_with_headers_and_bearer(
e2e_sm: async_sessionmaker[AsyncSession],
) -> None:
"""用例 2存 oauth + 档位路由指向 kimi-code → 建网关 → 适配器带 coding base/伪造头/bearer。"""
from ww_api.services.credentials import SqlCredentialStore, StoredRouting
from ww_api.services.kimi_oauth import TokenSet, encrypt_oauth_bundle
from ww_api.services.project_deps import build_gateway_for_tier
enc_key = _enc_key()
fresh = TokenSet(
access_token=_ACCESS_TOKEN,
refresh_token=_REFRESH_TOKEN,
expires_at=datetime.now(UTC) + timedelta(seconds=900), # 远未过期 → 不刷新
)
blob = encrypt_oauth_bundle(fresh, key=enc_key)
try:
await _cleanup_credentials(e2e_sm)
await _seed_oauth_credential(e2e_sm, blob)
# 档位路由writer → kimi-code:kimi-for-codingDB 真源)。
async with e2e_sm() as seed:
store = SqlCredentialStore(seed)
await store.upsert_routing(
StoredRouting(
tier="writer",
provider=KIMI_CODE_PROVIDER,
model="kimi-for-coding",
fallback=[],
)
)
# store 只 flush显式提交使独立的建网关 session 可见该路由(新契约)。
await seed.commit()
async with e2e_sm() as session:
store = SqlCredentialStore(session)
gateway = await build_gateway_for_tier(session, store, "writer")
adapter = gateway._adapters[KIMI_CODE_PROVIDER]
assert isinstance(adapter, KimiCodeAdapter)
client = adapter._client
# coding base。
assert str(client.base_url).rstrip("/") == KIMI_CODE_BASE_URL.rstrip("/")
# 伪造头UA + 全 X-Msh-*。
for key in _REQUIRED_HEADER_KEYS:
assert key in client.default_headers
assert client.default_headers["User-Agent"] == KIMI_CODE_USER_AGENT
# bearer = 存量 access tokenSDK 据 api_key 自动发 Authorization: Bearer
assert client.api_key == _ACCESS_TOKEN
# 未刷新 → DB oauth 凭据未变(仍是原 token
async with e2e_sm() as verify:
from ww_api.services.kimi_oauth import decrypt_oauth_bundle
row = (
await verify.execute(
select(ProviderCredential).where(
ProviderCredential.provider == KIMI_CODE_PROVIDER
)
)
).scalar_one()
assert row.oauth_enc is not None
assert decrypt_oauth_bundle(row.oauth_enc, key=enc_key).access_token == _ACCESS_TOKEN
finally:
await _cleanup_credentials(e2e_sm)
async def test_k1_refresh_on_build_updates_stored_token(
e2e_sm: async_sessionmaker[AsyncSession],
monkeypatch: pytest.MonkeyPatch,
) -> None:
"""用例 3存量 token 临近过期 → 建适配器触发刷新 → 用新 token + DB oauth_enc 更新为新包。
刷新经 fake httpx + fake `kimi_refresh`(不联网);新包真落 pgDB 真源)。
"""
from ww_api.services import project_deps
from ww_api.services.credentials import STUB_OWNER_ID, SqlCredentialStore
from ww_api.services.kimi_oauth import (
TokenSet,
decrypt_oauth_bundle,
encrypt_oauth_bundle,
)
from ww_api.services.project_deps import _build_provider_adapter
enc_key = _enc_key()
near = TokenSet(
access_token=_ACCESS_TOKEN,
refresh_token=_REFRESH_TOKEN,
expires_at=datetime.now(UTC) + timedelta(seconds=30), # 临近过期 → 触发刷新
)
blob = encrypt_oauth_bundle(near, key=enc_key)
refreshed = TokenSet(
access_token=_NEW_ACCESS_TOKEN,
refresh_token=_NEW_REFRESH_TOKEN,
expires_at=datetime.now(UTC) + timedelta(seconds=900),
)
refresh_calls: list[str] = []
async def _fake_refresh(_http: Any, refresh_token: str) -> TokenSet:
# 证明用存量 refresh token 刷新。
refresh_calls.append(refresh_token)
return refreshed
class _FakeAsyncClient:
async def __aenter__(self) -> _FakeAsyncClient:
return self
async def __aexit__(self, *exc: object) -> None:
return None
monkeypatch.setattr(project_deps, "kimi_refresh", _fake_refresh)
monkeypatch.setattr(
"ww_api.services.project_deps.httpx.AsyncClient", lambda *a, **k: _FakeAsyncClient()
)
try:
await _cleanup_credentials(e2e_sm)
await _seed_oauth_credential(e2e_sm, blob)
async with e2e_sm() as session:
store = SqlCredentialStore(session)
adapter = await _build_provider_adapter(store, KIMI_CODE_PROVIDER)
# refresh 被调,且用存量 refresh token。
assert refresh_calls == [_REFRESH_TOKEN]
# 适配器用新 access token。
assert isinstance(adapter, KimiCodeAdapter)
assert adapter._client.api_key == _NEW_ACCESS_TOKEN
# DB 真源oauth_enc 已更新为新包(下次建网关复用刷新结果)。
async with e2e_sm() as verify:
row = (
await verify.execute(
select(ProviderCredential).where(
ProviderCredential.owner_id == STUB_OWNER_ID,
ProviderCredential.provider == KIMI_CODE_PROVIDER,
)
)
).scalar_one()
assert row.oauth_enc is not None
stored = decrypt_oauth_bundle(row.oauth_enc, key=enc_key)
assert stored.access_token == _NEW_ACCESS_TOKEN
assert stored.refresh_token == _NEW_REFRESH_TOKEN
finally:
await _cleanup_credentials(e2e_sm)
async def test_k1_disconnect_clears_credential(
e2e_sm: async_sessionmaker[AsyncSession],
) -> None:
"""用例 4disconnect → 凭据行被清除DB 真源:之后无 oauth 凭据)。"""
from ww_api.main import create_app
from ww_api.services.credentials import STUB_OWNER_ID
from ww_api.services.kimi_oauth import TokenSet, encrypt_oauth_bundle
enc_key = _enc_key()
token = TokenSet(
access_token=_ACCESS_TOKEN,
refresh_token=_REFRESH_TOKEN,
expires_at=datetime.now(UTC) + timedelta(seconds=900),
)
blob = encrypt_oauth_bundle(token, key=enc_key)
app = create_app()
transport = httpx.ASGITransport(app=app)
try:
await _cleanup_credentials(e2e_sm)
await _seed_oauth_credential(e2e_sm, blob)
# 前置:凭据确实在库。
async with e2e_sm() as pre:
assert (
await pre.execute(
select(ProviderCredential).where(
ProviderCredential.owner_id == STUB_OWNER_ID,
ProviderCredential.provider == KIMI_CODE_PROVIDER,
)
)
).scalar_one_or_none() is not None
async with LifespanManager(app):
async with httpx.AsyncClient(transport=transport, base_url="http://test") as client:
resp = await client.post(_DISCONNECT_URL)
assert resp.status_code == 200
assert resp.json()["disconnected"] is True
# DB 真源:凭据行已清除。
async with e2e_sm() as verify:
assert (
await verify.execute(
select(ProviderCredential).where(
ProviderCredential.owner_id == STUB_OWNER_ID,
ProviderCredential.provider == KIMI_CODE_PROVIDER,
)
)
).scalar_one_or_none() is None
finally:
await _cleanup_credentials(e2e_sm)