mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-02 09:20:22 +08:00
8f26e1a31f
- 删除 aether-hub、aether-proxy 独立项目及其 Dockerfile/配置 - 删除 crates/aether-executor 和 crates/aether-gateway 全部模块 - 新增 apps/ 目录作为应用入口 - 将 hub 概念重构为 gateway tunnel transport - 将 executor 重构为 execution runtime - 新增 tunnel.rs 合约定义和 testkit tunnel/execution_runtime 模块 - 更新 Python 服务层和测试适配新架构命名
388 lines
13 KiB
Python
388 lines
13 KiB
Python
"""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.execution_runtime_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 _make_legacy_test_client(app: FastAPI) -> TestClient:
|
|
return TestClient(
|
|
app,
|
|
base_url="http://127.0.0.1",
|
|
headers={"x-aether-legacy-internal-gateway": "true"},
|
|
)
|
|
|
|
|
|
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 = _make_legacy_test_client(app)
|
|
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 = _make_legacy_test_client(app)
|
|
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 = _make_legacy_test_client(app)
|
|
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 = _make_legacy_test_client(app)
|
|
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 = _make_legacy_test_client(app)
|
|
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 = _make_legacy_test_client(app)
|
|
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 = _make_legacy_test_client(app)
|
|
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()
|