mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-02 17:30:23 +08:00
refactor: 拆分 Rust gateway/executor 大文件为模块目录结构,拆分 Python gateway.py 为子模块
Rust 侧: - executor.rs 拆分为 executor/ 目录 (plan_builders, stream, sync, submission) - 新增 kiro_stream/, local_finalize/, local_stream/ 模块目录 - 新增 video_tasks.rs - 测试文件 ai_execute/files/video 拆分为子目录 - handlers/headers/control/constants 扩展支持新模块 Python 侧: - gateway.py 拆分为 24 个子模块 (routes, shared, contract, chat, cli, video, files, finalize, reporting 等) - 新增 antigravity/gemini_cli/kiro 的 rust_http 适配层 - upstream_fetcher 增加 Rust sidecar 支持 - executor_plan/candidate/pipeline 适配调整 测试: - 对应拆分 test_internal_gateway_routes 为子目录 - 新增 rust_http 相关测试
This commit is contained in:
379
tests/api/internal_gateway/test_gateway_reporting_stream.py
Normal file
379
tests/api/internal_gateway/test_gateway_reporting_stream.py
Normal file
@@ -0,0 +1,379 @@
|
||||
"""Gateway internal report/finalize/trace route tests."""
|
||||
|
||||
import asyncio
|
||||
import base64
|
||||
import json
|
||||
import time
|
||||
from types import SimpleNamespace
|
||||
from typing import Any
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
import pytest
|
||||
from fastapi import BackgroundTasks, FastAPI
|
||||
from fastapi.responses import JSONResponse, StreamingResponse
|
||||
from fastapi.testclient import TestClient
|
||||
from sqlalchemy import create_engine
|
||||
from sqlalchemy.orm import sessionmaker
|
||||
|
||||
from src.api.internal.gateway import (
|
||||
CONTROL_ACTION_HEADER,
|
||||
CONTROL_ACTION_PROXY_PUBLIC,
|
||||
CONTROL_EXECUTED_HEADER,
|
||||
GatewayAuthContext,
|
||||
GatewayExecutionDecisionResponse,
|
||||
GatewayExecuteRequest,
|
||||
GatewayExecutionPlanResponse,
|
||||
GatewayResolveRequest,
|
||||
GatewayStreamReportRequest,
|
||||
GatewaySyncReportRequest,
|
||||
_dispatch_gateway_sync_telemetry,
|
||||
_build_gateway_sync_error_payload,
|
||||
_build_gateway_sync_telemetry_writer,
|
||||
_run_gateway_stream_report_background,
|
||||
_run_gateway_sync_report_background,
|
||||
_stream_executor_requires_python_rewrite,
|
||||
_build_claude_chat_sync_decision,
|
||||
_build_claude_chat_stream_decision,
|
||||
_build_claude_cli_sync_decision,
|
||||
_build_claude_cli_stream_decision,
|
||||
_build_gemini_files_download_stream_decision,
|
||||
_build_gemini_files_proxy_sync_decision,
|
||||
_build_gemini_chat_sync_decision,
|
||||
_build_gemini_chat_stream_decision,
|
||||
_build_gemini_cli_sync_decision,
|
||||
_build_gemini_cli_stream_decision,
|
||||
_build_openai_chat_sync_decision,
|
||||
_build_openai_chat_stream_decision,
|
||||
_build_openai_cli_stream_decision,
|
||||
_build_openai_video_content_stream_decision,
|
||||
_extract_gateway_sync_error_message,
|
||||
_record_gateway_direct_candidate_graph,
|
||||
_resolve_gateway_sync_error_status_code,
|
||||
_build_openai_chat_sync_plan,
|
||||
_build_openai_cli_sync_decision,
|
||||
_build_openai_cli_stream_plan,
|
||||
_build_openai_cli_sync_plan,
|
||||
_is_streaming_sync_payload,
|
||||
_resolve_auth_context,
|
||||
_resolve_gateway_sync_adapter,
|
||||
classify_gateway_route,
|
||||
router,
|
||||
)
|
||||
from src.database import get_db
|
||||
from src.models.database import Base, RequestCandidate
|
||||
from src.services.orchestration.candidate_resolver import CandidateResolver
|
||||
from src.services.scheduling.schemas import PoolCandidate, ProviderCandidate
|
||||
from src.services.request.executor_plan import ExecutionPlan, ExecutionPlanBody, PreparedExecutionPlan
|
||||
|
||||
|
||||
def _wait_until(predicate: Any, *, timeout: float = 1.0, interval: float = 0.01) -> None:
|
||||
deadline = time.time() + timeout
|
||||
while time.time() < deadline:
|
||||
if predicate():
|
||||
return
|
||||
time.sleep(interval)
|
||||
assert predicate()
|
||||
|
||||
|
||||
def test_report_stream_route_uses_lazy_session(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
app = FastAPI()
|
||||
app.include_router(router)
|
||||
monkeypatch.setattr("src.api.internal.gateway.ensure_loopback", lambda request: None)
|
||||
background_mock = AsyncMock(return_value=None)
|
||||
monkeypatch.setattr(
|
||||
"src.api.internal.gateway._run_gateway_stream_report_background_with_session",
|
||||
background_mock,
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
"src.api.internal.gateway.create_session",
|
||||
lambda: (_ for _ in ()).throw(AssertionError("create_session should not be called")),
|
||||
)
|
||||
|
||||
client = TestClient(app, base_url="http://127.0.0.1")
|
||||
response = client.post(
|
||||
"/api/internal/gateway/report-stream",
|
||||
json={
|
||||
"trace_id": "trace-report-stream-lazy-123",
|
||||
"report_kind": "openai_chat_stream_success",
|
||||
"report_context": {"user_id": "user-123", "api_key_id": "key-123"},
|
||||
"status_code": 200,
|
||||
"headers": {"content-type": "text/event-stream"},
|
||||
"body_base64": base64.b64encode(b"data: ok\n\n").decode("ascii"),
|
||||
},
|
||||
)
|
||||
|
||||
assert response.status_code == 200
|
||||
assert response.json() == {"ok": True}
|
||||
background_mock.assert_awaited_once()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_run_gateway_stream_report_background_ensures_request_candidate(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
db = object()
|
||||
ensure_mock = MagicMock(return_value=SimpleNamespace(id="cand-stream-123"))
|
||||
apply_mock = AsyncMock(return_value=None)
|
||||
monkeypatch.setattr(
|
||||
"src.api.internal.gateway._ensure_gateway_request_candidate",
|
||||
ensure_mock,
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
"src.api.internal.gateway._apply_gateway_stream_report",
|
||||
apply_mock,
|
||||
)
|
||||
|
||||
payload = GatewayStreamReportRequest(
|
||||
trace_id="trace-stream-candidate-123",
|
||||
report_kind="openai_chat_stream_success",
|
||||
report_context={
|
||||
"request_id": "req-stream-candidate-123",
|
||||
"candidate_id": "cand-stream-123",
|
||||
"provider_id": "provider-123",
|
||||
"endpoint_id": "endpoint-123",
|
||||
"key_id": "key-123",
|
||||
},
|
||||
status_code=200,
|
||||
headers={"content-type": "text/event-stream"},
|
||||
body_base64=base64.b64encode(b"data: ok\n\n").decode("ascii"),
|
||||
)
|
||||
|
||||
await _run_gateway_stream_report_background(payload, db)
|
||||
|
||||
ensure_mock.assert_called_once_with(
|
||||
db=db,
|
||||
report_context=dict(payload.report_context or {}),
|
||||
trace_id="trace-stream-candidate-123",
|
||||
initial_status="streaming",
|
||||
)
|
||||
apply_mock.assert_awaited_once_with(payload, db=db)
|
||||
|
||||
|
||||
def test_report_stream_route_records_openai_chat_stream_success(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
app = FastAPI()
|
||||
app.include_router(router)
|
||||
app.dependency_overrides[get_db] = lambda: object()
|
||||
monkeypatch.setattr("src.api.internal.gateway.ensure_loopback", lambda request: None)
|
||||
record_mock = AsyncMock(return_value=None)
|
||||
monkeypatch.setattr(
|
||||
"src.api.internal.gateway._record_gateway_openai_chat_stream_success",
|
||||
record_mock,
|
||||
)
|
||||
|
||||
client = TestClient(app, base_url="http://127.0.0.1")
|
||||
response = client.post(
|
||||
"/api/internal/gateway/report-stream",
|
||||
json={
|
||||
"trace_id": "trace-openai-chat-stream-report",
|
||||
"report_kind": "openai_chat_stream_success",
|
||||
"report_context": {"user_id": "user-123", "api_key_id": "key-123"},
|
||||
"status_code": 200,
|
||||
"headers": {"content-type": "text/event-stream"},
|
||||
"body_base64": base64.b64encode(b"data: hello\\n\\ndata: [DONE]\\n\\n").decode("ascii"),
|
||||
"telemetry": {"elapsed_ms": 123, "ttfb_ms": 22},
|
||||
},
|
||||
)
|
||||
|
||||
assert response.status_code == 200
|
||||
assert response.json() == {"ok": True}
|
||||
_wait_until(lambda: record_mock.await_count == 1)
|
||||
record_mock.assert_awaited_once()
|
||||
|
||||
|
||||
def test_report_stream_route_records_claude_chat_stream_success(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
app = FastAPI()
|
||||
app.include_router(router)
|
||||
app.dependency_overrides[get_db] = lambda: object()
|
||||
monkeypatch.setattr("src.api.internal.gateway.ensure_loopback", lambda request: None)
|
||||
record_mock = AsyncMock(return_value=None)
|
||||
monkeypatch.setattr(
|
||||
"src.api.internal.gateway._record_gateway_passthrough_chat_stream_success",
|
||||
record_mock,
|
||||
)
|
||||
|
||||
client = TestClient(app, base_url="http://127.0.0.1")
|
||||
response = client.post(
|
||||
"/api/internal/gateway/report-stream",
|
||||
json={
|
||||
"trace_id": "trace-claude-chat-stream-report",
|
||||
"report_kind": "claude_chat_stream_success",
|
||||
"report_context": {
|
||||
"user_id": "user-123",
|
||||
"api_key_id": "key-123",
|
||||
"client_api_format": "claude:chat",
|
||||
},
|
||||
"status_code": 200,
|
||||
"headers": {"content-type": "text/event-stream"},
|
||||
"body_base64": base64.b64encode(b"event: content_block_delta\n\n").decode("ascii"),
|
||||
"telemetry": {"elapsed_ms": 66},
|
||||
},
|
||||
)
|
||||
|
||||
assert response.status_code == 200
|
||||
assert response.json() == {"ok": True}
|
||||
_wait_until(lambda: record_mock.await_count == 1)
|
||||
record_mock.assert_awaited_once()
|
||||
|
||||
|
||||
def test_report_stream_route_records_gemini_chat_stream_success(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
app = FastAPI()
|
||||
app.include_router(router)
|
||||
app.dependency_overrides[get_db] = lambda: object()
|
||||
monkeypatch.setattr("src.api.internal.gateway.ensure_loopback", lambda request: None)
|
||||
record_mock = AsyncMock(return_value=None)
|
||||
monkeypatch.setattr(
|
||||
"src.api.internal.gateway._record_gateway_passthrough_chat_stream_success",
|
||||
record_mock,
|
||||
)
|
||||
|
||||
client = TestClient(app, base_url="http://127.0.0.1")
|
||||
response = client.post(
|
||||
"/api/internal/gateway/report-stream",
|
||||
json={
|
||||
"trace_id": "trace-gemini-chat-stream-report",
|
||||
"report_kind": "gemini_chat_stream_success",
|
||||
"report_context": {
|
||||
"user_id": "user-123",
|
||||
"api_key_id": "key-123",
|
||||
"client_api_format": "gemini:chat",
|
||||
},
|
||||
"status_code": 200,
|
||||
"headers": {"content-type": "text/event-stream"},
|
||||
"body_base64": base64.b64encode(b'data: {"candidates":[]}\n\n').decode("ascii"),
|
||||
"telemetry": {"elapsed_ms": 44},
|
||||
},
|
||||
)
|
||||
|
||||
assert response.status_code == 200
|
||||
assert response.json() == {"ok": True}
|
||||
_wait_until(lambda: record_mock.await_count == 1)
|
||||
record_mock.assert_awaited_once()
|
||||
|
||||
|
||||
def test_report_stream_route_records_openai_cli_stream_success(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
app = FastAPI()
|
||||
app.include_router(router)
|
||||
app.dependency_overrides[get_db] = lambda: object()
|
||||
monkeypatch.setattr("src.api.internal.gateway.ensure_loopback", lambda request: None)
|
||||
record_mock = AsyncMock(return_value=None)
|
||||
monkeypatch.setattr(
|
||||
"src.api.internal.gateway._record_gateway_passthrough_cli_stream_success",
|
||||
record_mock,
|
||||
)
|
||||
|
||||
client = TestClient(app, base_url="http://127.0.0.1")
|
||||
response = client.post(
|
||||
"/api/internal/gateway/report-stream",
|
||||
json={
|
||||
"trace_id": "trace-openai-cli-stream-report",
|
||||
"report_kind": "openai_cli_stream_success",
|
||||
"report_context": {
|
||||
"user_id": "user-123",
|
||||
"api_key_id": "key-123",
|
||||
"client_api_format": "openai:cli",
|
||||
},
|
||||
"status_code": 200,
|
||||
"headers": {"content-type": "text/event-stream"},
|
||||
"body_base64": base64.b64encode(
|
||||
b'event: response.completed\ndata: {"type":"response.completed"}\n\n'
|
||||
).decode("ascii"),
|
||||
"telemetry": {"elapsed_ms": 66},
|
||||
},
|
||||
)
|
||||
|
||||
assert response.status_code == 200
|
||||
assert response.json() == {"ok": True}
|
||||
_wait_until(lambda: record_mock.await_count == 1)
|
||||
record_mock.assert_awaited_once()
|
||||
|
||||
|
||||
def test_report_stream_route_records_claude_cli_stream_success(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
app = FastAPI()
|
||||
app.include_router(router)
|
||||
app.dependency_overrides[get_db] = lambda: object()
|
||||
monkeypatch.setattr("src.api.internal.gateway.ensure_loopback", lambda request: None)
|
||||
record_mock = AsyncMock(return_value=None)
|
||||
monkeypatch.setattr(
|
||||
"src.api.internal.gateway._record_gateway_passthrough_cli_stream_success",
|
||||
record_mock,
|
||||
)
|
||||
|
||||
client = TestClient(app, base_url="http://127.0.0.1")
|
||||
response = client.post(
|
||||
"/api/internal/gateway/report-stream",
|
||||
json={
|
||||
"trace_id": "trace-claude-cli-stream-report",
|
||||
"report_kind": "claude_cli_stream_success",
|
||||
"report_context": {
|
||||
"user_id": "user-123",
|
||||
"api_key_id": "key-123",
|
||||
"client_api_format": "claude:cli",
|
||||
},
|
||||
"status_code": 200,
|
||||
"headers": {"content-type": "text/event-stream"},
|
||||
"body_base64": base64.b64encode(
|
||||
b'event: message_start\ndata: {"type":"message_start"}\n\n'
|
||||
).decode("ascii"),
|
||||
"telemetry": {"elapsed_ms": 44},
|
||||
},
|
||||
)
|
||||
|
||||
assert response.status_code == 200
|
||||
assert response.json() == {"ok": True}
|
||||
_wait_until(lambda: record_mock.await_count == 1)
|
||||
record_mock.assert_awaited_once()
|
||||
|
||||
|
||||
def test_report_stream_route_records_gemini_cli_stream_success(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
app = FastAPI()
|
||||
app.include_router(router)
|
||||
app.dependency_overrides[get_db] = lambda: object()
|
||||
monkeypatch.setattr("src.api.internal.gateway.ensure_loopback", lambda request: None)
|
||||
record_mock = AsyncMock(return_value=None)
|
||||
monkeypatch.setattr(
|
||||
"src.api.internal.gateway._record_gateway_passthrough_cli_stream_success",
|
||||
record_mock,
|
||||
)
|
||||
|
||||
client = TestClient(app, base_url="http://127.0.0.1")
|
||||
response = client.post(
|
||||
"/api/internal/gateway/report-stream",
|
||||
json={
|
||||
"trace_id": "trace-gemini-cli-stream-report",
|
||||
"report_kind": "gemini_cli_stream_success",
|
||||
"report_context": {
|
||||
"user_id": "user-123",
|
||||
"api_key_id": "key-123",
|
||||
"client_api_format": "gemini:cli",
|
||||
},
|
||||
"status_code": 200,
|
||||
"headers": {"content-type": "text/event-stream"},
|
||||
"body_base64": base64.b64encode(b'data: {"candidates":[]}\n\n').decode("ascii"),
|
||||
"telemetry": {"elapsed_ms": 44},
|
||||
},
|
||||
)
|
||||
|
||||
assert response.status_code == 200
|
||||
assert response.json() == {"ok": True}
|
||||
_wait_until(lambda: record_mock.await_count == 1)
|
||||
record_mock.assert_awaited_once()
|
||||
Reference in New Issue
Block a user