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 服务层和测试适配新架构命名
326 lines
11 KiB
Python
326 lines
11 KiB
Python
from __future__ import annotations
|
|
|
|
import json
|
|
from collections.abc import AsyncGenerator
|
|
from types import SimpleNamespace
|
|
from unittest.mock import AsyncMock
|
|
|
|
import pytest
|
|
from fastapi import Response
|
|
from fastapi.responses import StreamingResponse
|
|
|
|
import src.api.public.gemini_files as gemini_files_mod
|
|
import src.services.proxy_node.resolver as resolver_mod
|
|
import src.services.request.execution_runtime_client as rust_client_mod
|
|
from src.api.public.gemini_files import UpstreamContext
|
|
from src.config.settings import config
|
|
from src.services.request.execution_runtime_plan import ExecutionProxySnapshot
|
|
from src.services.request.execution_runtime_client import (
|
|
ExecutionRuntimeStreamResult,
|
|
ExecutionRuntimeSyncResult,
|
|
)
|
|
|
|
|
|
class _DummyStreamResponseCtx:
|
|
def __init__(self) -> None:
|
|
self.closed = False
|
|
|
|
async def __aexit__(self, exc_type: object, exc: object, tb: object) -> None:
|
|
self.closed = True
|
|
|
|
|
|
class _FakeDBContext:
|
|
def __init__(self, db: object) -> None:
|
|
self._db = db
|
|
|
|
def __enter__(self) -> object:
|
|
return self._db
|
|
|
|
def __exit__(self, exc_type: object, exc: object, tb: object) -> bool:
|
|
return False
|
|
|
|
|
|
async def _iter_chunks(chunks: list[bytes]) -> AsyncGenerator[bytes]:
|
|
for chunk in chunks:
|
|
yield chunk
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_enrich_upstream_context_proxy_builds_tunnel_snapshot(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
ctx = UpstreamContext(
|
|
upstream_key="upstream-key",
|
|
base_url="https://generativelanguage.googleapis.com",
|
|
file_key_id="key-1",
|
|
user_id="user-1",
|
|
provider_id="prov-1",
|
|
endpoint_id="ep-1",
|
|
provider_proxy={"enabled": True, "node_id": "node-1"},
|
|
key_proxy=None,
|
|
)
|
|
|
|
monkeypatch.setattr(
|
|
resolver_mod,
|
|
"resolve_effective_proxy",
|
|
lambda provider_proxy, key_proxy=None: provider_proxy,
|
|
)
|
|
monkeypatch.setattr(
|
|
resolver_mod,
|
|
"get_system_proxy_config_async",
|
|
AsyncMock(return_value=None),
|
|
)
|
|
monkeypatch.setattr(
|
|
resolver_mod,
|
|
"resolve_delegate_config_async",
|
|
AsyncMock(return_value={"tunnel": True, "node_id": "node-1"}),
|
|
)
|
|
monkeypatch.setattr(
|
|
resolver_mod,
|
|
"resolve_proxy_info_async",
|
|
AsyncMock(return_value={"mode": "tunnel", "node_id": "node-1", "label": "Node 1"}),
|
|
)
|
|
|
|
async def _unexpected_build_proxy_url_async(proxy_config: object) -> str:
|
|
raise AssertionError(f"proxy url should not be built for tunnel: {proxy_config!r}")
|
|
|
|
monkeypatch.setattr(
|
|
resolver_mod,
|
|
"build_proxy_url_async",
|
|
_unexpected_build_proxy_url_async,
|
|
)
|
|
|
|
enriched = await gemini_files_mod._enrich_upstream_context_proxy(ctx)
|
|
|
|
assert enriched.proxy_config == {"enabled": True, "node_id": "node-1"}
|
|
assert enriched.delegate_config == {"tunnel": True, "node_id": "node-1"}
|
|
assert enriched.proxy_snapshot is not None
|
|
assert enriched.proxy_snapshot.mode == "tunnel"
|
|
assert enriched.proxy_snapshot.node_id == "node-1"
|
|
assert enriched.proxy_snapshot.url is None
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_proxy_request_passes_proxy_snapshot_to_rust_executor(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
monkeypatch.setattr(config, "executor_backend", "rust")
|
|
proxy_snapshot = ExecutionProxySnapshot(
|
|
enabled=True,
|
|
mode="http",
|
|
url="http://proxy.local:8080",
|
|
)
|
|
|
|
async def _fake_execute_sync_json(self: object, plan: object) -> ExecutionRuntimeSyncResult:
|
|
assert getattr(plan, "method") == "GET"
|
|
assert getattr(plan, "provider_id") == "prov-1"
|
|
assert getattr(plan, "endpoint_id") == "ep-1"
|
|
assert getattr(plan, "proxy").url == "http://proxy.local:8080"
|
|
return ExecutionRuntimeSyncResult(
|
|
status_code=200,
|
|
headers={"content-type": "application/json", "x-rust-files": "true"},
|
|
response_json={"files": [{"name": "files/abc"}]},
|
|
)
|
|
|
|
monkeypatch.setattr(
|
|
rust_client_mod.ExecutionRuntimeClient,
|
|
"execute_sync_json",
|
|
_fake_execute_sync_json,
|
|
)
|
|
response = await gemini_files_mod._proxy_request(
|
|
"GET",
|
|
"https://generativelanguage.googleapis.com/v1beta/files",
|
|
{"x-goog-api-key": "upstream-key"},
|
|
file_key_id="key-1",
|
|
user_id="user-1",
|
|
provider_id="prov-1",
|
|
endpoint_id="ep-1",
|
|
proxy=proxy_snapshot,
|
|
)
|
|
|
|
assert response.status_code == 200
|
|
assert response.headers["x-rust-files"] == "true"
|
|
assert json.loads(response.body) == {"files": [{"name": "files/abc"}]}
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_proxy_request_returns_503_when_rust_unavailable(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
async def _fake_try_rust_sync_proxy_request(*args: object, **kwargs: object) -> Response:
|
|
del args, kwargs
|
|
return gemini_files_mod._build_rust_unavailable_response()
|
|
|
|
monkeypatch.setattr(
|
|
gemini_files_mod,
|
|
"_try_rust_sync_proxy_request",
|
|
_fake_try_rust_sync_proxy_request,
|
|
)
|
|
response = await gemini_files_mod._proxy_request(
|
|
"GET",
|
|
"https://generativelanguage.googleapis.com/v1beta/files",
|
|
{"x-goog-api-key": "upstream-key"},
|
|
file_key_id="key-1",
|
|
user_id="user-1",
|
|
)
|
|
|
|
body = json.loads(response.body)
|
|
assert response.status_code == 503
|
|
assert body["error"]["code"] == 503
|
|
assert body["error"]["status"] == "UNAVAILABLE"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_download_file_uses_enriched_proxy_snapshot_for_regular_files(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
monkeypatch.setattr(config, "executor_backend", "rust")
|
|
dummy_ctx = _DummyStreamResponseCtx()
|
|
raw_ctx = UpstreamContext(
|
|
upstream_key="upstream-key",
|
|
base_url="https://generativelanguage.googleapis.com",
|
|
file_key_id="key-1",
|
|
user_id="user-1",
|
|
provider_id="prov-1",
|
|
endpoint_id="ep-1",
|
|
)
|
|
enriched_ctx = UpstreamContext(
|
|
upstream_key="upstream-key",
|
|
base_url="https://generativelanguage.googleapis.com",
|
|
file_key_id="key-1",
|
|
user_id="user-1",
|
|
provider_id="prov-1",
|
|
endpoint_id="ep-1",
|
|
proxy_config={"enabled": True, "url": "http://proxy.local:8080"},
|
|
delegate_config=None,
|
|
proxy_snapshot=ExecutionProxySnapshot(
|
|
enabled=True,
|
|
mode="http",
|
|
url="http://proxy.local:8080",
|
|
),
|
|
)
|
|
|
|
monkeypatch.setattr(gemini_files_mod, "_extract_gemini_api_key", lambda request: "client-key")
|
|
monkeypatch.setattr(
|
|
gemini_files_mod,
|
|
"create_session",
|
|
lambda: _FakeDBContext(SimpleNamespace()),
|
|
)
|
|
monkeypatch.setattr(
|
|
gemini_files_mod.AuthService,
|
|
"authenticate_api_key",
|
|
lambda db, key: (SimpleNamespace(id="user-1"), SimpleNamespace(id="user-api-key")),
|
|
)
|
|
monkeypatch.setattr(gemini_files_mod, "_ensure_balance_access", lambda db, user, api_key: None)
|
|
monkeypatch.setattr(
|
|
gemini_files_mod,
|
|
"_resolve_upstream_context",
|
|
AsyncMock(return_value=raw_ctx),
|
|
)
|
|
monkeypatch.setattr(
|
|
gemini_files_mod,
|
|
"_enrich_upstream_context_proxy",
|
|
AsyncMock(return_value=enriched_ctx),
|
|
)
|
|
|
|
async def _fake_execute_stream(self: object, plan: object) -> ExecutionRuntimeStreamResult:
|
|
assert getattr(plan, "method") == "GET"
|
|
assert getattr(plan, "url") == (
|
|
"https://generativelanguage.googleapis.com/v1beta/files/file-1:download?alt=media"
|
|
)
|
|
assert getattr(plan, "headers") == {"x-goog-api-key": "upstream-key"}
|
|
assert getattr(plan, "proxy").url == "http://proxy.local:8080"
|
|
return ExecutionRuntimeStreamResult(
|
|
status_code=200,
|
|
headers={"content-type": "application/octet-stream", "x-rust-files": "true"},
|
|
byte_iterator=_iter_chunks([b"file-", b"bytes"]),
|
|
response_ctx=dummy_ctx,
|
|
)
|
|
|
|
monkeypatch.setattr(
|
|
rust_client_mod.ExecutionRuntimeClient,
|
|
"execute_stream",
|
|
_fake_execute_stream,
|
|
)
|
|
response = await gemini_files_mod._download_file_response(
|
|
"file-1",
|
|
SimpleNamespace(
|
|
headers={},
|
|
query_params={"alt": "media"},
|
|
),
|
|
)
|
|
|
|
assert isinstance(response, StreamingResponse)
|
|
assert response.headers["x-rust-files"] == "true"
|
|
body = b"".join([chunk async for chunk in response.body_iterator])
|
|
assert body == b"file-bytes"
|
|
assert dummy_ctx.closed is True
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_download_file_returns_503_when_rust_stream_unavailable(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
monkeypatch.setattr(config, "executor_backend", "rust")
|
|
raw_ctx = UpstreamContext(
|
|
upstream_key="upstream-key",
|
|
base_url="https://generativelanguage.googleapis.com",
|
|
file_key_id="key-1",
|
|
user_id="user-1",
|
|
provider_id="prov-1",
|
|
endpoint_id="ep-1",
|
|
)
|
|
enriched_ctx = UpstreamContext(
|
|
upstream_key="upstream-key",
|
|
base_url="https://generativelanguage.googleapis.com",
|
|
file_key_id="key-1",
|
|
user_id="user-1",
|
|
provider_id="prov-1",
|
|
endpoint_id="ep-1",
|
|
)
|
|
|
|
monkeypatch.setattr(gemini_files_mod, "_extract_gemini_api_key", lambda request: "client-key")
|
|
monkeypatch.setattr(
|
|
gemini_files_mod,
|
|
"create_session",
|
|
lambda: _FakeDBContext(SimpleNamespace()),
|
|
)
|
|
monkeypatch.setattr(
|
|
gemini_files_mod.AuthService,
|
|
"authenticate_api_key",
|
|
lambda db, key: (SimpleNamespace(id="user-1"), SimpleNamespace(id="user-api-key")),
|
|
)
|
|
monkeypatch.setattr(gemini_files_mod, "_ensure_balance_access", lambda db, user, api_key: None)
|
|
monkeypatch.setattr(
|
|
gemini_files_mod,
|
|
"_resolve_upstream_context",
|
|
AsyncMock(return_value=raw_ctx),
|
|
)
|
|
monkeypatch.setattr(
|
|
gemini_files_mod,
|
|
"_enrich_upstream_context_proxy",
|
|
AsyncMock(return_value=enriched_ctx),
|
|
)
|
|
|
|
async def _fake_execute_stream(self: object, plan: object) -> ExecutionRuntimeStreamResult:
|
|
del self, plan
|
|
raise rust_client_mod.ExecutionRuntimeClientError("executor unavailable")
|
|
|
|
monkeypatch.setattr(
|
|
rust_client_mod.ExecutionRuntimeClient,
|
|
"execute_stream",
|
|
_fake_execute_stream,
|
|
)
|
|
response = await gemini_files_mod._download_file_response(
|
|
"file-1",
|
|
SimpleNamespace(
|
|
headers={},
|
|
query_params={"alt": "media"},
|
|
),
|
|
)
|
|
|
|
body = json.loads(response.body)
|
|
assert response.status_code == 503
|
|
assert body["error"]["code"] == 503
|
|
assert body["error"]["status"] == "UNAVAILABLE"
|