mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-01 17:00:21 +08:00
feat: 引入 Rust executor/gateway sidecar 及 Python 侧双后端适配
- 新增 Rust workspace crates: aether-contracts, aether-executor, aether-gateway - aether-executor: 支持 Unix Socket/TCP 双传输模式,处理同步/流式上游请求 - aether-gateway: 作为本地主入口代理,集成 /api/internal/gateway/resolve 认证预解析 - Python 侧新增 ExecutionPlan 契约和 RustExecutorClient,各 handler 支持 executor_backend=rust 时将可序列化请求转发给 Rust executor 执行 - 重构 dev.sh 支持 executor/gateway 进程编排与生命周期管理 - 新增 internal gateway 路由,提供 resolve/passthrough 端点 - handler 层(chat/cli/video/endpoint_checker 等)全面适配 Rust executor 回退逻辑 - pipeline 层支持 trusted auth context 跳过重复认证 - 新增 Rust CI workflow 及对应测试用例
This commit is contained in:
565
tests/api/handlers/base/test_chat_rust_executor.py
Normal file
565
tests/api/handlers/base/test_chat_rust_executor.py
Normal file
@@ -0,0 +1,565 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from types import SimpleNamespace
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
|
||||
import src.api.handlers.base.chat_sync_executor as chat_sync_mod
|
||||
from src.api.handlers.base.chat_sync_executor import ChatSyncExecutor
|
||||
from src.core.exceptions import EmbeddedErrorException
|
||||
from src.services.request.executor_plan import (
|
||||
ExecutionPlan,
|
||||
ExecutionPlanBody,
|
||||
PreparedExecutionPlan,
|
||||
)
|
||||
from src.services.request.rust_executor_client import (
|
||||
RustExecutorClientError,
|
||||
RustExecutorSyncResult,
|
||||
)
|
||||
|
||||
|
||||
class _FakeEnvelope:
|
||||
name = "fake-envelope"
|
||||
|
||||
def __init__(self) -> None:
|
||||
self.status_codes: list[int] = []
|
||||
self.postprocessed_payloads: list[dict[str, object]] = []
|
||||
|
||||
def on_http_status(self, *, base_url: str | None, status_code: int) -> None:
|
||||
self.status_codes.append(status_code)
|
||||
|
||||
def on_connection_error(self, *, base_url: str | None, exc: Exception) -> None:
|
||||
raise AssertionError("connection error hook should not be used in this test")
|
||||
|
||||
def unwrap_response(self, data: dict[str, object]) -> dict[str, object]:
|
||||
return dict(data["payload"]) # type: ignore[index]
|
||||
|
||||
def postprocess_unwrapped_response(self, *, model: str, data: dict[str, object]) -> None:
|
||||
self.postprocessed_payloads.append(dict(data))
|
||||
|
||||
|
||||
class _FakeNormalizer:
|
||||
def response_from_internal(
|
||||
self,
|
||||
internal_resp: object,
|
||||
*,
|
||||
requested_model: str,
|
||||
) -> dict[str, object]:
|
||||
return {
|
||||
"aggregated": True,
|
||||
"requested_model": requested_model,
|
||||
"internal_id": getattr(internal_resp, "id", "missing"),
|
||||
}
|
||||
|
||||
|
||||
def _make_prepared_plan() -> PreparedExecutionPlan:
|
||||
return PreparedExecutionPlan(
|
||||
contract=ExecutionPlan(
|
||||
request_id="req-test",
|
||||
candidate_id=None,
|
||||
provider_name="openai",
|
||||
provider_id="prov-1",
|
||||
endpoint_id="ep-1",
|
||||
key_id="key-1",
|
||||
method="POST",
|
||||
url="https://example.com/v1/chat/completions",
|
||||
headers={"content-type": "application/json"},
|
||||
body=ExecutionPlanBody(json_body={"model": "gpt-4.1"}),
|
||||
stream=False,
|
||||
provider_api_format="openai:chat",
|
||||
client_api_format="openai:chat",
|
||||
model_name="gpt-4.1",
|
||||
),
|
||||
payload={"model": "gpt-4.1"},
|
||||
headers={"content-type": "application/json"},
|
||||
upstream_is_stream=False,
|
||||
needs_conversion=False,
|
||||
provider_type="openai",
|
||||
request_timeout=30.0,
|
||||
)
|
||||
|
||||
|
||||
def _make_proxy_prepared_plan() -> PreparedExecutionPlan:
|
||||
prepared = _make_prepared_plan()
|
||||
prepared.contract.proxy = chat_sync_mod.ExecutionProxySnapshot(
|
||||
enabled=True,
|
||||
mode="http",
|
||||
label="proxy.internal",
|
||||
url="http://proxy.internal:8080",
|
||||
)
|
||||
prepared.proxy_config = {"url": "http://proxy.internal:8080"}
|
||||
return prepared
|
||||
|
||||
|
||||
def _make_tunnel_prepared_plan() -> PreparedExecutionPlan:
|
||||
prepared = _make_prepared_plan()
|
||||
prepared.contract.proxy = chat_sync_mod.ExecutionProxySnapshot(
|
||||
enabled=True,
|
||||
mode="tunnel",
|
||||
node_id="node-1",
|
||||
label="relay-node",
|
||||
)
|
||||
prepared.delegate_config = {"tunnel": True, "node_id": "node-1"}
|
||||
prepared.proxy_config = {"node_id": "node-1"}
|
||||
return prepared
|
||||
|
||||
|
||||
def _make_upstream_stream_prepared_plan() -> PreparedExecutionPlan:
|
||||
prepared = _make_prepared_plan()
|
||||
prepared.contract.stream = True
|
||||
prepared.upstream_is_stream = True
|
||||
return prepared
|
||||
|
||||
|
||||
def _make_tls_prepared_plan() -> PreparedExecutionPlan:
|
||||
prepared = _make_prepared_plan()
|
||||
prepared.contract.tls_profile = "claude_code_nodejs"
|
||||
prepared.provider_type = "claude_code"
|
||||
return prepared
|
||||
|
||||
|
||||
def _make_executor() -> ChatSyncExecutor:
|
||||
handler = SimpleNamespace(request_id="req-test")
|
||||
executor = ChatSyncExecutor(handler)
|
||||
executor._ctx.provider_api_format_for_error = "openai:chat"
|
||||
executor._ctx.client_api_format_for_error = "openai:chat"
|
||||
executor._ctx.needs_conversion_for_error = False
|
||||
return executor
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_execute_sync_plan_uses_rust_executor_when_available(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
executor = _make_executor()
|
||||
prepared_plan = _make_prepared_plan()
|
||||
|
||||
monkeypatch.setattr(chat_sync_mod.config, "executor_backend", "rust")
|
||||
|
||||
async def _fake_execute_sync_json(self: object, plan: ExecutionPlan) -> RustExecutorSyncResult:
|
||||
assert plan.request_id == "req-test"
|
||||
return RustExecutorSyncResult(
|
||||
status_code=200,
|
||||
response_json={"id": "chatcmpl-1"},
|
||||
headers={"content-type": "application/json"},
|
||||
)
|
||||
|
||||
async def _should_not_fallback(**kwargs: object) -> dict[str, object]:
|
||||
raise AssertionError("local execution should not be used")
|
||||
|
||||
monkeypatch.setattr(
|
||||
chat_sync_mod.RustExecutorClient,
|
||||
"execute_sync_json",
|
||||
_fake_execute_sync_json,
|
||||
)
|
||||
monkeypatch.setattr(executor, "_execute_sync_plan_locally", _should_not_fallback)
|
||||
|
||||
response = await executor._execute_sync_plan(
|
||||
prepared_plan=prepared_plan,
|
||||
provider=SimpleNamespace(name="provider"),
|
||||
model="gpt-4.1",
|
||||
)
|
||||
|
||||
assert response == {"id": "chatcmpl-1"}
|
||||
assert executor._ctx.status_code == 200
|
||||
assert executor._ctx.response_json == {"id": "chatcmpl-1"}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_execute_sync_plan_allows_supported_proxy_urls_for_rust(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
executor = _make_executor()
|
||||
prepared_plan = _make_proxy_prepared_plan()
|
||||
|
||||
monkeypatch.setattr(chat_sync_mod.config, "executor_backend", "rust")
|
||||
|
||||
async def _fake_execute_sync_json(self: object, plan: ExecutionPlan) -> RustExecutorSyncResult:
|
||||
assert plan.proxy is not None
|
||||
assert plan.proxy.url == "http://proxy.internal:8080"
|
||||
return RustExecutorSyncResult(
|
||||
status_code=200,
|
||||
response_json={"id": "chatcmpl-proxy"},
|
||||
headers={"content-type": "application/json"},
|
||||
)
|
||||
|
||||
async def _should_not_fallback(**kwargs: object) -> dict[str, object]:
|
||||
raise AssertionError("local execution should not be used")
|
||||
|
||||
monkeypatch.setattr(
|
||||
chat_sync_mod.RustExecutorClient,
|
||||
"execute_sync_json",
|
||||
_fake_execute_sync_json,
|
||||
)
|
||||
monkeypatch.setattr(executor, "_execute_sync_plan_locally", _should_not_fallback)
|
||||
|
||||
response = await executor._execute_sync_plan(
|
||||
prepared_plan=prepared_plan,
|
||||
provider=SimpleNamespace(name="provider"),
|
||||
model="gpt-4.1",
|
||||
)
|
||||
|
||||
assert response == {"id": "chatcmpl-proxy"}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_execute_sync_plan_allows_tunnel_delegate_for_rust(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
executor = _make_executor()
|
||||
prepared_plan = _make_tunnel_prepared_plan()
|
||||
|
||||
monkeypatch.setattr(chat_sync_mod.config, "executor_backend", "rust")
|
||||
|
||||
async def _fake_execute_sync_json(self: object, plan: ExecutionPlan) -> RustExecutorSyncResult:
|
||||
assert plan.proxy is not None
|
||||
assert plan.proxy.mode == "tunnel"
|
||||
assert plan.proxy.node_id == "node-1"
|
||||
return RustExecutorSyncResult(
|
||||
status_code=200,
|
||||
response_json={"id": "chatcmpl-tunnel"},
|
||||
headers={"content-type": "application/json"},
|
||||
)
|
||||
|
||||
async def _should_not_fallback(**kwargs: object) -> dict[str, object]:
|
||||
raise AssertionError("local execution should not be used")
|
||||
|
||||
monkeypatch.setattr(
|
||||
chat_sync_mod.RustExecutorClient,
|
||||
"execute_sync_json",
|
||||
_fake_execute_sync_json,
|
||||
)
|
||||
monkeypatch.setattr(executor, "_execute_sync_plan_locally", _should_not_fallback)
|
||||
|
||||
response = await executor._execute_sync_plan(
|
||||
prepared_plan=prepared_plan,
|
||||
provider=SimpleNamespace(name="provider"),
|
||||
model="gpt-4.1",
|
||||
)
|
||||
|
||||
assert response == {"id": "chatcmpl-tunnel"}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_execute_sync_plan_allows_tls_profile_for_rust(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
executor = _make_executor()
|
||||
prepared_plan = _make_tls_prepared_plan()
|
||||
|
||||
monkeypatch.setattr(chat_sync_mod.config, "executor_backend", "rust")
|
||||
|
||||
async def _fake_execute_sync_json(self: object, plan: ExecutionPlan) -> RustExecutorSyncResult:
|
||||
assert plan.tls_profile == "claude_code_nodejs"
|
||||
return RustExecutorSyncResult(
|
||||
status_code=200,
|
||||
response_json={"id": "chatcmpl-tls"},
|
||||
headers={"content-type": "application/json"},
|
||||
)
|
||||
|
||||
async def _should_not_fallback(**kwargs: object) -> dict[str, object]:
|
||||
raise AssertionError("local execution should not be used")
|
||||
|
||||
monkeypatch.setattr(
|
||||
chat_sync_mod.RustExecutorClient,
|
||||
"execute_sync_json",
|
||||
_fake_execute_sync_json,
|
||||
)
|
||||
monkeypatch.setattr(executor, "_execute_sync_plan_locally", _should_not_fallback)
|
||||
|
||||
response = await executor._execute_sync_plan(
|
||||
prepared_plan=prepared_plan,
|
||||
provider=SimpleNamespace(name="provider"),
|
||||
model="claude-3.7-sonnet",
|
||||
)
|
||||
|
||||
assert response == {"id": "chatcmpl-tls"}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_execute_sync_plan_applies_envelope_postprocessing_after_rust(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
executor = _make_executor()
|
||||
prepared_plan = _make_prepared_plan()
|
||||
prepared_plan.envelope = _FakeEnvelope()
|
||||
|
||||
monkeypatch.setattr(chat_sync_mod.config, "executor_backend", "rust")
|
||||
|
||||
async def _fake_execute_sync_json(self: object, plan: ExecutionPlan) -> RustExecutorSyncResult:
|
||||
return RustExecutorSyncResult(
|
||||
status_code=200,
|
||||
response_json={"payload": {"id": "wrapped-1", "message": "ok"}},
|
||||
headers={"content-type": "application/json"},
|
||||
)
|
||||
|
||||
async def _should_not_fallback(**kwargs: object) -> dict[str, object]:
|
||||
raise AssertionError("local execution should not be used")
|
||||
|
||||
monkeypatch.setattr(
|
||||
chat_sync_mod.RustExecutorClient,
|
||||
"execute_sync_json",
|
||||
_fake_execute_sync_json,
|
||||
)
|
||||
monkeypatch.setattr(executor, "_execute_sync_plan_locally", _should_not_fallback)
|
||||
|
||||
response = await executor._execute_sync_plan(
|
||||
prepared_plan=prepared_plan,
|
||||
provider=SimpleNamespace(name="provider"),
|
||||
model="gpt-4.1",
|
||||
)
|
||||
|
||||
assert response == {"id": "wrapped-1", "message": "ok"}
|
||||
assert prepared_plan.envelope.status_codes == [200]
|
||||
assert prepared_plan.envelope.postprocessed_payloads == [
|
||||
{"id": "wrapped-1", "message": "ok"}
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_execute_sync_plan_applies_format_conversion_after_rust(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
executor = _make_executor()
|
||||
prepared_plan = _make_prepared_plan()
|
||||
prepared_plan.needs_conversion = True
|
||||
prepared_plan.contract.provider_api_format = "gemini:chat"
|
||||
prepared_plan.contract.client_api_format = "openai:chat"
|
||||
executor._ctx.provider_api_format_for_error = "gemini:chat"
|
||||
executor._ctx.client_api_format_for_error = "openai:chat"
|
||||
executor._ctx.needs_conversion_for_error = True
|
||||
|
||||
monkeypatch.setattr(chat_sync_mod.config, "executor_backend", "rust")
|
||||
|
||||
class _FakeRegistry:
|
||||
def convert_response(
|
||||
self,
|
||||
response_json: dict[str, object],
|
||||
source_format: str,
|
||||
target_format: str,
|
||||
*,
|
||||
requested_model: str,
|
||||
) -> dict[str, object]:
|
||||
assert source_format == "gemini:chat"
|
||||
assert target_format == "openai:chat"
|
||||
assert requested_model == "gpt-4.1"
|
||||
return {
|
||||
"converted": True,
|
||||
"source_id": response_json["provider_id"],
|
||||
}
|
||||
|
||||
async def _fake_execute_sync_json(self: object, plan: ExecutionPlan) -> RustExecutorSyncResult:
|
||||
assert plan.provider_api_format == "gemini:chat"
|
||||
return RustExecutorSyncResult(
|
||||
status_code=200,
|
||||
response_json={"provider_id": "gemini-1"},
|
||||
headers={"content-type": "application/json"},
|
||||
)
|
||||
|
||||
async def _should_not_fallback(**kwargs: object) -> dict[str, object]:
|
||||
raise AssertionError("local execution should not be used")
|
||||
|
||||
monkeypatch.setattr(
|
||||
chat_sync_mod.RustExecutorClient,
|
||||
"execute_sync_json",
|
||||
_fake_execute_sync_json,
|
||||
)
|
||||
monkeypatch.setattr(chat_sync_mod, "get_format_converter_registry", lambda: _FakeRegistry())
|
||||
monkeypatch.setattr(executor, "_execute_sync_plan_locally", _should_not_fallback)
|
||||
|
||||
response = await executor._execute_sync_plan(
|
||||
prepared_plan=prepared_plan,
|
||||
provider=SimpleNamespace(name="provider"),
|
||||
model="gpt-4.1",
|
||||
)
|
||||
|
||||
assert response == {"converted": True, "source_id": "gemini-1"}
|
||||
assert executor._ctx.provider_response_json == {"provider_id": "gemini-1"}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_execute_sync_plan_aggregates_upstream_stream_after_rust(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
executor = _make_executor()
|
||||
prepared_plan = _make_upstream_stream_prepared_plan()
|
||||
|
||||
monkeypatch.setattr(chat_sync_mod.config, "executor_backend", "rust")
|
||||
|
||||
class _FakeRegistry:
|
||||
def get_normalizer(self, format_id: str) -> _FakeNormalizer:
|
||||
assert format_id == "openai:chat"
|
||||
return _FakeNormalizer()
|
||||
|
||||
captured_chunks: list[bytes] = []
|
||||
|
||||
async def _fake_aggregate(
|
||||
byte_iter: object,
|
||||
*,
|
||||
provider_api_format: str,
|
||||
provider_name: str,
|
||||
model: str,
|
||||
request_id: str,
|
||||
envelope: object = None,
|
||||
provider_parser: object = None,
|
||||
) -> object:
|
||||
async for chunk in byte_iter: # type: ignore[attr-defined]
|
||||
captured_chunks.append(chunk)
|
||||
assert provider_api_format == "openai:chat"
|
||||
assert provider_name == "provider"
|
||||
assert model == "gpt-4.1"
|
||||
assert request_id == "req-test"
|
||||
return SimpleNamespace(id="agg-1")
|
||||
|
||||
async def _fake_execute_sync_json(self: object, plan: ExecutionPlan) -> RustExecutorSyncResult:
|
||||
assert plan.stream is True
|
||||
return RustExecutorSyncResult(
|
||||
status_code=200,
|
||||
response_body_bytes=b"data: {\"id\":\"chunk-1\"}\n\ndata: [DONE]\n\n",
|
||||
headers={"content-type": "text/event-stream"},
|
||||
)
|
||||
|
||||
async def _should_not_fallback(**kwargs: object) -> dict[str, object]:
|
||||
raise AssertionError("local execution should not be used")
|
||||
|
||||
monkeypatch.setattr(
|
||||
chat_sync_mod.RustExecutorClient,
|
||||
"execute_sync_json",
|
||||
_fake_execute_sync_json,
|
||||
)
|
||||
monkeypatch.setattr(chat_sync_mod, "get_format_converter_registry", lambda: _FakeRegistry())
|
||||
monkeypatch.setattr(
|
||||
"src.api.handlers.base.upstream_stream_bridge.aggregate_upstream_stream_to_internal_response",
|
||||
_fake_aggregate,
|
||||
)
|
||||
monkeypatch.setattr(executor, "_execute_sync_plan_locally", _should_not_fallback)
|
||||
|
||||
response = await executor._execute_sync_plan(
|
||||
prepared_plan=prepared_plan,
|
||||
provider=SimpleNamespace(name="provider"),
|
||||
model="gpt-4.1",
|
||||
)
|
||||
|
||||
assert response == {
|
||||
"aggregated": True,
|
||||
"requested_model": "gpt-4.1",
|
||||
"internal_id": "agg-1",
|
||||
}
|
||||
assert captured_chunks == [b"data: {\"id\":\"chunk-1\"}\n\ndata: [DONE]\n\n"]
|
||||
assert executor._ctx.status_code == 200
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_execute_sync_plan_turns_rust_http_error_into_httpx_status_error(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
executor = _make_executor()
|
||||
prepared_plan = _make_prepared_plan()
|
||||
|
||||
monkeypatch.setattr(chat_sync_mod.config, "executor_backend", "rust")
|
||||
|
||||
async def _fake_execute_sync_json(self: object, plan: ExecutionPlan) -> RustExecutorSyncResult:
|
||||
assert plan.url.endswith("/chat/completions")
|
||||
return RustExecutorSyncResult(
|
||||
status_code=429,
|
||||
response_json={"error": {"message": "slow down"}},
|
||||
headers={"content-type": "application/json"},
|
||||
)
|
||||
|
||||
async def _should_not_fallback(**kwargs: object) -> dict[str, object]:
|
||||
raise AssertionError("local execution should not be used")
|
||||
|
||||
monkeypatch.setattr(
|
||||
chat_sync_mod.RustExecutorClient,
|
||||
"execute_sync_json",
|
||||
_fake_execute_sync_json,
|
||||
)
|
||||
monkeypatch.setattr(executor, "_execute_sync_plan_locally", _should_not_fallback)
|
||||
|
||||
with pytest.raises(httpx.HTTPStatusError) as exc_info:
|
||||
await executor._execute_sync_plan(
|
||||
prepared_plan=prepared_plan,
|
||||
provider=SimpleNamespace(name="provider"),
|
||||
model="gpt-4.1",
|
||||
)
|
||||
|
||||
assert exc_info.value.response.status_code == 429
|
||||
assert '"message": "slow down"' in exc_info.value.upstream_response # type: ignore[attr-defined]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_execute_sync_plan_preserves_embedded_error_semantics_from_rust(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
executor = _make_executor()
|
||||
prepared_plan = _make_prepared_plan()
|
||||
|
||||
monkeypatch.setattr(chat_sync_mod.config, "executor_backend", "rust")
|
||||
|
||||
async def _fake_execute_sync_json(self: object, plan: ExecutionPlan) -> RustExecutorSyncResult:
|
||||
assert plan.provider_api_format == "openai:chat"
|
||||
return RustExecutorSyncResult(
|
||||
status_code=200,
|
||||
response_json={
|
||||
"error": {
|
||||
"message": "bad request",
|
||||
"type": "invalid_request_error",
|
||||
"code": 400,
|
||||
}
|
||||
},
|
||||
headers={"content-type": "application/json"},
|
||||
)
|
||||
|
||||
monkeypatch.setattr(
|
||||
chat_sync_mod.RustExecutorClient,
|
||||
"execute_sync_json",
|
||||
_fake_execute_sync_json,
|
||||
)
|
||||
|
||||
with pytest.raises(EmbeddedErrorException) as exc_info:
|
||||
await executor._execute_sync_plan(
|
||||
prepared_plan=prepared_plan,
|
||||
provider=SimpleNamespace(name="provider"),
|
||||
model="gpt-4.1",
|
||||
)
|
||||
|
||||
assert exc_info.value.error_message == "bad request"
|
||||
assert exc_info.value.error_code == 400
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_execute_sync_plan_falls_back_to_local_when_rust_unavailable(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
executor = _make_executor()
|
||||
prepared_plan = _make_prepared_plan()
|
||||
|
||||
monkeypatch.setattr(chat_sync_mod.config, "executor_backend", "rust")
|
||||
|
||||
async def _fake_execute_sync_json(self: object, plan: ExecutionPlan) -> RustExecutorSyncResult:
|
||||
assert plan.request_id == "req-test"
|
||||
raise RustExecutorClientError("executor down")
|
||||
|
||||
fallback_called = False
|
||||
|
||||
async def _fake_local_execute(**kwargs: object) -> dict[str, object]:
|
||||
nonlocal fallback_called
|
||||
fallback_called = True
|
||||
return {"id": "local-fallback"}
|
||||
|
||||
monkeypatch.setattr(
|
||||
chat_sync_mod.RustExecutorClient,
|
||||
"execute_sync_json",
|
||||
_fake_execute_sync_json,
|
||||
)
|
||||
monkeypatch.setattr(executor, "_execute_sync_plan_locally", _fake_local_execute)
|
||||
|
||||
response = await executor._execute_sync_plan(
|
||||
prepared_plan=prepared_plan,
|
||||
provider=SimpleNamespace(name="provider"),
|
||||
model="gpt-4.1",
|
||||
)
|
||||
|
||||
assert fallback_called is True
|
||||
assert response == {"id": "local-fallback"}
|
||||
553
tests/api/handlers/base/test_chat_rust_stream.py
Normal file
553
tests/api/handlers/base/test_chat_rust_stream.py
Normal file
@@ -0,0 +1,553 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import AsyncGenerator
|
||||
from types import SimpleNamespace
|
||||
from typing import Any
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
|
||||
import src.api.handlers.base.chat_handler_base as chatmod
|
||||
import src.services.proxy_node.resolver as proxymod
|
||||
from src.api.handlers.base.chat_handler_base import ChatHandlerBase
|
||||
from src.api.handlers.base.stream_context import StreamContext
|
||||
from src.services.request.rust_executor_client import (
|
||||
RustExecutorClientError,
|
||||
RustExecutorStreamResult,
|
||||
)
|
||||
|
||||
|
||||
class _DummyAuthInfo:
|
||||
auth_header = "authorization"
|
||||
auth_value = "Bearer test"
|
||||
decrypted_auth_config = None
|
||||
|
||||
def as_tuple(self) -> tuple[str, str]:
|
||||
return self.auth_header, self.auth_value
|
||||
|
||||
|
||||
class _PassBuilder:
|
||||
def build(self, request_body: dict[str, Any], *args: Any, **kwargs: Any) -> Any:
|
||||
return request_body, {"content-type": "application/json"}
|
||||
|
||||
|
||||
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 _FakeStreamProcessor:
|
||||
def __init__(self) -> None:
|
||||
self.prefetched_chunks: list[bytes] | None = None
|
||||
self.response_ctx: _DummyStreamResponseCtx | None = None
|
||||
|
||||
async def prefetch_and_check_error(
|
||||
self,
|
||||
byte_iterator: Any,
|
||||
provider: Any,
|
||||
endpoint: Any,
|
||||
ctx: Any,
|
||||
max_prefetch_lines: int = 5,
|
||||
max_prefetch_bytes: int = 65536,
|
||||
) -> list[bytes]:
|
||||
del provider, endpoint, ctx, max_prefetch_lines, max_prefetch_bytes
|
||||
first = await anext(byte_iterator)
|
||||
self.prefetched_chunks = [first]
|
||||
return self.prefetched_chunks
|
||||
|
||||
async def create_response_stream(
|
||||
self,
|
||||
ctx: Any,
|
||||
byte_iterator: Any,
|
||||
response_ctx: _DummyStreamResponseCtx,
|
||||
prefetched_chunks: list[bytes] | None = None,
|
||||
*,
|
||||
start_time: float | None = None,
|
||||
) -> AsyncGenerator[bytes]:
|
||||
del ctx, start_time
|
||||
self.response_ctx = response_ctx
|
||||
try:
|
||||
for chunk in prefetched_chunks or []:
|
||||
yield chunk
|
||||
async for chunk in byte_iterator:
|
||||
yield chunk
|
||||
finally:
|
||||
await response_ctx.__aexit__(None, None, None)
|
||||
|
||||
|
||||
class _DummyChatHandler(ChatHandlerBase):
|
||||
FORMAT_ID = "openai:chat"
|
||||
|
||||
def __init__(self) -> None:
|
||||
self.request_id = "req-test"
|
||||
self.api_key = SimpleNamespace(id="user-key-1")
|
||||
self._request_builder = _PassBuilder()
|
||||
self.allowed_api_formats = ["openai:chat"]
|
||||
self.api_family = None
|
||||
self.endpoint_kind = None
|
||||
self.start_time = 0.0
|
||||
|
||||
async def _convert_request(self, request: Any) -> Any:
|
||||
return request
|
||||
|
||||
def _extract_usage(self, response: dict) -> dict[str, int]:
|
||||
return {}
|
||||
|
||||
async def _get_mapped_model(
|
||||
self,
|
||||
source_model: str,
|
||||
provider_id: str,
|
||||
api_format: str | None = None,
|
||||
) -> str | None:
|
||||
del source_model, provider_id, api_format
|
||||
return None
|
||||
|
||||
def apply_mapped_model(self, request_body: dict[str, Any], mapped_model: str) -> dict[str, Any]:
|
||||
out = dict(request_body)
|
||||
out["model"] = mapped_model
|
||||
return out
|
||||
|
||||
def prepare_provider_request_body(self, request_body: dict[str, Any]) -> dict[str, Any]:
|
||||
return dict(request_body)
|
||||
|
||||
def finalize_provider_request(
|
||||
self,
|
||||
request_body: dict[str, Any],
|
||||
*,
|
||||
mapped_model: str | None,
|
||||
provider_api_format: str | None,
|
||||
) -> dict[str, Any]:
|
||||
del mapped_model, provider_api_format
|
||||
return dict(request_body)
|
||||
|
||||
def get_model_for_url(
|
||||
self,
|
||||
request_body: dict[str, Any],
|
||||
mapped_model: str | None,
|
||||
) -> str | None:
|
||||
return mapped_model or str(request_body.get("model") or "")
|
||||
|
||||
|
||||
def _patch_stream_setup(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
*,
|
||||
proxy_info: dict[str, Any] | None = None,
|
||||
delegate_config: dict[str, Any] | None = None,
|
||||
) -> None:
|
||||
async def _fake_get_provider_auth(endpoint: Any, key: Any) -> _DummyAuthInfo:
|
||||
del endpoint, key
|
||||
return _DummyAuthInfo()
|
||||
|
||||
async def _fake_resolve_proxy_info(proxy_config: Any) -> Any:
|
||||
del proxy_config
|
||||
return proxy_info
|
||||
|
||||
async def _fake_resolve_delegate(proxy_config: Any) -> Any:
|
||||
del proxy_config
|
||||
return delegate_config
|
||||
|
||||
async def _fake_get_system_proxy() -> None:
|
||||
return None
|
||||
|
||||
monkeypatch.setattr(chatmod, "get_provider_auth", _fake_get_provider_auth)
|
||||
monkeypatch.setattr(
|
||||
chatmod,
|
||||
"get_provider_behavior",
|
||||
lambda **kwargs: SimpleNamespace(
|
||||
envelope=None,
|
||||
same_format_variant=None,
|
||||
cross_format_variant=None,
|
||||
),
|
||||
)
|
||||
monkeypatch.setattr(chatmod, "build_provider_url", lambda *args, **kwargs: "https://upstream.test/v1/chat/completions")
|
||||
monkeypatch.setattr(chatmod, "get_upstream_stream_policy", lambda *args, **kwargs: None)
|
||||
monkeypatch.setattr(
|
||||
chatmod,
|
||||
"resolve_upstream_is_stream",
|
||||
lambda *, client_is_stream, policy: client_is_stream,
|
||||
)
|
||||
monkeypatch.setattr(chatmod, "enforce_stream_mode_for_upstream", lambda *args, **kwargs: None)
|
||||
monkeypatch.setattr(
|
||||
chatmod,
|
||||
"maybe_patch_request_with_prompt_cache_key",
|
||||
lambda request_body, **kwargs: request_body,
|
||||
)
|
||||
monkeypatch.setattr(proxymod, "resolve_effective_proxy", lambda provider_proxy, key_proxy=None: None)
|
||||
monkeypatch.setattr(proxymod, "resolve_proxy_info_async", _fake_resolve_proxy_info)
|
||||
monkeypatch.setattr(proxymod, "get_proxy_label", lambda proxy_info: "direct")
|
||||
monkeypatch.setattr(proxymod, "resolve_delegate_config_async", _fake_resolve_delegate)
|
||||
monkeypatch.setattr(proxymod, "get_system_proxy_config_async", _fake_get_system_proxy)
|
||||
monkeypatch.setattr(proxymod, "build_proxy_url_async", _fake_get_system_proxy)
|
||||
|
||||
|
||||
async def _iter_chunks(chunks: list[bytes]) -> AsyncGenerator[bytes]:
|
||||
for chunk in chunks:
|
||||
yield chunk
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_execute_stream_request_uses_rust_executor_when_available(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
_patch_stream_setup(monkeypatch)
|
||||
monkeypatch.setattr(chatmod.config, "executor_backend", "rust")
|
||||
|
||||
handler = _DummyChatHandler()
|
||||
stream_processor = _FakeStreamProcessor()
|
||||
ctx = StreamContext(model="gpt-test", api_format="openai:chat")
|
||||
ctx.client_api_format = "openai:chat"
|
||||
|
||||
provider = SimpleNamespace(name="provider", id="provider-1", provider_type="", proxy=None)
|
||||
endpoint = SimpleNamespace(id="endpoint-1", api_format="openai:chat", base_url="https://x")
|
||||
key = SimpleNamespace(id="key-1", proxy=None)
|
||||
candidate = SimpleNamespace(
|
||||
request_candidate_id="cand-1",
|
||||
mapping_matched_model=None,
|
||||
needs_conversion=False,
|
||||
output_limit=None,
|
||||
)
|
||||
|
||||
dummy_ctx = _DummyStreamResponseCtx()
|
||||
|
||||
async def _fake_execute_stream(self: object, plan: object) -> RustExecutorStreamResult:
|
||||
assert getattr(plan, "stream") is True
|
||||
return RustExecutorStreamResult(
|
||||
status_code=200,
|
||||
headers={"content-type": "text/event-stream", "x-upstream-test": "true"},
|
||||
byte_iterator=_iter_chunks(
|
||||
[
|
||||
b"data: {\"id\":\"chunk-1\"}\n\n",
|
||||
b"data: [DONE]\n\n",
|
||||
]
|
||||
),
|
||||
response_ctx=dummy_ctx,
|
||||
)
|
||||
|
||||
monkeypatch.setattr(chatmod.RustExecutorClient, "execute_stream", _fake_execute_stream)
|
||||
|
||||
stream = await handler._execute_stream_request(
|
||||
ctx,
|
||||
stream_processor,
|
||||
provider,
|
||||
endpoint,
|
||||
key,
|
||||
{"model": "gpt-test", "messages": [{"role": "user", "content": "hello"}]},
|
||||
{},
|
||||
candidate=candidate,
|
||||
)
|
||||
|
||||
received = [chunk async for chunk in stream]
|
||||
|
||||
assert received == [
|
||||
b"data: {\"id\":\"chunk-1\"}\n\n",
|
||||
b"data: [DONE]\n\n",
|
||||
]
|
||||
assert ctx.status_code == 200
|
||||
assert ctx.response_headers["x-upstream-test"] == "true"
|
||||
assert stream_processor.prefetched_chunks == [b"data: {\"id\":\"chunk-1\"}\n\n"]
|
||||
assert dummy_ctx.closed is True
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_execute_stream_request_accepts_async_generator_stream_processor(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
_patch_stream_setup(monkeypatch)
|
||||
monkeypatch.setattr(chatmod.config, "executor_backend", "rust")
|
||||
|
||||
handler = _DummyChatHandler()
|
||||
stream_processor = _FakeStreamProcessor()
|
||||
ctx = StreamContext(model="gpt-test", api_format="openai:chat")
|
||||
ctx.client_api_format = "openai:chat"
|
||||
|
||||
provider = SimpleNamespace(name="provider", id="provider-1", provider_type="", proxy=None)
|
||||
endpoint = SimpleNamespace(id="endpoint-1", api_format="openai:chat", base_url="https://x")
|
||||
key = SimpleNamespace(id="key-1", proxy=None)
|
||||
candidate = SimpleNamespace(
|
||||
request_candidate_id="cand-1",
|
||||
mapping_matched_model=None,
|
||||
needs_conversion=False,
|
||||
output_limit=None,
|
||||
)
|
||||
|
||||
dummy_ctx = _DummyStreamResponseCtx()
|
||||
|
||||
async def _fake_execute_stream(self: object, plan: object) -> RustExecutorStreamResult:
|
||||
assert getattr(plan, "stream") is True
|
||||
return RustExecutorStreamResult(
|
||||
status_code=200,
|
||||
headers={"content-type": "text/event-stream"},
|
||||
byte_iterator=_iter_chunks(
|
||||
[
|
||||
b"data: {\"id\":\"chunk-1\"}\n\n",
|
||||
b"data: [DONE]\n\n",
|
||||
]
|
||||
),
|
||||
response_ctx=dummy_ctx,
|
||||
)
|
||||
|
||||
monkeypatch.setattr(chatmod.RustExecutorClient, "execute_stream", _fake_execute_stream)
|
||||
|
||||
stream = await handler._execute_stream_request(
|
||||
ctx,
|
||||
stream_processor,
|
||||
provider,
|
||||
endpoint,
|
||||
key,
|
||||
{"model": "gpt-test", "messages": [{"role": "user", "content": "hello"}]},
|
||||
{},
|
||||
candidate=candidate,
|
||||
)
|
||||
|
||||
received = [chunk async for chunk in stream]
|
||||
|
||||
assert received == [
|
||||
b"data: {\"id\":\"chunk-1\"}\n\n",
|
||||
b"data: [DONE]\n\n",
|
||||
]
|
||||
assert dummy_ctx.closed is True
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_execute_stream_request_allows_tunnel_delegate_for_rust(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
_patch_stream_setup(
|
||||
monkeypatch,
|
||||
proxy_info={"node_id": "node-1", "node_name": "relay-node", "mode": "tunnel"},
|
||||
delegate_config={"tunnel": True, "node_id": "node-1"},
|
||||
)
|
||||
monkeypatch.setattr(chatmod.config, "executor_backend", "rust")
|
||||
|
||||
handler = _DummyChatHandler()
|
||||
stream_processor = _FakeStreamProcessor()
|
||||
ctx = StreamContext(model="gpt-test", api_format="openai:chat")
|
||||
ctx.client_api_format = "openai:chat"
|
||||
|
||||
provider = SimpleNamespace(
|
||||
name="provider",
|
||||
id="provider-1",
|
||||
provider_type="",
|
||||
proxy={"enabled": True, "node_id": "node-1"},
|
||||
)
|
||||
endpoint = SimpleNamespace(id="endpoint-1", api_format="openai:chat", base_url="https://x")
|
||||
key = SimpleNamespace(id="key-1", proxy=None)
|
||||
candidate = SimpleNamespace(
|
||||
request_candidate_id="cand-1",
|
||||
mapping_matched_model=None,
|
||||
needs_conversion=False,
|
||||
output_limit=None,
|
||||
)
|
||||
|
||||
dummy_ctx = _DummyStreamResponseCtx()
|
||||
|
||||
async def _fake_execute_stream(self: object, plan: object) -> RustExecutorStreamResult:
|
||||
assert getattr(plan, "proxy") is not None
|
||||
assert getattr(plan.proxy, "mode") == "tunnel"
|
||||
assert getattr(plan.proxy, "node_id") == "node-1"
|
||||
return RustExecutorStreamResult(
|
||||
status_code=200,
|
||||
headers={"content-type": "text/event-stream"},
|
||||
byte_iterator=_iter_chunks([b"data: [DONE]\n\n"]),
|
||||
response_ctx=dummy_ctx,
|
||||
)
|
||||
|
||||
monkeypatch.setattr(chatmod.RustExecutorClient, "execute_stream", _fake_execute_stream)
|
||||
|
||||
stream = await handler._execute_stream_request(
|
||||
ctx,
|
||||
stream_processor,
|
||||
provider,
|
||||
endpoint,
|
||||
key,
|
||||
{"model": "gpt-test", "messages": [{"role": "user", "content": "hello"}]},
|
||||
{},
|
||||
candidate=candidate,
|
||||
)
|
||||
|
||||
received = [chunk async for chunk in stream]
|
||||
|
||||
assert received == [b"data: [DONE]\n\n"]
|
||||
assert dummy_ctx.closed is True
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_execute_stream_request_allows_tls_profile_for_rust(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
_patch_stream_setup(monkeypatch)
|
||||
monkeypatch.setattr(chatmod.config, "executor_backend", "rust")
|
||||
|
||||
handler = _DummyChatHandler()
|
||||
stream_processor = _FakeStreamProcessor()
|
||||
ctx = StreamContext(model="gpt-test", api_format="openai:chat")
|
||||
ctx.client_api_format = "openai:chat"
|
||||
|
||||
provider = SimpleNamespace(name="provider", id="provider-1", provider_type="", proxy=None)
|
||||
endpoint = SimpleNamespace(id="endpoint-1", api_format="openai:chat", base_url="https://x")
|
||||
key = SimpleNamespace(id="key-1", proxy=None)
|
||||
candidate = SimpleNamespace(
|
||||
request_candidate_id="cand-1",
|
||||
mapping_matched_model=None,
|
||||
needs_conversion=False,
|
||||
output_limit=None,
|
||||
)
|
||||
|
||||
dummy_ctx = _DummyStreamResponseCtx()
|
||||
|
||||
async def _fake_prepare_provider_request(self: object, **kwargs: Any) -> object:
|
||||
del self, kwargs
|
||||
return chatmod.ProviderRequestResult(
|
||||
request_body={"model": "gpt-test", "messages": [{"role": "user", "content": "hello"}]},
|
||||
url_model="gpt-test",
|
||||
mapped_model=None,
|
||||
envelope=None,
|
||||
extra_headers={},
|
||||
upstream_is_stream=True,
|
||||
needs_conversion=False,
|
||||
provider_api_format="openai:chat",
|
||||
client_api_format="openai:chat",
|
||||
auth_info=_DummyAuthInfo(),
|
||||
tls_profile="claude_code_nodejs",
|
||||
)
|
||||
|
||||
monkeypatch.setattr(
|
||||
_DummyChatHandler,
|
||||
"_prepare_provider_request",
|
||||
_fake_prepare_provider_request,
|
||||
)
|
||||
|
||||
async def _fake_execute_stream(self: object, plan: object) -> RustExecutorStreamResult:
|
||||
assert getattr(plan, "tls_profile") == "claude_code_nodejs"
|
||||
return RustExecutorStreamResult(
|
||||
status_code=200,
|
||||
headers={"content-type": "text/event-stream"},
|
||||
byte_iterator=_iter_chunks([b"data: [DONE]\n\n"]),
|
||||
response_ctx=dummy_ctx,
|
||||
)
|
||||
|
||||
monkeypatch.setattr(chatmod.RustExecutorClient, "execute_stream", _fake_execute_stream)
|
||||
|
||||
stream = await handler._execute_stream_request(
|
||||
ctx,
|
||||
stream_processor,
|
||||
provider,
|
||||
endpoint,
|
||||
key,
|
||||
{"model": "gpt-test", "messages": [{"role": "user", "content": "hello"}]},
|
||||
{},
|
||||
candidate=candidate,
|
||||
)
|
||||
|
||||
received = [chunk async for chunk in stream]
|
||||
|
||||
assert received == [b"data: [DONE]\n\n"]
|
||||
assert dummy_ctx.closed is True
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_execute_stream_request_turns_rust_upstream_error_into_http_status_error(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
_patch_stream_setup(monkeypatch)
|
||||
monkeypatch.setattr(chatmod.config, "executor_backend", "rust")
|
||||
|
||||
handler = _DummyChatHandler()
|
||||
stream_processor = _FakeStreamProcessor()
|
||||
ctx = StreamContext(model="gpt-test", api_format="openai:chat")
|
||||
ctx.client_api_format = "openai:chat"
|
||||
|
||||
provider = SimpleNamespace(name="provider", id="provider-1", provider_type="", proxy=None)
|
||||
endpoint = SimpleNamespace(id="endpoint-1", api_format="openai:chat", base_url="https://x")
|
||||
key = SimpleNamespace(id="key-1", proxy=None)
|
||||
candidate = SimpleNamespace(
|
||||
request_candidate_id="cand-1",
|
||||
mapping_matched_model=None,
|
||||
needs_conversion=False,
|
||||
output_limit=None,
|
||||
)
|
||||
|
||||
dummy_ctx = _DummyStreamResponseCtx()
|
||||
|
||||
async def _fake_execute_stream(self: object, plan: object) -> RustExecutorStreamResult:
|
||||
assert getattr(plan, "stream") is True
|
||||
return RustExecutorStreamResult(
|
||||
status_code=429,
|
||||
headers={"content-type": "application/json"},
|
||||
byte_iterator=_iter_chunks([b'{"error":{"message":"slow down"}}']),
|
||||
response_ctx=dummy_ctx,
|
||||
)
|
||||
|
||||
monkeypatch.setattr(chatmod.RustExecutorClient, "execute_stream", _fake_execute_stream)
|
||||
|
||||
with pytest.raises(httpx.HTTPStatusError) as exc_info:
|
||||
await handler._execute_stream_request(
|
||||
ctx,
|
||||
stream_processor,
|
||||
provider,
|
||||
endpoint,
|
||||
key,
|
||||
{"model": "gpt-test", "messages": [{"role": "user", "content": "hello"}]},
|
||||
{},
|
||||
candidate=candidate,
|
||||
)
|
||||
|
||||
assert exc_info.value.response.status_code == 429
|
||||
assert "slow down" in exc_info.value.upstream_response # type: ignore[attr-defined]
|
||||
assert dummy_ctx.closed is True
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_execute_stream_request_falls_back_to_python_when_rust_unavailable(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
_patch_stream_setup(monkeypatch)
|
||||
monkeypatch.setattr(chatmod.config, "executor_backend", "rust")
|
||||
|
||||
handler = _DummyChatHandler()
|
||||
ctx = StreamContext(model="gpt-test", api_format="openai:chat")
|
||||
ctx.client_api_format = "openai:chat"
|
||||
|
||||
provider = SimpleNamespace(name="provider", id="provider-1", provider_type="", proxy=None)
|
||||
endpoint = SimpleNamespace(id="endpoint-1", api_format="openai:chat", base_url="https://x")
|
||||
key = SimpleNamespace(id="key-1", proxy=None)
|
||||
candidate = SimpleNamespace(
|
||||
request_candidate_id="cand-1",
|
||||
mapping_matched_model=None,
|
||||
needs_conversion=False,
|
||||
output_limit=None,
|
||||
)
|
||||
|
||||
async def _fake_execute_stream(self: object, plan: object) -> RustExecutorStreamResult:
|
||||
del plan
|
||||
raise RustExecutorClientError("executor down")
|
||||
|
||||
class _FakeHTTPClient:
|
||||
def stream(self, **kwargs: Any) -> Any:
|
||||
raise RuntimeError("local-http-client-used")
|
||||
|
||||
async def _fake_get_upstream_client(*args: Any, **kwargs: Any) -> _FakeHTTPClient:
|
||||
return _FakeHTTPClient()
|
||||
|
||||
monkeypatch.setattr(chatmod.RustExecutorClient, "execute_stream", _fake_execute_stream)
|
||||
monkeypatch.setattr(
|
||||
"src.clients.http_client.HTTPClientPool.get_upstream_client",
|
||||
_fake_get_upstream_client,
|
||||
)
|
||||
|
||||
with pytest.raises(RuntimeError) as exc_info:
|
||||
await handler._execute_stream_request(
|
||||
ctx,
|
||||
object(),
|
||||
provider,
|
||||
endpoint,
|
||||
key,
|
||||
{"model": "gpt-test", "messages": [{"role": "user", "content": "hello"}]},
|
||||
{},
|
||||
candidate=candidate,
|
||||
)
|
||||
|
||||
assert "local-http-client-used" in str(exc_info.value)
|
||||
427
tests/api/handlers/base/test_cli_rust_executor.py
Normal file
427
tests/api/handlers/base/test_cli_rust_executor.py
Normal file
@@ -0,0 +1,427 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from collections.abc import AsyncGenerator
|
||||
from types import SimpleNamespace
|
||||
from typing import Any
|
||||
|
||||
import pytest
|
||||
|
||||
import src.api.handlers.base.cli_stream_mixin as cli_stream_mod
|
||||
import src.api.handlers.base.cli_sync_mixin as cli_sync_mod
|
||||
import src.services.proxy_node.resolver as proxymod
|
||||
import src.services.task as taskmod
|
||||
from src.api.handlers.base.cli_stream_mixin import CliStreamMixin
|
||||
from src.api.handlers.base.cli_sync_mixin import CliSyncMixin
|
||||
from src.api.handlers.base.stream_context import StreamContext
|
||||
from src.services.request.rust_executor_client import (
|
||||
RustExecutorStreamResult,
|
||||
RustExecutorSyncResult,
|
||||
)
|
||||
|
||||
|
||||
class _DummyParser:
|
||||
def extract_usage_from_response(self, response: dict[str, Any]) -> dict[str, int]:
|
||||
del response
|
||||
return {
|
||||
"input_tokens": 0,
|
||||
"output_tokens": 0,
|
||||
"cache_read_tokens": 0,
|
||||
"cache_creation_tokens": 0,
|
||||
}
|
||||
|
||||
def extract_text_content(self, response: dict[str, Any]) -> str:
|
||||
return str(response.get("id") or "")
|
||||
|
||||
|
||||
class _DummyTelemetry:
|
||||
async def record_success(self, **kwargs: Any) -> int:
|
||||
del kwargs
|
||||
return 0
|
||||
|
||||
async def record_failure(self, **kwargs: Any) -> None:
|
||||
del kwargs
|
||||
|
||||
|
||||
class _DummySyncHandler(CliSyncMixin):
|
||||
FORMAT_ID = "openai:cli"
|
||||
|
||||
def __init__(self) -> None:
|
||||
self.db = None
|
||||
self.redis = None
|
||||
self.user = SimpleNamespace(id="user-1")
|
||||
self.api_key = SimpleNamespace(id="user-key-1")
|
||||
self.request_id = "req-cli-sync"
|
||||
self.client_ip = "127.0.0.1"
|
||||
self.user_agent = "pytest"
|
||||
self.start_time = 0.0
|
||||
self.allowed_api_formats = ["openai:cli"]
|
||||
self.primary_api_format = "openai:cli"
|
||||
self.api_family = None
|
||||
self.endpoint_kind = None
|
||||
self.telemetry = _DummyTelemetry()
|
||||
self.perf_metrics = None
|
||||
self._parser = _DummyParser()
|
||||
|
||||
@property
|
||||
def parser(self) -> _DummyParser:
|
||||
return self._parser
|
||||
|
||||
def _create_pending_usage(self, **kwargs: object) -> bool:
|
||||
del kwargs
|
||||
return True
|
||||
|
||||
def _build_request_metadata(self, http_request: Any | None = None) -> dict[str, Any]:
|
||||
del http_request
|
||||
return {}
|
||||
|
||||
def _merge_scheduling_metadata(
|
||||
self,
|
||||
request_metadata: dict[str, Any] | None,
|
||||
**kwargs: Any,
|
||||
) -> dict[str, Any]:
|
||||
del kwargs
|
||||
return dict(request_metadata or {})
|
||||
|
||||
def _resolve_capability_requirements(
|
||||
self,
|
||||
model_name: str,
|
||||
request_headers: dict[str, str] | None = None,
|
||||
request_body: dict[str, Any] | None = None,
|
||||
) -> dict[str, bool]:
|
||||
del model_name, request_headers, request_body
|
||||
return {}
|
||||
|
||||
async def _resolve_preferred_key_ids(
|
||||
self,
|
||||
model_name: str,
|
||||
request_body: dict[str, Any] | None = None,
|
||||
) -> list[str] | None:
|
||||
del model_name, request_body
|
||||
return None
|
||||
|
||||
def extract_model_from_request(
|
||||
self,
|
||||
request_body: dict[str, Any],
|
||||
path_params: dict[str, Any] | None = None,
|
||||
) -> str:
|
||||
del path_params
|
||||
return str(request_body.get("model") or "unknown")
|
||||
|
||||
async def _get_mapped_model(self, source_model: str, provider_id: str) -> str | None:
|
||||
del source_model, provider_id
|
||||
return None
|
||||
|
||||
async def _build_upstream_request(self, **kwargs: Any) -> Any:
|
||||
payload = dict(kwargs["request_body"])
|
||||
return SimpleNamespace(
|
||||
payload=payload,
|
||||
headers={"content-type": "application/json"},
|
||||
url="https://upstream.test/v1/responses",
|
||||
url_model=str(payload.get("model") or ""),
|
||||
envelope=None,
|
||||
upstream_is_stream=False,
|
||||
tls_profile=None,
|
||||
selected_base_url=None,
|
||||
)
|
||||
|
||||
def _extract_response_metadata(self, response_json: dict[str, Any]) -> dict[str, Any]:
|
||||
return {"id": response_json.get("id")}
|
||||
|
||||
|
||||
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 _DummyCliStreamHandler(CliStreamMixin):
|
||||
FORMAT_ID = "openai:cli"
|
||||
|
||||
def __init__(self, *, upstream_is_stream: bool) -> None:
|
||||
self.request_id = "req-cli-stream"
|
||||
self.api_key = SimpleNamespace(id="user-key-1")
|
||||
self._upstream_is_stream = upstream_is_stream
|
||||
|
||||
async def _get_mapped_model(self, source_model: str, provider_id: str) -> str | None:
|
||||
del source_model, provider_id
|
||||
return None
|
||||
|
||||
async def _build_upstream_request(self, **kwargs: Any) -> Any:
|
||||
payload = dict(kwargs["request_body"])
|
||||
return SimpleNamespace(
|
||||
payload=payload,
|
||||
headers={"content-type": "application/json"},
|
||||
url="https://upstream.test/v1/responses",
|
||||
url_model=str(payload.get("model") or ""),
|
||||
envelope=None,
|
||||
upstream_is_stream=self._upstream_is_stream,
|
||||
tls_profile=None,
|
||||
selected_base_url=None,
|
||||
)
|
||||
|
||||
def apply_mapped_model(self, request_body: dict[str, Any], mapped_model: str) -> dict[str, Any]:
|
||||
out = dict(request_body)
|
||||
out["model"] = mapped_model
|
||||
return out
|
||||
|
||||
def _extract_response_metadata(self, response_json: dict[str, Any]) -> dict[str, Any]:
|
||||
return {"id": response_json.get("id")}
|
||||
|
||||
def _record_converted_chunks(self, ctx: Any, converted_events: Any) -> None:
|
||||
del ctx, converted_events
|
||||
|
||||
def _mark_first_output(self, ctx: Any, output_state: dict[str, Any]) -> None:
|
||||
del ctx
|
||||
output_state["first_yield"] = False
|
||||
|
||||
async def _prefetch_and_check_embedded_error(
|
||||
self,
|
||||
byte_iterator: Any,
|
||||
provider: Any,
|
||||
endpoint: Any,
|
||||
ctx: Any,
|
||||
) -> list[bytes]:
|
||||
del provider, endpoint, ctx
|
||||
first = await anext(byte_iterator)
|
||||
return [first]
|
||||
|
||||
async def _create_response_stream_with_prefetch(
|
||||
self,
|
||||
ctx: Any,
|
||||
byte_iterator: Any,
|
||||
response_ctx: _DummyStreamResponseCtx,
|
||||
prefetched_chunks: list[bytes],
|
||||
) -> AsyncGenerator[bytes]:
|
||||
del ctx
|
||||
|
||||
async def _gen() -> AsyncGenerator[bytes]:
|
||||
try:
|
||||
for chunk in prefetched_chunks:
|
||||
yield chunk
|
||||
async for chunk in byte_iterator:
|
||||
yield chunk
|
||||
finally:
|
||||
await response_ctx.__aexit__(None, None, None)
|
||||
|
||||
return _gen()
|
||||
|
||||
|
||||
async def _iter_chunks(chunks: list[bytes]) -> AsyncGenerator[bytes]:
|
||||
for chunk in chunks:
|
||||
yield chunk
|
||||
|
||||
|
||||
def _patch_proxy_resolver(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
monkeypatch.setattr(proxymod, "resolve_effective_proxy", lambda provider_proxy, key_proxy=None: None)
|
||||
monkeypatch.setattr(proxymod, "get_proxy_label", lambda proxy_info: "direct")
|
||||
|
||||
async def _fake_resolve_proxy_info(proxy_config: Any) -> Any:
|
||||
del proxy_config
|
||||
return None
|
||||
|
||||
async def _fake_resolve_delegate(proxy_config: Any) -> Any:
|
||||
del proxy_config
|
||||
return None
|
||||
|
||||
async def _fake_build_proxy_url(proxy_config: Any) -> Any:
|
||||
del proxy_config
|
||||
return None
|
||||
|
||||
monkeypatch.setattr(proxymod, "resolve_proxy_info_async", _fake_resolve_proxy_info)
|
||||
monkeypatch.setattr(proxymod, "resolve_delegate_config_async", _fake_resolve_delegate)
|
||||
monkeypatch.setattr(proxymod, "build_proxy_url_async", _fake_build_proxy_url)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_cli_process_sync_uses_rust_executor_when_available(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
handler = _DummySyncHandler()
|
||||
monkeypatch.setattr(cli_sync_mod.config, "executor_backend", "rust")
|
||||
_patch_proxy_resolver(monkeypatch)
|
||||
|
||||
class _FakeTaskService:
|
||||
def __init__(self, db: Any, redis: Any) -> None:
|
||||
del db, redis
|
||||
|
||||
async def execute(self, **kwargs: Any) -> Any:
|
||||
candidate = SimpleNamespace(
|
||||
request_candidate_id="cand-1",
|
||||
mapping_matched_model=None,
|
||||
needs_conversion=False,
|
||||
output_limit=None,
|
||||
)
|
||||
provider = SimpleNamespace(
|
||||
name="provider",
|
||||
id="provider-1",
|
||||
provider_type="",
|
||||
proxy=None,
|
||||
request_timeout=None,
|
||||
stream_first_byte_timeout=None,
|
||||
)
|
||||
endpoint = SimpleNamespace(id="endpoint-1", api_format="openai:cli")
|
||||
key = SimpleNamespace(id="key-1", api_key="sk-test", proxy=None)
|
||||
response = await kwargs["request_func"](provider, endpoint, key, candidate)
|
||||
return SimpleNamespace(
|
||||
response=response,
|
||||
provider_name="provider",
|
||||
request_candidate_id="cand-1",
|
||||
provider_id="provider-1",
|
||||
endpoint_id="endpoint-1",
|
||||
key_id="key-1",
|
||||
pool_summary=None,
|
||||
)
|
||||
|
||||
async def _fake_execute_sync_json(self: object, plan: object) -> RustExecutorSyncResult:
|
||||
assert getattr(plan, "provider_api_format") == "openai:cli"
|
||||
return RustExecutorSyncResult(
|
||||
status_code=200,
|
||||
response_json={"id": "resp-rust-cli"},
|
||||
headers={"content-type": "application/json"},
|
||||
)
|
||||
|
||||
monkeypatch.setattr(taskmod, "TaskService", _FakeTaskService)
|
||||
monkeypatch.setattr(
|
||||
cli_sync_mod.RustExecutorClient,
|
||||
"execute_sync_json",
|
||||
_fake_execute_sync_json,
|
||||
)
|
||||
|
||||
response = await handler.process_sync(
|
||||
original_request_body={"model": "gpt-4.1", "input": "hello"},
|
||||
original_headers={},
|
||||
)
|
||||
|
||||
assert response.status_code == 200
|
||||
assert json.loads(response.body) == {"id": "resp-rust-cli"}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_cli_execute_stream_request_uses_rust_sync_bridge(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
handler = _DummyCliStreamHandler(upstream_is_stream=False)
|
||||
ctx = StreamContext(model="gpt-test", api_format="openai:cli")
|
||||
ctx.client_api_format = "openai:cli"
|
||||
|
||||
monkeypatch.setattr(cli_stream_mod.config, "executor_backend", "rust")
|
||||
_patch_proxy_resolver(monkeypatch)
|
||||
|
||||
async def _fake_execute_sync_json(self: object, plan: object) -> RustExecutorSyncResult:
|
||||
assert getattr(plan, "stream") is False
|
||||
return RustExecutorSyncResult(
|
||||
status_code=200,
|
||||
response_json={"id": "sync-bridge-rust"},
|
||||
headers={"content-type": "application/json"},
|
||||
)
|
||||
|
||||
async def _fake_streamify(**kwargs: Any) -> AsyncGenerator[bytes]:
|
||||
assert kwargs["response_json"] == {"id": "sync-bridge-rust"}
|
||||
yield b"data: cli-bridge\n\n"
|
||||
|
||||
monkeypatch.setattr(
|
||||
cli_stream_mod.RustExecutorClient,
|
||||
"execute_sync_json",
|
||||
_fake_execute_sync_json,
|
||||
)
|
||||
monkeypatch.setattr(handler, "_streamify_sync_response", _fake_streamify)
|
||||
|
||||
provider = SimpleNamespace(
|
||||
name="provider",
|
||||
id="provider-1",
|
||||
provider_type="",
|
||||
proxy=None,
|
||||
request_timeout=None,
|
||||
stream_first_byte_timeout=None,
|
||||
)
|
||||
endpoint = SimpleNamespace(id="endpoint-1", api_format="openai:cli", base_url="https://x")
|
||||
key = SimpleNamespace(id="key-1", proxy=None, auth_type="", api_key="sk-test")
|
||||
candidate = SimpleNamespace(
|
||||
request_candidate_id="cand-1",
|
||||
mapping_matched_model=None,
|
||||
needs_conversion=False,
|
||||
output_limit=None,
|
||||
)
|
||||
|
||||
stream = await handler._execute_stream_request(
|
||||
ctx,
|
||||
provider,
|
||||
endpoint,
|
||||
key,
|
||||
{"model": "gpt-test", "input": "hello"},
|
||||
{},
|
||||
candidate=candidate,
|
||||
)
|
||||
chunks = [chunk async for chunk in stream]
|
||||
|
||||
assert chunks == [b"data: cli-bridge\n\n"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_cli_execute_stream_request_uses_rust_native_stream(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
handler = _DummyCliStreamHandler(upstream_is_stream=True)
|
||||
ctx = StreamContext(model="gpt-test", api_format="openai:cli")
|
||||
ctx.client_api_format = "openai:cli"
|
||||
|
||||
monkeypatch.setattr(cli_stream_mod.config, "executor_backend", "rust")
|
||||
_patch_proxy_resolver(monkeypatch)
|
||||
|
||||
async def _fake_execute_stream(self: object, plan: object) -> RustExecutorStreamResult:
|
||||
assert getattr(plan, "stream") is True
|
||||
return RustExecutorStreamResult(
|
||||
status_code=200,
|
||||
headers={"content-type": "text/event-stream", "x-upstream-test": "true"},
|
||||
byte_iterator=_iter_chunks(
|
||||
[
|
||||
b"data: {\"id\":\"chunk-1\"}\n\n",
|
||||
b"data: [DONE]\n\n",
|
||||
]
|
||||
),
|
||||
response_ctx=_DummyStreamResponseCtx(),
|
||||
)
|
||||
|
||||
monkeypatch.setattr(
|
||||
cli_stream_mod.RustExecutorClient,
|
||||
"execute_stream",
|
||||
_fake_execute_stream,
|
||||
)
|
||||
|
||||
provider = SimpleNamespace(
|
||||
name="provider",
|
||||
id="provider-1",
|
||||
provider_type="",
|
||||
proxy=None,
|
||||
request_timeout=None,
|
||||
stream_first_byte_timeout=None,
|
||||
)
|
||||
endpoint = SimpleNamespace(id="endpoint-1", api_format="openai:cli", base_url="https://x")
|
||||
key = SimpleNamespace(id="key-1", proxy=None, api_key="sk-test")
|
||||
candidate = SimpleNamespace(
|
||||
request_candidate_id="cand-1",
|
||||
mapping_matched_model=None,
|
||||
needs_conversion=False,
|
||||
output_limit=None,
|
||||
)
|
||||
|
||||
stream = await handler._execute_stream_request(
|
||||
ctx,
|
||||
provider,
|
||||
endpoint,
|
||||
key,
|
||||
{"model": "gpt-test", "input": "hello"},
|
||||
{},
|
||||
candidate=candidate,
|
||||
)
|
||||
chunks = [chunk async for chunk in stream]
|
||||
|
||||
assert chunks == [
|
||||
b"data: {\"id\":\"chunk-1\"}\n\n",
|
||||
b"data: [DONE]\n\n",
|
||||
]
|
||||
assert ctx.status_code == 200
|
||||
assert ctx.response_headers["x-upstream-test"] == "true"
|
||||
168
tests/api/handlers/base/test_endpoint_checker_rust.py
Normal file
168
tests/api/handlers/base/test_endpoint_checker_rust.py
Normal file
@@ -0,0 +1,168 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any
|
||||
from unittest.mock import AsyncMock
|
||||
|
||||
import pytest
|
||||
|
||||
from src.api.handlers.base.endpoint_checker import EndpointCheckRequest, HttpRequestExecutor
|
||||
from src.services.request.rust_executor_client import (
|
||||
RustExecutorStreamResult,
|
||||
RustExecutorSyncResult,
|
||||
)
|
||||
|
||||
|
||||
class _DummyStreamContext:
|
||||
def __init__(self) -> None:
|
||||
self.closed = False
|
||||
|
||||
async def __aexit__(self, exc_type: object, exc: object, tb: object) -> None:
|
||||
self.closed = True
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_endpoint_checker_sync_prefers_rust_executor(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
from src.api.handlers.base import endpoint_checker as mod
|
||||
from src.services.request import rust_executor_client as rust_mod
|
||||
|
||||
monkeypatch.setattr(mod.config, "executor_backend", "rust")
|
||||
|
||||
executor = HttpRequestExecutor(timeout=15.0)
|
||||
proxy_snapshot = object()
|
||||
monkeypatch.setattr(
|
||||
executor,
|
||||
"_build_rust_proxy_snapshot",
|
||||
AsyncMock(return_value=proxy_snapshot),
|
||||
)
|
||||
|
||||
captured: dict[str, Any] = {}
|
||||
|
||||
async def _fake_execute_sync_json(
|
||||
self: object,
|
||||
plan: Any,
|
||||
) -> RustExecutorSyncResult:
|
||||
captured["plan"] = plan
|
||||
return RustExecutorSyncResult(
|
||||
status_code=200,
|
||||
response_json={"id": "resp_1", "usage": {"prompt_tokens": 1, "completion_tokens": 2}},
|
||||
headers={"content-type": "application/json"},
|
||||
)
|
||||
|
||||
monkeypatch.setattr(rust_mod.RustExecutorClient, "execute_sync_json", _fake_execute_sync_json)
|
||||
|
||||
result = await executor.execute(
|
||||
EndpointCheckRequest(
|
||||
url="https://upstream.test/v1/chat/completions",
|
||||
headers={"authorization": "Bearer test"},
|
||||
json_body={"model": "gpt-test", "messages": [{"role": "user", "content": "hi"}]},
|
||||
api_format="openai:chat",
|
||||
provider_name="openai",
|
||||
model_name="gpt-test",
|
||||
api_key_id="key_1",
|
||||
provider_id="provider_1",
|
||||
)
|
||||
)
|
||||
|
||||
assert result.status_code == 200
|
||||
assert result.response_data == {
|
||||
"id": "resp_1",
|
||||
"usage": {"prompt_tokens": 1, "completion_tokens": 2},
|
||||
}
|
||||
assert captured["plan"].proxy is proxy_snapshot
|
||||
assert captured["plan"].method == "POST"
|
||||
assert captured["plan"].url == "https://upstream.test/v1/chat/completions"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_endpoint_checker_stream_prefers_rust_executor(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
from src.api.handlers.base import endpoint_checker as mod
|
||||
from src.services.request import rust_executor_client as rust_mod
|
||||
|
||||
monkeypatch.setattr(mod.config, "executor_backend", "rust")
|
||||
|
||||
executor = HttpRequestExecutor(timeout=15.0)
|
||||
monkeypatch.setattr(
|
||||
executor,
|
||||
"_build_rust_proxy_snapshot",
|
||||
AsyncMock(return_value=None),
|
||||
)
|
||||
|
||||
async def _byte_iter() -> Any:
|
||||
yield b'data: {"choices":[{"delta":{"content":"Hel'
|
||||
yield b'lo"}}]}\n\n'
|
||||
yield b'data: {"choices":[{"delta":{"content":" world"},"finish_reason":"stop"}]}\n\n'
|
||||
|
||||
stream_ctx = _DummyStreamContext()
|
||||
|
||||
async def _fake_execute_stream(self: object, plan: Any) -> RustExecutorStreamResult:
|
||||
return RustExecutorStreamResult(
|
||||
status_code=200,
|
||||
headers={"content-type": "text/event-stream"},
|
||||
byte_iterator=_byte_iter(),
|
||||
response_ctx=stream_ctx,
|
||||
)
|
||||
|
||||
monkeypatch.setattr(rust_mod.RustExecutorClient, "execute_stream", _fake_execute_stream)
|
||||
|
||||
result = await executor.execute(
|
||||
EndpointCheckRequest(
|
||||
url="https://upstream.test/v1/chat/completions",
|
||||
headers={"authorization": "Bearer test"},
|
||||
json_body={
|
||||
"model": "gpt-test",
|
||||
"messages": [{"role": "user", "content": "hi"}],
|
||||
"stream": True,
|
||||
},
|
||||
api_format="openai:chat",
|
||||
provider_name="openai",
|
||||
model_name="gpt-test",
|
||||
is_stream=True,
|
||||
)
|
||||
)
|
||||
|
||||
assert result.status_code == 200
|
||||
assert result.response_data == {
|
||||
"choices": [{"delta": {"content": " world"}, "finish_reason": "stop"}]
|
||||
}
|
||||
assert stream_ctx.closed is True
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_endpoint_checker_proxy_snapshot_falls_back_to_system_proxy(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
from src.services.proxy_node import resolver as resolver_mod
|
||||
|
||||
executor = HttpRequestExecutor()
|
||||
|
||||
monkeypatch.setattr(
|
||||
resolver_mod,
|
||||
"get_system_proxy_config_async",
|
||||
AsyncMock(return_value={"enabled": True, "url": "http://system-proxy.test:8080"}),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
resolver_mod,
|
||||
"resolve_delegate_config_async",
|
||||
AsyncMock(return_value=None),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
resolver_mod,
|
||||
"build_proxy_url_async",
|
||||
AsyncMock(return_value="http://system-proxy.test:8080"),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
resolver_mod,
|
||||
"resolve_proxy_info_async",
|
||||
AsyncMock(return_value={"mode": "http", "label": "system-proxy"}),
|
||||
)
|
||||
|
||||
snapshot = await executor._build_rust_proxy_snapshot(None)
|
||||
|
||||
assert snapshot is not None
|
||||
assert snapshot.enabled is True
|
||||
assert snapshot.url == "http://system-proxy.test:8080"
|
||||
assert snapshot.mode == "http"
|
||||
226
tests/api/handlers/gemini/test_video_rust_handler.py
Normal file
226
tests/api/handlers/gemini/test_video_rust_handler.py
Normal file
@@ -0,0 +1,226 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import AsyncGenerator
|
||||
from datetime import datetime, timezone
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import AsyncMock
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
from fastapi.responses import JSONResponse, StreamingResponse
|
||||
|
||||
import src.api.handlers.gemini.video_handler as video_mod
|
||||
import src.services.proxy_node.resolver as resolver_mod
|
||||
import src.services.request.rust_executor_client as rust_client_mod
|
||||
from src.api.handlers.gemini.video_handler import GeminiVeoHandler
|
||||
from src.core.api_format.conversion.internal_video import VideoStatus
|
||||
from src.services.request.rust_executor_client import RustExecutorStreamResult
|
||||
|
||||
|
||||
class _DummyStreamResponseCtx:
|
||||
def __init__(self) -> None:
|
||||
self.closed = False
|
||||
|
||||
async def __aexit__(self, exc_type: object, exc: object, tb: object) -> None:
|
||||
self.closed = True
|
||||
|
||||
|
||||
async def _iter_chunks(chunks: list[bytes]) -> AsyncGenerator[bytes]:
|
||||
for chunk in chunks:
|
||||
yield chunk
|
||||
|
||||
|
||||
def _make_handler() -> GeminiVeoHandler:
|
||||
return GeminiVeoHandler(
|
||||
db=SimpleNamespace(),
|
||||
user=SimpleNamespace(id="user-1"),
|
||||
api_key=SimpleNamespace(id="api-key-1"),
|
||||
request_id="req-gemini-video-test",
|
||||
client_ip="127.0.0.1",
|
||||
user_agent="pytest",
|
||||
start_time=0.0,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_handle_create_task_uses_rust_sync_helper(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
handler = _make_handler()
|
||||
monkeypatch.setattr(video_mod.UsageService, "create_pending_usage", lambda *args, **kwargs: None)
|
||||
monkeypatch.setattr(
|
||||
handler._normalizer,
|
||||
"video_request_to_internal",
|
||||
lambda body: SimpleNamespace(model=str(body.get("model") or "veo-3")),
|
||||
)
|
||||
|
||||
candidate = SimpleNamespace(provider=SimpleNamespace(name="provider-1", id="prov-1"))
|
||||
endpoint = SimpleNamespace(
|
||||
id="ep-1",
|
||||
provider_id="prov-1",
|
||||
api_family="gemini",
|
||||
endpoint_kind="video",
|
||||
base_url="https://generativelanguage.googleapis.com",
|
||||
body_rules=None,
|
||||
)
|
||||
provider_key = SimpleNamespace(id="key-1")
|
||||
|
||||
monkeypatch.setattr(
|
||||
handler,
|
||||
"_resolve_upstream_key",
|
||||
AsyncMock(return_value=("upstream-key", endpoint, provider_key, None)),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
handler,
|
||||
"_build_upstream_url",
|
||||
lambda base_url, model: f"https://generativelanguage.googleapis.com/v1beta/models/{model}:predictLongRunning",
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
handler,
|
||||
"_build_upstream_headers",
|
||||
lambda original_headers, upstream_key, endpoint, auth_info, **kwargs: {
|
||||
"x-goog-api-key": upstream_key
|
||||
},
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
video_mod.HTTPClientPool,
|
||||
"get_default_client_async",
|
||||
AsyncMock(side_effect=AssertionError("python fallback should not run")),
|
||||
)
|
||||
|
||||
async def _fake_rust_sync(**kwargs: object) -> httpx.Response:
|
||||
assert kwargs["method"] == "POST"
|
||||
assert kwargs["provider_id"] == "prov-1"
|
||||
assert kwargs["endpoint_id"] == "ep-1"
|
||||
assert kwargs["key_id"] == "key-1"
|
||||
assert kwargs["body"] == {"model": "veo-3", "prompt": "hello"}
|
||||
return httpx.Response(
|
||||
200,
|
||||
request=httpx.Request("POST", str(kwargs["url"])),
|
||||
json={"name": "operations/ext-1"},
|
||||
)
|
||||
|
||||
monkeypatch.setattr(handler, "_try_rust_sync_http_response", _fake_rust_sync)
|
||||
|
||||
async def _fake_submit_with_failover(**kwargs: object) -> JSONResponse:
|
||||
response = await kwargs["submit_func"](candidate)
|
||||
assert response.status_code == 200
|
||||
assert response.json()["name"] == "operations/ext-1"
|
||||
return JSONResponse(status_code=400, content={"error": {"message": "stop"}})
|
||||
|
||||
monkeypatch.setattr(handler, "_submit_with_failover", _fake_submit_with_failover)
|
||||
|
||||
response = await handler.handle_create_task(
|
||||
http_request=SimpleNamespace(
|
||||
headers={},
|
||||
url=SimpleNamespace(scheme="https", netloc="example.com"),
|
||||
),
|
||||
original_headers={},
|
||||
original_request_body={"model": "veo-3", "prompt": "hello"},
|
||||
)
|
||||
|
||||
assert response.status_code == 400
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_handle_download_content_uses_rust_executor_with_proxy_snapshot(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
monkeypatch.setattr(video_mod.config, "executor_backend", "rust")
|
||||
handler = _make_handler()
|
||||
dummy_ctx = _DummyStreamResponseCtx()
|
||||
|
||||
monkeypatch.setattr(
|
||||
handler,
|
||||
"_get_task_by_external_id",
|
||||
lambda task_id: SimpleNamespace(
|
||||
id=task_id,
|
||||
status=VideoStatus.COMPLETED.value,
|
||||
video_url="https://storage.example.com/video.mp4",
|
||||
video_expires_at=datetime.now(timezone.utc).replace(year=2099),
|
||||
model="veo-3",
|
||||
),
|
||||
)
|
||||
|
||||
endpoint = SimpleNamespace(
|
||||
id="ep-1",
|
||||
provider_id="prov-1",
|
||||
api_family="gemini",
|
||||
endpoint_kind="video",
|
||||
base_url="https://generativelanguage.googleapis.com",
|
||||
proxy={"enabled": True, "url": "http://proxy.local:8080"},
|
||||
)
|
||||
key = SimpleNamespace(id="key-1", api_key="encrypted", proxy=None)
|
||||
monkeypatch.setattr(handler, "_get_endpoint_and_key", lambda task: (endpoint, key))
|
||||
monkeypatch.setattr(video_mod.crypto_service, "decrypt", lambda _: "upstream-key")
|
||||
monkeypatch.setattr(
|
||||
video_mod,
|
||||
"get_provider_auth",
|
||||
AsyncMock(return_value=None),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
video_mod,
|
||||
"resolve_provider_proxy",
|
||||
lambda endpoint, key: {"enabled": True, "url": "http://proxy.local:8080"},
|
||||
)
|
||||
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=None),
|
||||
)
|
||||
|
||||
async def _fake_build_proxy_url_async(proxy_config: object) -> str:
|
||||
assert proxy_config == {"enabled": True, "url": "http://proxy.local:8080"}
|
||||
return "http://proxy.local:8080"
|
||||
|
||||
async def _fake_resolve_proxy_info_async(proxy_config: object) -> dict[str, str]:
|
||||
assert proxy_config == {"enabled": True, "url": "http://proxy.local:8080"}
|
||||
return {"url": "http://proxy.local:8080"}
|
||||
|
||||
monkeypatch.setattr(resolver_mod, "build_proxy_url_async", _fake_build_proxy_url_async)
|
||||
monkeypatch.setattr(
|
||||
resolver_mod,
|
||||
"resolve_proxy_info_async",
|
||||
_fake_resolve_proxy_info_async,
|
||||
)
|
||||
|
||||
async def _fake_execute_stream(self: object, plan: object) -> RustExecutorStreamResult:
|
||||
assert getattr(plan, "method") == "GET"
|
||||
assert getattr(plan, "url") == "https://storage.example.com/video.mp4"
|
||||
assert getattr(plan, "headers") == {"x-goog-api-key": "upstream-key"}
|
||||
assert getattr(plan, "proxy").url == "http://proxy.local:8080"
|
||||
return RustExecutorStreamResult(
|
||||
status_code=200,
|
||||
headers={"content-type": "video/mp4", "x-rust-download": "true"},
|
||||
byte_iterator=_iter_chunks([b"gemini-", b"video"]),
|
||||
response_ctx=dummy_ctx,
|
||||
)
|
||||
|
||||
monkeypatch.setattr(rust_client_mod.RustExecutorClient, "execute_stream", _fake_execute_stream)
|
||||
monkeypatch.setattr(
|
||||
video_mod.HTTPClientPool,
|
||||
"get_default_client_async",
|
||||
AsyncMock(side_effect=AssertionError("python fallback should not run")),
|
||||
)
|
||||
|
||||
response = await handler.handle_download_content(
|
||||
task_id="operations/ext-1",
|
||||
http_request=SimpleNamespace(),
|
||||
original_headers={},
|
||||
query_params=None,
|
||||
)
|
||||
|
||||
assert isinstance(response, StreamingResponse)
|
||||
body = b"".join([chunk async for chunk in response.body_iterator])
|
||||
assert body == b"gemini-video"
|
||||
assert dummy_ctx.closed is True
|
||||
198
tests/api/handlers/openai/test_video_rust_download.py
Normal file
198
tests/api/handlers/openai/test_video_rust_download.py
Normal file
@@ -0,0 +1,198 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import AsyncGenerator
|
||||
from types import SimpleNamespace
|
||||
from typing import Any
|
||||
|
||||
import pytest
|
||||
from fastapi.responses import Response, StreamingResponse
|
||||
|
||||
import src.api.handlers.openai.video_handler as video_mod
|
||||
from src.api.handlers.openai.video_handler import OpenAIVideoHandler
|
||||
from src.core.api_format.conversion.internal_video import VideoStatus
|
||||
from src.services.request.rust_executor_client import (
|
||||
RustExecutorClientError,
|
||||
RustExecutorStreamResult,
|
||||
)
|
||||
|
||||
|
||||
class _DummyStreamResponseCtx:
|
||||
def __init__(self) -> None:
|
||||
self.closed = False
|
||||
|
||||
async def __aexit__(self, exc_type: object, exc: object, tb: object) -> None:
|
||||
self.closed = True
|
||||
|
||||
|
||||
async def _iter_chunks(chunks: list[bytes]) -> AsyncGenerator[bytes]:
|
||||
for chunk in chunks:
|
||||
yield chunk
|
||||
|
||||
|
||||
def _make_handler() -> OpenAIVideoHandler:
|
||||
return OpenAIVideoHandler(
|
||||
db=SimpleNamespace(),
|
||||
user=SimpleNamespace(id="user-1"),
|
||||
api_key=SimpleNamespace(id="api-key-1"),
|
||||
request_id="req-video-test",
|
||||
client_ip="127.0.0.1",
|
||||
user_agent="pytest",
|
||||
start_time=0.0,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_handle_download_content_uses_rust_executor_for_direct_video_url(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
monkeypatch.setattr(video_mod.config, "executor_backend", "rust")
|
||||
handler = _make_handler()
|
||||
dummy_ctx = _DummyStreamResponseCtx()
|
||||
|
||||
monkeypatch.setattr(
|
||||
handler,
|
||||
"_get_task",
|
||||
lambda task_id: SimpleNamespace(
|
||||
id=task_id,
|
||||
status=VideoStatus.COMPLETED.value,
|
||||
video_url="https://cdn.example.com/video.mp4",
|
||||
model="sora-2",
|
||||
),
|
||||
)
|
||||
|
||||
async def _fake_execute_stream(self: object, plan: object) -> RustExecutorStreamResult:
|
||||
assert getattr(plan, "method") == "GET"
|
||||
assert getattr(plan, "url") == "https://cdn.example.com/video.mp4"
|
||||
assert getattr(plan, "body").json_body is None
|
||||
assert getattr(plan, "body").body_bytes_b64 is None
|
||||
return RustExecutorStreamResult(
|
||||
status_code=200,
|
||||
headers={"content-type": "video/mp4", "x-rust-download": "true"},
|
||||
byte_iterator=_iter_chunks([b"video-", b"bytes"]),
|
||||
response_ctx=dummy_ctx,
|
||||
)
|
||||
|
||||
monkeypatch.setattr(video_mod.RustExecutorClient, "execute_stream", _fake_execute_stream)
|
||||
|
||||
response = await handler.handle_download_content(
|
||||
task_id="task-1",
|
||||
http_request=SimpleNamespace(),
|
||||
original_headers={},
|
||||
query_params={"variant": "video"},
|
||||
)
|
||||
|
||||
assert isinstance(response, StreamingResponse)
|
||||
assert response.headers["x-rust-download"] == "true"
|
||||
body = b"".join([chunk async for chunk in response.body_iterator])
|
||||
assert body == b"video-bytes"
|
||||
assert dummy_ctx.closed is True
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_handle_download_content_uses_rust_executor_for_upstream_content_endpoint(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
monkeypatch.setattr(video_mod.config, "executor_backend", "rust")
|
||||
handler = _make_handler()
|
||||
dummy_ctx = _DummyStreamResponseCtx()
|
||||
|
||||
monkeypatch.setattr(
|
||||
handler,
|
||||
"_get_task",
|
||||
lambda task_id: SimpleNamespace(
|
||||
id=task_id,
|
||||
status=VideoStatus.COMPLETED.value,
|
||||
video_url=None,
|
||||
external_task_id="ext-1",
|
||||
model="sora-2",
|
||||
),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
handler,
|
||||
"_get_endpoint_and_key",
|
||||
lambda task: (
|
||||
SimpleNamespace(id="ep-1", provider_id="prov-1", base_url="https://api.openai.com"),
|
||||
SimpleNamespace(id="key-1", api_key="encrypted"),
|
||||
),
|
||||
)
|
||||
monkeypatch.setattr(video_mod.crypto_service, "decrypt", lambda _: "upstream-key")
|
||||
monkeypatch.setattr(
|
||||
handler,
|
||||
"_build_upstream_url",
|
||||
lambda base_url, suffix=None: "https://api.openai.com/v1/videos/ext-1/content",
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
handler,
|
||||
"_build_upstream_headers",
|
||||
lambda original_headers, upstream_key, endpoint: {"authorization": f"Bearer {upstream_key}"},
|
||||
)
|
||||
|
||||
async def _fake_execute_stream(self: object, plan: object) -> RustExecutorStreamResult:
|
||||
assert getattr(plan, "method") == "GET"
|
||||
assert getattr(plan, "url") == "https://api.openai.com/v1/videos/ext-1/content"
|
||||
assert getattr(plan, "headers") == {"authorization": "Bearer upstream-key"}
|
||||
assert getattr(plan, "provider_id") == "prov-1"
|
||||
assert getattr(plan, "endpoint_id") == "ep-1"
|
||||
assert getattr(plan, "key_id") == "key-1"
|
||||
return RustExecutorStreamResult(
|
||||
status_code=200,
|
||||
headers={"content-type": "video/mp4", "x-rust-download": "true"},
|
||||
byte_iterator=_iter_chunks([b"upstream-", b"video"]),
|
||||
response_ctx=dummy_ctx,
|
||||
)
|
||||
|
||||
monkeypatch.setattr(video_mod.RustExecutorClient, "execute_stream", _fake_execute_stream)
|
||||
|
||||
response = await handler.handle_download_content(
|
||||
task_id="task-1",
|
||||
http_request=SimpleNamespace(),
|
||||
original_headers={},
|
||||
query_params=None,
|
||||
)
|
||||
|
||||
assert isinstance(response, StreamingResponse)
|
||||
body = b"".join([chunk async for chunk in response.body_iterator])
|
||||
assert body == b"upstream-video"
|
||||
assert dummy_ctx.closed is True
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_handle_download_content_falls_back_to_python_proxy_when_rust_unavailable(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
monkeypatch.setattr(video_mod.config, "executor_backend", "rust")
|
||||
handler = _make_handler()
|
||||
|
||||
monkeypatch.setattr(
|
||||
handler,
|
||||
"_get_task",
|
||||
lambda task_id: SimpleNamespace(
|
||||
id=task_id,
|
||||
status=VideoStatus.COMPLETED.value,
|
||||
video_url="https://cdn.example.com/video.mp4",
|
||||
model="sora-2",
|
||||
),
|
||||
)
|
||||
|
||||
async def _failing_execute_stream(self: object, plan: object) -> RustExecutorStreamResult:
|
||||
del self, plan
|
||||
raise RustExecutorClientError("executor down")
|
||||
|
||||
fallback_response = Response(content=b"python-fallback", media_type="video/mp4")
|
||||
|
||||
async def _fake_proxy_direct_url(url: str, task_id: str) -> Response:
|
||||
assert url == "https://cdn.example.com/video.mp4"
|
||||
assert task_id == "task-1"
|
||||
return fallback_response
|
||||
|
||||
monkeypatch.setattr(video_mod.RustExecutorClient, "execute_stream", _failing_execute_stream)
|
||||
monkeypatch.setattr(handler, "_proxy_direct_url", _fake_proxy_direct_url)
|
||||
|
||||
response = await handler.handle_download_content(
|
||||
task_id="task-1",
|
||||
http_request=SimpleNamespace(),
|
||||
original_headers={},
|
||||
query_params={"variant": "video"},
|
||||
)
|
||||
|
||||
assert response is fallback_response
|
||||
103
tests/api/handlers/openai/test_video_rust_sync.py
Normal file
103
tests/api/handlers/openai/test_video_rust_sync.py
Normal file
@@ -0,0 +1,103 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import AsyncMock
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
from fastapi.responses import JSONResponse
|
||||
|
||||
import src.api.handlers.openai.video_handler as video_mod
|
||||
from src.api.handlers.openai.video_handler import OpenAIVideoHandler
|
||||
|
||||
|
||||
def _make_handler() -> OpenAIVideoHandler:
|
||||
return OpenAIVideoHandler(
|
||||
db=SimpleNamespace(),
|
||||
user=SimpleNamespace(id="user-1"),
|
||||
api_key=SimpleNamespace(id="api-key-1"),
|
||||
request_id="req-video-sync-test",
|
||||
client_ip="127.0.0.1",
|
||||
user_agent="pytest",
|
||||
start_time=0.0,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_handle_create_task_uses_rust_sync_helper(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
handler = _make_handler()
|
||||
monkeypatch.setattr(video_mod.UsageService, "create_pending_usage", lambda *args, **kwargs: None)
|
||||
monkeypatch.setattr(
|
||||
handler._normalizer,
|
||||
"video_request_to_internal",
|
||||
lambda body: SimpleNamespace(model=str(body.get("model") or "sora-2")),
|
||||
)
|
||||
|
||||
candidate = SimpleNamespace(provider=SimpleNamespace(name="provider-1", id="prov-1"))
|
||||
endpoint = SimpleNamespace(
|
||||
id="ep-1",
|
||||
api_family="openai",
|
||||
endpoint_kind="video",
|
||||
base_url="https://api.openai.com",
|
||||
body_rules=None,
|
||||
)
|
||||
provider_key = SimpleNamespace(id="key-1")
|
||||
|
||||
monkeypatch.setattr(
|
||||
handler,
|
||||
"_resolve_upstream_key",
|
||||
AsyncMock(return_value=("upstream-key", endpoint, provider_key)),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
handler,
|
||||
"_build_upstream_url",
|
||||
lambda base_url: "https://api.openai.com/v1/videos",
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
handler,
|
||||
"_build_upstream_headers",
|
||||
lambda original_headers, upstream_key, endpoint, **kwargs: {
|
||||
"authorization": f"Bearer {upstream_key}"
|
||||
},
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
video_mod.HTTPClientPool,
|
||||
"get_default_client_async",
|
||||
AsyncMock(side_effect=AssertionError("python fallback should not run")),
|
||||
)
|
||||
|
||||
async def _fake_rust_sync(**kwargs: object) -> httpx.Response:
|
||||
assert kwargs["method"] == "POST"
|
||||
assert kwargs["url"] == "https://api.openai.com/v1/videos"
|
||||
assert kwargs["provider_id"] == "prov-1"
|
||||
assert kwargs["endpoint_id"] == "ep-1"
|
||||
assert kwargs["key_id"] == "key-1"
|
||||
assert kwargs["body"] == {"model": "sora-2", "prompt": "hello"}
|
||||
return httpx.Response(
|
||||
200,
|
||||
request=httpx.Request("POST", str(kwargs["url"])),
|
||||
json={"id": "ext-1"},
|
||||
)
|
||||
|
||||
create_failed = AsyncMock()
|
||||
monkeypatch.setattr(handler, "_try_rust_sync_http_response", _fake_rust_sync)
|
||||
monkeypatch.setattr(handler, "_create_failed_task_and_usage", create_failed)
|
||||
|
||||
async def _fake_submit_with_failover(**kwargs: object) -> JSONResponse:
|
||||
response = await kwargs["submit_func"](candidate)
|
||||
assert response.status_code == 200
|
||||
assert response.json()["id"] == "ext-1"
|
||||
return JSONResponse(status_code=400, content={"error": {"message": "stop"}})
|
||||
|
||||
monkeypatch.setattr(handler, "_submit_with_failover", _fake_submit_with_failover)
|
||||
|
||||
response = await handler.handle_create_task(
|
||||
http_request=SimpleNamespace(),
|
||||
original_headers={},
|
||||
original_request_body={"model": "sora-2", "prompt": "hello"},
|
||||
)
|
||||
|
||||
assert response.status_code == 400
|
||||
create_failed.assert_awaited_once()
|
||||
296
tests/api/public/test_gemini_files_rust.py
Normal file
296
tests/api/public/test_gemini_files_rust.py
Normal file
@@ -0,0 +1,296 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from collections.abc import AsyncGenerator
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import AsyncMock
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
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.rust_executor_client as rust_client_mod
|
||||
from src.api.public.gemini_files import UpstreamContext
|
||||
from src.config.settings import config
|
||||
from src.services.request.executor_plan import ExecutionProxySnapshot
|
||||
from src.services.request.rust_executor_client import (
|
||||
RustExecutorStreamResult,
|
||||
RustExecutorSyncResult,
|
||||
)
|
||||
|
||||
|
||||
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) -> RustExecutorSyncResult:
|
||||
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 RustExecutorSyncResult(
|
||||
status_code=200,
|
||||
headers={"content-type": "application/json", "x-rust-files": "true"},
|
||||
response_json={"files": [{"name": "files/abc"}]},
|
||||
)
|
||||
|
||||
monkeypatch.setattr(
|
||||
rust_client_mod.RustExecutorClient,
|
||||
"execute_sync_json",
|
||||
_fake_execute_sync_json,
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
gemini_files_mod.HTTPClientPool,
|
||||
"get_upstream_client",
|
||||
AsyncMock(side_effect=AssertionError("python fallback should not run")),
|
||||
)
|
||||
|
||||
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_fallback_uses_upstream_client_proxy_context(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
class _FakeClient:
|
||||
async def get(self, url: str, headers: dict[str, str]) -> httpx.Response:
|
||||
assert url == "https://generativelanguage.googleapis.com/v1beta/files"
|
||||
assert headers == {"x-goog-api-key": "upstream-key"}
|
||||
return httpx.Response(
|
||||
200,
|
||||
request=httpx.Request("GET", url),
|
||||
json={"files": []},
|
||||
)
|
||||
|
||||
async def _fake_try_rust_sync_proxy_request(*args: object, **kwargs: object) -> None:
|
||||
del args, kwargs
|
||||
return None
|
||||
|
||||
async def _fake_get_upstream_client(
|
||||
delegate_cfg: dict[str, object] | None,
|
||||
*,
|
||||
proxy_config: dict[str, object] | None = None,
|
||||
tls_profile: str | None = None,
|
||||
) -> _FakeClient:
|
||||
assert delegate_cfg == {"tunnel": True, "node_id": "node-1"}
|
||||
assert proxy_config == {"enabled": True, "node_id": "node-1"}
|
||||
assert tls_profile is None
|
||||
return _FakeClient()
|
||||
|
||||
monkeypatch.setattr(
|
||||
gemini_files_mod,
|
||||
"_try_rust_sync_proxy_request",
|
||||
_fake_try_rust_sync_proxy_request,
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
gemini_files_mod.HTTPClientPool,
|
||||
"get_upstream_client",
|
||||
_fake_get_upstream_client,
|
||||
)
|
||||
|
||||
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",
|
||||
proxy_config={"enabled": True, "node_id": "node-1"},
|
||||
delegate_config={"tunnel": True, "node_id": "node-1"},
|
||||
)
|
||||
|
||||
assert response.status_code == 200
|
||||
assert json.loads(response.body) == {"files": []}
|
||||
|
||||
|
||||
@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) -> RustExecutorStreamResult:
|
||||
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 RustExecutorStreamResult(
|
||||
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.RustExecutorClient,
|
||||
"execute_stream",
|
||||
_fake_execute_stream,
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
gemini_files_mod.HTTPClientPool,
|
||||
"get_upstream_client",
|
||||
AsyncMock(side_effect=AssertionError("python fallback should not run")),
|
||||
)
|
||||
|
||||
response = await gemini_files_mod.download_file(
|
||||
"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
|
||||
142
tests/api/public/test_system_catalog_rust.py
Normal file
142
tests/api/public/test_system_catalog_rust.py
Normal file
@@ -0,0 +1,142 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
|
||||
|
||||
def _build_provider_fixture() -> SimpleNamespace:
|
||||
endpoint = SimpleNamespace(
|
||||
id="endpoint_1",
|
||||
base_url="https://upstream.test",
|
||||
api_format="openai:chat",
|
||||
is_active=True,
|
||||
)
|
||||
key = SimpleNamespace(
|
||||
id="key_1",
|
||||
is_active=True,
|
||||
api_formats=None,
|
||||
)
|
||||
return SimpleNamespace(
|
||||
id="provider_1",
|
||||
name="openai",
|
||||
endpoints=[endpoint],
|
||||
api_keys=[key],
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_test_connection_prefers_rust_executor(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
from src.api.public import system_catalog as mod
|
||||
|
||||
provider = _build_provider_fixture()
|
||||
|
||||
monkeypatch.setattr(mod, "_select_provider", lambda _db, _provider_name: provider)
|
||||
monkeypatch.setattr(mod, "build_test_request_body", lambda *_args, **_kwargs: {"model": "gpt-test"})
|
||||
monkeypatch.setattr(mod, "get_provider_auth", AsyncMock(return_value=None))
|
||||
|
||||
class _DummyBuilder:
|
||||
def build(self, *_args: object, **_kwargs: object) -> tuple[dict[str, str], dict[str, str]]:
|
||||
return {"model": "gpt-test"}, {"authorization": "Bearer test"}
|
||||
|
||||
monkeypatch.setattr(mod, "PassthroughRequestBuilder", lambda: _DummyBuilder())
|
||||
monkeypatch.setattr(
|
||||
mod,
|
||||
"build_provider_url",
|
||||
lambda *_args, **_kwargs: "https://upstream.test/v1/chat/completions",
|
||||
)
|
||||
|
||||
proxy_snapshot = object()
|
||||
monkeypatch.setattr(
|
||||
mod,
|
||||
"_build_test_connection_transport_context",
|
||||
AsyncMock(return_value=({"enabled": True}, {"node_id": "node-1", "tunnel": True}, proxy_snapshot)),
|
||||
)
|
||||
|
||||
rust_response = httpx.Response(
|
||||
200,
|
||||
request=httpx.Request("POST", "https://upstream.test/v1/chat/completions"),
|
||||
json={"id": "resp_rust"},
|
||||
)
|
||||
rust_call = AsyncMock(return_value=rust_response)
|
||||
monkeypatch.setattr(mod, "_try_rust_test_connection_response", rust_call)
|
||||
|
||||
get_upstream_client = AsyncMock(side_effect=AssertionError("python upstream client should not be used"))
|
||||
monkeypatch.setattr(mod.HTTPClientPool, "get_upstream_client", get_upstream_client)
|
||||
|
||||
result = await mod.test_connection(
|
||||
request=SimpleNamespace(query_params={}),
|
||||
db=MagicMock(),
|
||||
provider=None,
|
||||
model="gpt-test",
|
||||
api_format=None,
|
||||
)
|
||||
|
||||
assert result["status"] == "success"
|
||||
assert result["response_id"] == "resp_rust"
|
||||
rust_call.assert_awaited_once()
|
||||
get_upstream_client.assert_not_awaited()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_test_connection_fallback_uses_transport_aware_client(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
from src.api.public import system_catalog as mod
|
||||
|
||||
provider = _build_provider_fixture()
|
||||
|
||||
monkeypatch.setattr(mod, "_select_provider", lambda _db, _provider_name: provider)
|
||||
monkeypatch.setattr(mod, "build_test_request_body", lambda *_args, **_kwargs: {"model": "gpt-test"})
|
||||
monkeypatch.setattr(mod, "get_provider_auth", AsyncMock(return_value=None))
|
||||
|
||||
class _DummyBuilder:
|
||||
def build(self, *_args: object, **_kwargs: object) -> tuple[dict[str, str], dict[str, str]]:
|
||||
return {"model": "gpt-test"}, {"authorization": "Bearer test"}
|
||||
|
||||
monkeypatch.setattr(mod, "PassthroughRequestBuilder", lambda: _DummyBuilder())
|
||||
monkeypatch.setattr(
|
||||
mod,
|
||||
"build_provider_url",
|
||||
lambda *_args, **_kwargs: "https://upstream.test/v1/chat/completions",
|
||||
)
|
||||
|
||||
proxy_config = {"enabled": True, "url": "http://proxy.test:8080"}
|
||||
delegate_cfg = {"node_id": "node-1", "tunnel": True}
|
||||
monkeypatch.setattr(
|
||||
mod,
|
||||
"_build_test_connection_transport_context",
|
||||
AsyncMock(return_value=(proxy_config, delegate_cfg, None)),
|
||||
)
|
||||
monkeypatch.setattr(mod, "_try_rust_test_connection_response", AsyncMock(return_value=None))
|
||||
|
||||
upstream_response = httpx.Response(
|
||||
200,
|
||||
request=httpx.Request("POST", "https://upstream.test/v1/chat/completions"),
|
||||
json={"id": "resp_python"},
|
||||
)
|
||||
upstream_client = MagicMock()
|
||||
upstream_client.post = AsyncMock(return_value=upstream_response)
|
||||
get_upstream_client = AsyncMock(return_value=upstream_client)
|
||||
monkeypatch.setattr(mod.HTTPClientPool, "get_upstream_client", get_upstream_client)
|
||||
|
||||
result = await mod.test_connection(
|
||||
request=SimpleNamespace(query_params={}),
|
||||
db=MagicMock(),
|
||||
provider=None,
|
||||
model="gpt-test",
|
||||
api_format=None,
|
||||
)
|
||||
|
||||
assert result["status"] == "success"
|
||||
assert result["response_id"] == "resp_python"
|
||||
get_upstream_client.assert_awaited_once_with(delegate_cfg, proxy_config=proxy_config)
|
||||
upstream_client.post.assert_awaited_once_with(
|
||||
"https://upstream.test/v1/chat/completions",
|
||||
json={"model": "gpt-test"},
|
||||
headers={"authorization": "Bearer test"},
|
||||
)
|
||||
633
tests/api/test_internal_gateway_routes.py
Normal file
633
tests/api/test_internal_gateway_routes.py
Normal file
@@ -0,0 +1,633 @@
|
||||
import base64
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import AsyncMock
|
||||
|
||||
import pytest
|
||||
from fastapi import FastAPI
|
||||
from fastapi.testclient import TestClient
|
||||
from fastapi.responses import JSONResponse, StreamingResponse
|
||||
|
||||
from src.api.internal.gateway import (
|
||||
CONTROL_ACTION_HEADER,
|
||||
CONTROL_ACTION_PROXY_PUBLIC,
|
||||
CONTROL_EXECUTED_HEADER,
|
||||
GatewayResolveRequest,
|
||||
_is_streaming_sync_payload,
|
||||
_resolve_auth_context,
|
||||
_resolve_gateway_sync_adapter,
|
||||
classify_gateway_route,
|
||||
router,
|
||||
)
|
||||
from src.database import get_db
|
||||
|
||||
|
||||
def test_classify_openai_chat_route_as_ai_public() -> None:
|
||||
decision = classify_gateway_route("POST", "/v1/chat/completions")
|
||||
|
||||
assert decision.route_class == "ai_public"
|
||||
assert decision.route_family == "openai"
|
||||
assert decision.route_kind == "chat"
|
||||
assert decision.auth_endpoint_signature == "openai:chat"
|
||||
assert decision.executor_candidate is True
|
||||
assert decision.action == "proxy_public"
|
||||
|
||||
|
||||
def test_classify_gemini_files_download_route_as_ai_public() -> None:
|
||||
decision = classify_gateway_route("GET", "/v1beta/files/file-123:download")
|
||||
|
||||
assert decision.route_class == "ai_public"
|
||||
assert decision.route_family == "gemini"
|
||||
assert decision.route_kind == "files"
|
||||
assert decision.auth_endpoint_signature == "gemini:chat"
|
||||
assert decision.executor_candidate is True
|
||||
|
||||
|
||||
def test_classify_gemini_files_nested_metadata_route_as_ai_public() -> None:
|
||||
decision = classify_gateway_route("GET", "/v1beta/files/files/abc-123")
|
||||
|
||||
assert decision.route_class == "ai_public"
|
||||
assert decision.route_family == "gemini"
|
||||
assert decision.route_kind == "files"
|
||||
assert decision.executor_candidate is True
|
||||
|
||||
|
||||
def test_classify_gemini_video_operation_route_as_ai_public() -> None:
|
||||
decision = classify_gateway_route("GET", "/v1beta/models/veo-3/operations/op-123")
|
||||
|
||||
assert decision.route_class == "ai_public"
|
||||
assert decision.route_family == "gemini"
|
||||
assert decision.route_kind == "video"
|
||||
assert decision.auth_endpoint_signature == "gemini:video"
|
||||
assert decision.executor_candidate is True
|
||||
|
||||
|
||||
def test_classify_non_ai_route_as_passthrough() -> None:
|
||||
decision = classify_gateway_route("GET", "/api/admin/system/info")
|
||||
|
||||
assert decision.route_class == "passthrough"
|
||||
assert decision.route_family is None
|
||||
assert decision.route_kind is None
|
||||
assert decision.executor_candidate is False
|
||||
|
||||
|
||||
def test_classify_claude_cli_route_from_bearer_header() -> None:
|
||||
decision = classify_gateway_route(
|
||||
"POST",
|
||||
"/v1/messages",
|
||||
{"authorization": "Bearer sk-cli"},
|
||||
)
|
||||
|
||||
assert decision.route_class == "ai_public"
|
||||
assert decision.route_family == "claude"
|
||||
assert decision.route_kind == "cli"
|
||||
assert decision.auth_endpoint_signature == "claude:cli"
|
||||
|
||||
|
||||
def test_classify_gemini_cli_route_from_user_agent() -> None:
|
||||
decision = classify_gateway_route(
|
||||
"POST",
|
||||
"/v1beta/models/gemini-2.5-pro:generateContent",
|
||||
{"user-agent": "GeminiCLI/1.2.3"},
|
||||
)
|
||||
|
||||
assert decision.route_class == "ai_public"
|
||||
assert decision.route_family == "gemini"
|
||||
assert decision.route_kind == "cli"
|
||||
assert decision.auth_endpoint_signature == "gemini:cli"
|
||||
|
||||
|
||||
def test_resolve_sync_adapter_for_openai_chat_route() -> None:
|
||||
decision = classify_gateway_route("POST", "/v1/chat/completions")
|
||||
|
||||
adapter, path_params = _resolve_gateway_sync_adapter(decision, "/v1/chat/completions")
|
||||
|
||||
assert adapter is not None
|
||||
assert adapter.name == "openai.chat"
|
||||
assert path_params == {}
|
||||
|
||||
|
||||
def test_resolve_sync_adapter_for_gemini_route_extracts_model_path_params() -> None:
|
||||
decision = classify_gateway_route("POST", "/v1beta/models/gemini-2.5-pro:generateContent")
|
||||
|
||||
adapter, path_params = _resolve_gateway_sync_adapter(
|
||||
decision,
|
||||
"/v1beta/models/gemini-2.5-pro:generateContent",
|
||||
)
|
||||
|
||||
assert adapter is not None
|
||||
assert adapter.name == "gemini.chat"
|
||||
assert path_params == {"model": "gemini-2.5-pro", "stream": False}
|
||||
|
||||
|
||||
def test_resolve_sync_adapter_rejects_non_sync_files_route() -> None:
|
||||
decision = classify_gateway_route("GET", "/v1beta/files/file-123:download")
|
||||
|
||||
adapter, path_params = _resolve_gateway_sync_adapter(decision, "/v1beta/files/file-123:download")
|
||||
|
||||
assert adapter is None
|
||||
assert path_params == {}
|
||||
|
||||
|
||||
def test_is_streaming_sync_payload_detects_body_and_path_stream_flags() -> None:
|
||||
assert _is_streaming_sync_payload({"stream": True}, {}) is True
|
||||
assert _is_streaming_sync_payload({}, {"stream": True}) is True
|
||||
assert _is_streaming_sync_payload({"stream": False}, {"stream": False}) is False
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_resolve_auth_context_from_openai_bearer_header(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
payload = GatewayResolveRequest(
|
||||
method="POST",
|
||||
path="/v1/chat/completions",
|
||||
headers={"Authorization": "Bearer sk-test"},
|
||||
)
|
||||
decision = classify_gateway_route(payload.method, payload.path, payload.headers)
|
||||
monkeypatch.setattr(
|
||||
"src.api.internal.gateway.AuthService.authenticate_api_key_threadsafe",
|
||||
AsyncMock(
|
||||
return_value=SimpleNamespace(
|
||||
user=SimpleNamespace(id="user-123"),
|
||||
api_key=SimpleNamespace(id="key-123"),
|
||||
balance_remaining=42.5,
|
||||
access_allowed=True,
|
||||
)
|
||||
),
|
||||
)
|
||||
|
||||
auth_context = await _resolve_auth_context(payload, decision)
|
||||
|
||||
assert auth_context == {
|
||||
"user_id": "user-123",
|
||||
"api_key_id": "key-123",
|
||||
"balance_remaining": 42.5,
|
||||
"access_allowed": True,
|
||||
}
|
||||
|
||||
|
||||
def test_execute_sync_route_returns_controlled_response(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)
|
||||
|
||||
class FakeAdapter:
|
||||
mode = SimpleNamespace(value="standard")
|
||||
allowed_api_formats = ["openai:chat"]
|
||||
|
||||
def authorize(self, context: object) -> None:
|
||||
self.authorized_context = context
|
||||
|
||||
async def handle(self, context: object) -> JSONResponse:
|
||||
assert getattr(context, "path_params", {}) == {}
|
||||
return JSONResponse(status_code=201, content={"ok": True, "request_id": context.request_id})
|
||||
|
||||
fake_adapter = FakeAdapter()
|
||||
fake_pipeline = SimpleNamespace(_check_user_rate_limit=AsyncMock(return_value=None))
|
||||
|
||||
monkeypatch.setattr(
|
||||
"src.api.internal.gateway._resolve_gateway_sync_adapter",
|
||||
lambda decision, path: (fake_adapter, {}),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
"src.api.internal.gateway._load_gateway_auth_models",
|
||||
lambda db, auth_context: (
|
||||
SimpleNamespace(id="user-123"),
|
||||
SimpleNamespace(id="key-123"),
|
||||
),
|
||||
)
|
||||
monkeypatch.setattr("src.api.internal.gateway.get_pipeline", lambda: fake_pipeline)
|
||||
|
||||
client = TestClient(app, base_url="http://127.0.0.1")
|
||||
response = client.post(
|
||||
"/api/internal/gateway/execute-sync",
|
||||
json={
|
||||
"trace_id": "trace-sync-123",
|
||||
"method": "POST",
|
||||
"path": "/v1/chat/completions",
|
||||
"headers": {"user-agent": "pytest"},
|
||||
"body_json": {"model": "gpt-5", "messages": []},
|
||||
"auth_context": {
|
||||
"user_id": "user-123",
|
||||
"api_key_id": "key-123",
|
||||
"balance_remaining": 12.5,
|
||||
"access_allowed": True,
|
||||
},
|
||||
},
|
||||
)
|
||||
|
||||
assert response.status_code == 201
|
||||
assert response.headers[CONTROL_EXECUTED_HEADER] == "true"
|
||||
assert response.json()["ok"] is True
|
||||
assert response.json()["request_id"] == "trace-sync-123"
|
||||
|
||||
|
||||
def test_execute_sync_route_falls_back_for_stream_payload() -> None:
|
||||
app = FastAPI()
|
||||
app.include_router(router)
|
||||
app.dependency_overrides[get_db] = lambda: object()
|
||||
monkeypatch = pytest.MonkeyPatch()
|
||||
monkeypatch.setattr("src.api.internal.gateway.ensure_loopback", lambda request: None)
|
||||
|
||||
client = TestClient(app, base_url="http://127.0.0.1")
|
||||
response = client.post(
|
||||
"/api/internal/gateway/execute-sync",
|
||||
json={
|
||||
"trace_id": "trace-stream-123",
|
||||
"method": "POST",
|
||||
"path": "/v1/chat/completions",
|
||||
"headers": {"user-agent": "pytest"},
|
||||
"body_json": {"model": "gpt-5", "messages": [], "stream": True},
|
||||
"auth_context": {
|
||||
"user_id": "user-123",
|
||||
"api_key_id": "key-123",
|
||||
"access_allowed": True,
|
||||
},
|
||||
},
|
||||
)
|
||||
|
||||
assert response.status_code == 409
|
||||
assert response.headers[CONTROL_ACTION_HEADER] == CONTROL_ACTION_PROXY_PUBLIC
|
||||
assert response.json() == {"action": CONTROL_ACTION_PROXY_PUBLIC}
|
||||
monkeypatch.undo()
|
||||
|
||||
|
||||
def test_execute_stream_route_returns_controlled_stream(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)
|
||||
|
||||
class FakeAdapter:
|
||||
mode = SimpleNamespace(value="standard")
|
||||
allowed_api_formats = ["openai:chat"]
|
||||
|
||||
def authorize(self, context: object) -> None:
|
||||
self.authorized_context = context
|
||||
|
||||
async def handle(self, context: object) -> StreamingResponse:
|
||||
async def _iter() -> object:
|
||||
yield b"data: one\n\n"
|
||||
yield b"data: [DONE]\n\n"
|
||||
|
||||
return StreamingResponse(_iter(), media_type="text/event-stream")
|
||||
|
||||
fake_pipeline = SimpleNamespace(_check_user_rate_limit=AsyncMock(return_value=None))
|
||||
monkeypatch.setattr(
|
||||
"src.api.internal.gateway._resolve_gateway_sync_adapter",
|
||||
lambda decision, path: (FakeAdapter(), {}),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
"src.api.internal.gateway._load_gateway_auth_models",
|
||||
lambda db, auth_context: (
|
||||
SimpleNamespace(id="user-123"),
|
||||
SimpleNamespace(id="key-123"),
|
||||
),
|
||||
)
|
||||
monkeypatch.setattr("src.api.internal.gateway.get_pipeline", lambda: fake_pipeline)
|
||||
|
||||
client = TestClient(app, base_url="http://127.0.0.1")
|
||||
response = client.post(
|
||||
"/api/internal/gateway/execute-stream",
|
||||
json={
|
||||
"trace_id": "trace-stream-123",
|
||||
"method": "POST",
|
||||
"path": "/v1/chat/completions",
|
||||
"headers": {"user-agent": "pytest"},
|
||||
"body_json": {"model": "gpt-5", "messages": [], "stream": True},
|
||||
"auth_context": {
|
||||
"user_id": "user-123",
|
||||
"api_key_id": "key-123",
|
||||
"balance_remaining": 12.5,
|
||||
"access_allowed": True,
|
||||
},
|
||||
},
|
||||
)
|
||||
|
||||
assert response.status_code == 200
|
||||
assert response.headers[CONTROL_EXECUTED_HEADER] == "true"
|
||||
assert response.text == "data: one\n\ndata: [DONE]\n\n"
|
||||
|
||||
|
||||
def test_execute_stream_route_falls_back_for_sync_payload(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)
|
||||
|
||||
client = TestClient(app, base_url="http://127.0.0.1")
|
||||
response = client.post(
|
||||
"/api/internal/gateway/execute-stream",
|
||||
json={
|
||||
"trace_id": "trace-sync-456",
|
||||
"method": "POST",
|
||||
"path": "/v1/chat/completions",
|
||||
"headers": {"user-agent": "pytest"},
|
||||
"body_json": {"model": "gpt-5", "messages": []},
|
||||
"auth_context": {
|
||||
"user_id": "user-123",
|
||||
"api_key_id": "key-123",
|
||||
"access_allowed": True,
|
||||
},
|
||||
},
|
||||
)
|
||||
|
||||
assert response.status_code == 409
|
||||
assert response.headers[CONTROL_ACTION_HEADER] == CONTROL_ACTION_PROXY_PUBLIC
|
||||
assert response.json() == {"action": CONTROL_ACTION_PROXY_PUBLIC}
|
||||
|
||||
|
||||
def test_execute_sync_route_handles_gemini_files_list(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)
|
||||
monkeypatch.setattr(
|
||||
"src.api.internal.gateway._load_gateway_auth_models",
|
||||
lambda db, auth_context: (
|
||||
SimpleNamespace(id="user-123"),
|
||||
SimpleNamespace(id="key-123"),
|
||||
),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
"src.api.internal.gateway.get_pipeline",
|
||||
lambda: SimpleNamespace(_check_user_rate_limit=AsyncMock(return_value=None)),
|
||||
)
|
||||
|
||||
async def _fake_list_files(
|
||||
request: object,
|
||||
pageSize: int | None = None,
|
||||
pageToken: str | None = None,
|
||||
) -> JSONResponse:
|
||||
del request
|
||||
assert pageSize == 25
|
||||
assert pageToken == "page-2"
|
||||
return JSONResponse(status_code=200, content={"files": [{"name": "files/abc"}]})
|
||||
|
||||
monkeypatch.setattr("src.api.public.gemini_files.list_files", _fake_list_files)
|
||||
|
||||
client = TestClient(app, base_url="http://127.0.0.1")
|
||||
response = client.post(
|
||||
"/api/internal/gateway/execute-sync",
|
||||
json={
|
||||
"trace_id": "trace-files-list",
|
||||
"method": "GET",
|
||||
"path": "/v1beta/files",
|
||||
"query_string": "pageSize=25&pageToken=page-2",
|
||||
"headers": {"x-goog-api-key": "client-key"},
|
||||
"auth_context": {
|
||||
"user_id": "user-123",
|
||||
"api_key_id": "key-123",
|
||||
"access_allowed": True,
|
||||
},
|
||||
},
|
||||
)
|
||||
|
||||
assert response.status_code == 200
|
||||
assert response.headers[CONTROL_EXECUTED_HEADER] == "true"
|
||||
assert response.json() == {"files": [{"name": "files/abc"}]}
|
||||
|
||||
|
||||
def test_execute_sync_route_handles_gemini_files_upload_raw_body(
|
||||
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)
|
||||
monkeypatch.setattr(
|
||||
"src.api.internal.gateway._load_gateway_auth_models",
|
||||
lambda db, auth_context: (
|
||||
SimpleNamespace(id="user-123"),
|
||||
SimpleNamespace(id="key-123"),
|
||||
),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
"src.api.internal.gateway.get_pipeline",
|
||||
lambda: SimpleNamespace(_check_user_rate_limit=AsyncMock(return_value=None)),
|
||||
)
|
||||
|
||||
async def _fake_upload_file(request: object) -> JSONResponse:
|
||||
body = await request.body()
|
||||
assert body == b"upload-bytes"
|
||||
assert request.headers["content-type"] == "application/octet-stream"
|
||||
return JSONResponse(status_code=201, content={"uploaded": True})
|
||||
|
||||
monkeypatch.setattr("src.api.public.gemini_files.upload_file", _fake_upload_file)
|
||||
|
||||
client = TestClient(app, base_url="http://127.0.0.1")
|
||||
response = client.post(
|
||||
"/api/internal/gateway/execute-sync",
|
||||
json={
|
||||
"trace_id": "trace-files-upload",
|
||||
"method": "POST",
|
||||
"path": "/upload/v1beta/files",
|
||||
"query_string": "uploadType=resumable",
|
||||
"headers": {
|
||||
"x-goog-api-key": "client-key",
|
||||
"content-type": "application/octet-stream",
|
||||
},
|
||||
"body_base64": base64.b64encode(b"upload-bytes").decode("ascii"),
|
||||
"auth_context": {
|
||||
"user_id": "user-123",
|
||||
"api_key_id": "key-123",
|
||||
"access_allowed": True,
|
||||
},
|
||||
},
|
||||
)
|
||||
|
||||
assert response.status_code == 201
|
||||
assert response.headers[CONTROL_EXECUTED_HEADER] == "true"
|
||||
assert response.json() == {"uploaded": True}
|
||||
|
||||
|
||||
def test_execute_sync_route_handles_openai_video_remix_with_original_request(
|
||||
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)
|
||||
|
||||
class FakeVideoAdapter:
|
||||
mode = SimpleNamespace(value="standard")
|
||||
allowed_api_formats = ["openai:video"]
|
||||
|
||||
def authorize(self, context: object) -> None:
|
||||
self.authorized_context = context
|
||||
|
||||
async def handle(self, context: object) -> JSONResponse:
|
||||
assert context.request.method == "POST"
|
||||
assert context.request.url.path == "/v1/videos/task-123/remix"
|
||||
assert context.path_params == {"task_id": "task-123"}
|
||||
assert await context.ensure_json_body_async() == {"prompt": "remix this"}
|
||||
return JSONResponse(status_code=200, content={"video": True})
|
||||
|
||||
monkeypatch.setattr(
|
||||
"src.api.internal.gateway._resolve_gateway_sync_adapter",
|
||||
lambda decision, path: (FakeVideoAdapter(), {"task_id": "task-123"}),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
"src.api.internal.gateway._load_gateway_auth_models",
|
||||
lambda db, auth_context: (
|
||||
SimpleNamespace(id="user-123"),
|
||||
SimpleNamespace(id="key-123"),
|
||||
),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
"src.api.internal.gateway.get_pipeline",
|
||||
lambda: SimpleNamespace(_check_user_rate_limit=AsyncMock(return_value=None)),
|
||||
)
|
||||
|
||||
client = TestClient(app, base_url="http://127.0.0.1")
|
||||
response = client.post(
|
||||
"/api/internal/gateway/execute-sync",
|
||||
json={
|
||||
"trace_id": "trace-video-remix",
|
||||
"method": "POST",
|
||||
"path": "/v1/videos/task-123/remix",
|
||||
"headers": {
|
||||
"content-type": "application/json",
|
||||
"user-agent": "pytest",
|
||||
},
|
||||
"body_json": {"prompt": "remix this"},
|
||||
"auth_context": {
|
||||
"user_id": "user-123",
|
||||
"api_key_id": "key-123",
|
||||
"access_allowed": True,
|
||||
},
|
||||
},
|
||||
)
|
||||
|
||||
assert response.status_code == 200
|
||||
assert response.headers[CONTROL_EXECUTED_HEADER] == "true"
|
||||
assert response.json() == {"video": True}
|
||||
|
||||
|
||||
def test_plan_stream_route_returns_executor_plan_for_gemini_files_download(
|
||||
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)
|
||||
|
||||
fake_plan = {
|
||||
"request_id": "plan-123",
|
||||
"provider_id": "provider-123",
|
||||
"endpoint_id": "endpoint-123",
|
||||
"key_id": "key-123",
|
||||
"method": "GET",
|
||||
"url": "https://example.com/v1beta/files/file-123:download",
|
||||
"headers": {"x-goog-api-key": "upstream-key"},
|
||||
"body": {},
|
||||
"stream": True,
|
||||
"provider_api_format": "gemini:files",
|
||||
"client_api_format": "gemini:files",
|
||||
"model_name": "gemini-files",
|
||||
}
|
||||
|
||||
monkeypatch.setattr(
|
||||
"src.api.internal.gateway._load_gateway_auth_models",
|
||||
lambda db, auth_context: (
|
||||
SimpleNamespace(id="user-123"),
|
||||
SimpleNamespace(id="key-123"),
|
||||
),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
"src.api.internal.gateway.get_pipeline",
|
||||
lambda: SimpleNamespace(_check_user_rate_limit=AsyncMock(return_value=None)),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
"src.api.internal.gateway._build_gemini_files_download_stream_plan",
|
||||
AsyncMock(return_value=fake_plan),
|
||||
)
|
||||
|
||||
client = TestClient(app, base_url="http://127.0.0.1")
|
||||
response = client.post(
|
||||
"/api/internal/gateway/plan-stream",
|
||||
json={
|
||||
"trace_id": "trace-files-plan",
|
||||
"method": "GET",
|
||||
"path": "/v1beta/files/file-123:download",
|
||||
"query_string": "alt=media",
|
||||
"headers": {"x-goog-api-key": "client-key"},
|
||||
"auth_context": {
|
||||
"user_id": "user-123",
|
||||
"api_key_id": "key-123",
|
||||
"access_allowed": True,
|
||||
},
|
||||
},
|
||||
)
|
||||
|
||||
assert response.status_code == 200
|
||||
assert response.json() == {
|
||||
"action": "executor_stream",
|
||||
"plan_kind": "gemini_files_download",
|
||||
"plan": fake_plan,
|
||||
}
|
||||
|
||||
|
||||
def test_plan_stream_route_returns_executor_plan_for_openai_video_content(
|
||||
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)
|
||||
|
||||
fake_plan = {
|
||||
"request_id": "plan-video-123",
|
||||
"provider_id": "provider-video-123",
|
||||
"endpoint_id": "endpoint-video-123",
|
||||
"key_id": "key-video-123",
|
||||
"method": "GET",
|
||||
"url": "https://api.openai.com/v1/videos/ext-123/content",
|
||||
"headers": {"authorization": "Bearer upstream-key"},
|
||||
"body": {},
|
||||
"stream": True,
|
||||
"provider_api_format": "openai:video",
|
||||
"client_api_format": "openai:video",
|
||||
"model_name": "sora-2",
|
||||
}
|
||||
|
||||
monkeypatch.setattr(
|
||||
"src.api.internal.gateway._load_gateway_auth_models",
|
||||
lambda db, auth_context: (
|
||||
SimpleNamespace(id="user-123"),
|
||||
SimpleNamespace(id="key-123"),
|
||||
),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
"src.api.internal.gateway.get_pipeline",
|
||||
lambda: SimpleNamespace(_check_user_rate_limit=AsyncMock(return_value=None)),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
"src.api.internal.gateway._build_openai_video_content_stream_plan",
|
||||
AsyncMock(return_value=fake_plan),
|
||||
)
|
||||
|
||||
client = TestClient(app, base_url="http://127.0.0.1")
|
||||
response = client.post(
|
||||
"/api/internal/gateway/plan-stream",
|
||||
json={
|
||||
"trace_id": "trace-video-plan",
|
||||
"method": "GET",
|
||||
"path": "/v1/videos/task-123/content",
|
||||
"query_string": "variant=video",
|
||||
"headers": {"authorization": "Bearer client-key"},
|
||||
"auth_context": {
|
||||
"user_id": "user-123",
|
||||
"api_key_id": "key-123",
|
||||
"access_allowed": True,
|
||||
},
|
||||
},
|
||||
)
|
||||
|
||||
assert response.status_code == 200
|
||||
assert response.json() == {
|
||||
"action": "executor_stream",
|
||||
"plan_kind": "openai_video_content",
|
||||
"plan": fake_plan,
|
||||
}
|
||||
@@ -529,6 +529,101 @@ class TestPipelineAuthentication:
|
||||
assert exc_info.value.status_code == 403
|
||||
assert "锁定" in str(exc_info.value.detail)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_authenticate_client_reuses_trusted_gateway_auth_context(
|
||||
self, pipeline: ApiRequestPipeline
|
||||
) -> None:
|
||||
mock_request = MagicMock()
|
||||
mock_request.headers = {
|
||||
"x-aether-gateway": "rust-phase3b",
|
||||
"x-aether-auth-user-id": "user-123",
|
||||
"x-aether-auth-api-key-id": "key-123",
|
||||
"x-aether-auth-balance-remaining": "42.5",
|
||||
"x-aether-auth-access-allowed": "true",
|
||||
}
|
||||
mock_request.client = SimpleNamespace(host="127.0.0.1")
|
||||
mock_request.url.path = "/v1/chat/completions"
|
||||
mock_request.state = MagicMock()
|
||||
|
||||
db_user = MagicMock()
|
||||
db_user.id = "user-123"
|
||||
db_user.is_active = True
|
||||
db_user.is_deleted = False
|
||||
db_api_key = MagicMock()
|
||||
db_api_key.id = "key-123"
|
||||
db_api_key.user_id = "user-123"
|
||||
db_api_key.is_active = True
|
||||
db_api_key.is_locked = False
|
||||
db_api_key.is_standalone = False
|
||||
db_api_key.expires_at = None
|
||||
|
||||
mock_db = MagicMock()
|
||||
user_query = MagicMock()
|
||||
user_query.filter.return_value.first.return_value = db_user
|
||||
api_key_query = MagicMock()
|
||||
api_key_query.filter.return_value.first.return_value = db_api_key
|
||||
mock_db.query.side_effect = [user_query, api_key_query]
|
||||
|
||||
mock_adapter = MagicMock()
|
||||
mock_adapter.extract_api_key.side_effect = AssertionError("trusted auth should short-circuit")
|
||||
|
||||
with patch.object(
|
||||
pipeline.auth_service,
|
||||
"authenticate_api_key_threadsafe",
|
||||
new_callable=AsyncMock,
|
||||
side_effect=AssertionError("threadsafe auth should not be called"),
|
||||
):
|
||||
user, api_key = await pipeline._authenticate_client(mock_request, mock_db, mock_adapter)
|
||||
|
||||
assert user == db_user
|
||||
assert api_key == db_api_key
|
||||
assert mock_request.state.user_id == "user-123"
|
||||
assert mock_request.state.api_key_id == "key-123"
|
||||
assert mock_request.state.prefetched_balance_remaining == 42.5
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_authenticate_client_trusted_gateway_balance_denied(
|
||||
self, pipeline: ApiRequestPipeline
|
||||
) -> None:
|
||||
mock_request = MagicMock()
|
||||
mock_request.headers = {
|
||||
"x-aether-gateway": "rust-phase3b",
|
||||
"x-aether-auth-user-id": "user-123",
|
||||
"x-aether-auth-api-key-id": "key-123",
|
||||
"x-aether-auth-balance-remaining": "0",
|
||||
"x-aether-auth-access-allowed": "false",
|
||||
}
|
||||
mock_request.client = SimpleNamespace(host="127.0.0.1")
|
||||
mock_request.url.path = "/v1/chat/completions"
|
||||
mock_request.state = MagicMock()
|
||||
|
||||
db_user = MagicMock()
|
||||
db_user.id = "user-123"
|
||||
db_user.is_active = True
|
||||
db_user.is_deleted = False
|
||||
db_api_key = MagicMock()
|
||||
db_api_key.id = "key-123"
|
||||
db_api_key.user_id = "user-123"
|
||||
db_api_key.is_active = True
|
||||
db_api_key.is_locked = False
|
||||
db_api_key.is_standalone = False
|
||||
db_api_key.expires_at = None
|
||||
|
||||
mock_db = MagicMock()
|
||||
user_query = MagicMock()
|
||||
user_query.filter.return_value.first.return_value = db_user
|
||||
api_key_query = MagicMock()
|
||||
api_key_query.filter.return_value.first.return_value = db_api_key
|
||||
mock_db.query.side_effect = [user_query, api_key_query]
|
||||
|
||||
mock_adapter = MagicMock()
|
||||
mock_adapter.extract_api_key.side_effect = AssertionError("trusted auth should short-circuit")
|
||||
|
||||
from src.core.exceptions import BalanceInsufficientException
|
||||
|
||||
with pytest.raises(BalanceInsufficientException):
|
||||
await pipeline._authenticate_client(mock_request, mock_db, mock_adapter)
|
||||
|
||||
|
||||
class TestPipelineUserRateLimit:
|
||||
@pytest.fixture
|
||||
|
||||
899
tests/e2e_rust_executor.py
Normal file
899
tests/e2e_rust_executor.py
Normal file
@@ -0,0 +1,899 @@
|
||||
"""
|
||||
aether-executor 本地端到端测试
|
||||
|
||||
测试流程:
|
||||
1. 启动本地假上游 HTTP 服务
|
||||
2. 启动 aether-executor(Unix Socket)
|
||||
3. 用 Python RustExecutorClient 发送 ExecutionPlan
|
||||
4. 验证执行结果与上游响应一致
|
||||
|
||||
运行:
|
||||
cargo build -p aether-executor
|
||||
uv run python tests/e2e_rust_executor.py
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import base64
|
||||
import contextlib
|
||||
import gzip
|
||||
import json
|
||||
import signal
|
||||
import subprocess
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
import httpx
|
||||
|
||||
from src.services.request.executor_plan import (
|
||||
ExecutionPlan,
|
||||
ExecutionPlanBody,
|
||||
ExecutionProxySnapshot,
|
||||
ExecutionPlanTimeouts,
|
||||
)
|
||||
from src.services.request.rust_executor_client import RustExecutorClient
|
||||
|
||||
|
||||
async def _upstream_app(scope, receive, send) -> None: # type: ignore[no-untyped-def]
|
||||
assert scope["type"] == "http"
|
||||
body = b""
|
||||
while True:
|
||||
message = await receive()
|
||||
if message["type"] != "http.request":
|
||||
continue
|
||||
body += message.get("body", b"")
|
||||
if not message.get("more_body", False):
|
||||
break
|
||||
|
||||
header_map = {
|
||||
key.decode("latin1").lower(): value.decode("latin1")
|
||||
for key, value in scope.get("headers", [])
|
||||
}
|
||||
path = str(scope.get("path") or "")
|
||||
content_encoding = header_map.get("content-encoding")
|
||||
decoded_body = gzip.decompress(body) if content_encoding == "gzip" else body
|
||||
|
||||
if path == "/v1/raw":
|
||||
response_body = json.dumps(
|
||||
{
|
||||
"ok": True,
|
||||
"path": path,
|
||||
"received_body_text": decoded_body.decode("utf-8", errors="replace"),
|
||||
"received_content_type": header_map.get("content-type"),
|
||||
"received_content_encoding": content_encoding,
|
||||
},
|
||||
ensure_ascii=False,
|
||||
).encode("utf-8")
|
||||
await send(
|
||||
{
|
||||
"type": "http.response.start",
|
||||
"status": 200,
|
||||
"headers": [
|
||||
[b"content-type", b"application/json"],
|
||||
[b"x-upstream-test", b"true"],
|
||||
],
|
||||
}
|
||||
)
|
||||
await send(
|
||||
{
|
||||
"type": "http.response.body",
|
||||
"body": response_body,
|
||||
"more_body": False,
|
||||
}
|
||||
)
|
||||
return
|
||||
|
||||
if path == "/v1/download":
|
||||
await send(
|
||||
{
|
||||
"type": "http.response.start",
|
||||
"status": 200,
|
||||
"headers": [
|
||||
[b"content-type", b"video/mp4"],
|
||||
[b"x-download-test", b"true"],
|
||||
],
|
||||
}
|
||||
)
|
||||
await send(
|
||||
{
|
||||
"type": "http.response.body",
|
||||
"body": b"video-download-bytes",
|
||||
"more_body": False,
|
||||
}
|
||||
)
|
||||
return
|
||||
|
||||
if path == "/v1/delete" and scope.get("method") == "DELETE":
|
||||
await send(
|
||||
{
|
||||
"type": "http.response.start",
|
||||
"status": 204,
|
||||
"headers": [[b"x-delete-test", b"true"]],
|
||||
}
|
||||
)
|
||||
await send(
|
||||
{
|
||||
"type": "http.response.body",
|
||||
"body": b"",
|
||||
"more_body": False,
|
||||
}
|
||||
)
|
||||
return
|
||||
|
||||
payload = json.loads(decoded_body.decode("utf-8"))
|
||||
if payload.get("stream") is True:
|
||||
stream_body = (
|
||||
"data: "
|
||||
+ json.dumps(
|
||||
{
|
||||
"id": "chatcmpl-stream-1",
|
||||
"object": "chat.completion.chunk",
|
||||
"model": payload.get("model"),
|
||||
"choices": [
|
||||
{
|
||||
"index": 0,
|
||||
"delta": {"role": "assistant", "content": "hello"},
|
||||
"finish_reason": None,
|
||||
}
|
||||
],
|
||||
},
|
||||
ensure_ascii=False,
|
||||
)
|
||||
+ "\n\n"
|
||||
+ "data: "
|
||||
+ json.dumps(
|
||||
{
|
||||
"id": "chatcmpl-stream-1",
|
||||
"object": "chat.completion.chunk",
|
||||
"model": payload.get("model"),
|
||||
"choices": [
|
||||
{
|
||||
"index": 0,
|
||||
"delta": {},
|
||||
"finish_reason": "stop",
|
||||
}
|
||||
],
|
||||
"usage": {
|
||||
"prompt_tokens": 3,
|
||||
"completion_tokens": 1,
|
||||
"total_tokens": 4,
|
||||
},
|
||||
},
|
||||
ensure_ascii=False,
|
||||
)
|
||||
+ "\n\n"
|
||||
+ "data: [DONE]\n\n"
|
||||
).encode("utf-8")
|
||||
await send(
|
||||
{
|
||||
"type": "http.response.start",
|
||||
"status": 200,
|
||||
"headers": [
|
||||
[b"content-type", b"text/event-stream"],
|
||||
[b"x-upstream-test", b"true"],
|
||||
],
|
||||
}
|
||||
)
|
||||
await send(
|
||||
{
|
||||
"type": "http.response.body",
|
||||
"body": stream_body,
|
||||
"more_body": False,
|
||||
}
|
||||
)
|
||||
return
|
||||
|
||||
response = {
|
||||
"ok": True,
|
||||
"received_model": payload.get("model"),
|
||||
"received_messages": payload.get("messages", []),
|
||||
"path": path,
|
||||
"received_content_encoding": content_encoding,
|
||||
}
|
||||
response_body = json.dumps(response, ensure_ascii=False).encode("utf-8")
|
||||
|
||||
await send(
|
||||
{
|
||||
"type": "http.response.start",
|
||||
"status": 200,
|
||||
"headers": [
|
||||
[b"content-type", b"application/json"],
|
||||
[b"x-upstream-test", b"true"],
|
||||
],
|
||||
}
|
||||
)
|
||||
await send(
|
||||
{
|
||||
"type": "http.response.body",
|
||||
"body": response_body,
|
||||
"more_body": False,
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
async def _relay_app(scope, receive, send) -> None: # type: ignore[no-untyped-def]
|
||||
assert scope["type"] == "http"
|
||||
body = b""
|
||||
while True:
|
||||
message = await receive()
|
||||
if message["type"] != "http.request":
|
||||
continue
|
||||
body += message.get("body", b"")
|
||||
if not message.get("more_body", False):
|
||||
break
|
||||
|
||||
if len(body) < 4:
|
||||
await send(
|
||||
{
|
||||
"type": "http.response.start",
|
||||
"status": 400,
|
||||
"headers": [[b"content-type", b"text/plain; charset=utf-8"]],
|
||||
}
|
||||
)
|
||||
await send({"type": "http.response.body", "body": b"invalid relay envelope", "more_body": False})
|
||||
return
|
||||
|
||||
path = str(scope.get("path") or "")
|
||||
node_id = path.rsplit("/", 1)[-1]
|
||||
meta_len = int.from_bytes(body[:4], "big")
|
||||
if meta_len <= 0 or len(body) < 4 + meta_len:
|
||||
await send(
|
||||
{
|
||||
"type": "http.response.start",
|
||||
"status": 400,
|
||||
"headers": [[b"content-type", b"text/plain; charset=utf-8"]],
|
||||
}
|
||||
)
|
||||
await send({"type": "http.response.body", "body": b"invalid relay metadata", "more_body": False})
|
||||
return
|
||||
meta = json.loads(body[4 : 4 + meta_len].decode("utf-8"))
|
||||
request_json = json.loads(body[4 + meta_len :].decode("utf-8"))
|
||||
|
||||
if request_json.get("stream") is True:
|
||||
stream_body = (
|
||||
"data: "
|
||||
+ json.dumps(
|
||||
{
|
||||
"id": "relay-stream-1",
|
||||
"object": "chat.completion.chunk",
|
||||
"model": request_json.get("model"),
|
||||
"choices": [{"index": 0, "delta": {"content": "relay"}, "finish_reason": None}],
|
||||
},
|
||||
ensure_ascii=False,
|
||||
)
|
||||
+ "\n\n"
|
||||
+ "data: [DONE]\n\n"
|
||||
).encode("utf-8")
|
||||
await send(
|
||||
{
|
||||
"type": "http.response.start",
|
||||
"status": 200,
|
||||
"headers": [
|
||||
[b"content-type", b"text/event-stream"],
|
||||
[b"x-relay-node", node_id.encode("utf-8")],
|
||||
],
|
||||
}
|
||||
)
|
||||
await send({"type": "http.response.body", "body": stream_body, "more_body": False})
|
||||
return
|
||||
|
||||
response_body = json.dumps(
|
||||
{
|
||||
"ok": True,
|
||||
"via_tunnel": True,
|
||||
"node_id": node_id,
|
||||
"meta_url": meta.get("url"),
|
||||
"received_model": request_json.get("model"),
|
||||
},
|
||||
ensure_ascii=False,
|
||||
).encode("utf-8")
|
||||
await send(
|
||||
{
|
||||
"type": "http.response.start",
|
||||
"status": 200,
|
||||
"headers": [
|
||||
[b"content-type", b"application/json"],
|
||||
[b"x-relay-node", node_id.encode("utf-8")],
|
||||
],
|
||||
}
|
||||
)
|
||||
await send(
|
||||
{
|
||||
"type": "http.response.body",
|
||||
"body": response_body,
|
||||
"more_body": False,
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
async def _start_proxy_server() -> tuple[asyncio.AbstractServer, int, asyncio.Future[str]]:
|
||||
loop = asyncio.get_running_loop()
|
||||
request_line_future: asyncio.Future[str] = loop.create_future()
|
||||
|
||||
async def _handle_proxy(
|
||||
reader: asyncio.StreamReader,
|
||||
writer: asyncio.StreamWriter,
|
||||
) -> None:
|
||||
try:
|
||||
header_bytes = await reader.readuntil(b"\r\n\r\n")
|
||||
header_text = header_bytes.decode("latin-1")
|
||||
lines = header_text.split("\r\n")
|
||||
if lines and not request_line_future.done():
|
||||
request_line_future.set_result(lines[0])
|
||||
|
||||
content_length = 0
|
||||
for line in lines[1:]:
|
||||
if not line:
|
||||
continue
|
||||
if line.lower().startswith("content-length:"):
|
||||
content_length = int(line.split(":", 1)[1].strip() or "0")
|
||||
break
|
||||
|
||||
if content_length:
|
||||
await reader.readexactly(content_length)
|
||||
|
||||
body = json.dumps(
|
||||
{
|
||||
"ok": True,
|
||||
"via_proxy": True,
|
||||
"request_line": lines[0] if lines else "",
|
||||
},
|
||||
ensure_ascii=False,
|
||||
).encode("utf-8")
|
||||
writer.write(
|
||||
b"HTTP/1.1 200 OK\r\n"
|
||||
b"content-type: application/json\r\n"
|
||||
b"x-proxy-test: true\r\n"
|
||||
+ f"content-length: {len(body)}\r\n\r\n".encode("ascii")
|
||||
+ body
|
||||
)
|
||||
await writer.drain()
|
||||
finally:
|
||||
writer.close()
|
||||
with contextlib.suppress(Exception):
|
||||
await writer.wait_closed()
|
||||
|
||||
server = await asyncio.start_server(_handle_proxy, "127.0.0.1", 0)
|
||||
port = server.sockets[0].getsockname()[1]
|
||||
return server, port, request_line_future
|
||||
|
||||
|
||||
async def run_test() -> bool:
|
||||
repo_root = Path(__file__).resolve().parents[1]
|
||||
executor_binary = repo_root / "target" / "debug" / "aether-executor"
|
||||
if not executor_binary.is_file():
|
||||
print(f"FAIL: executor binary not found at {executor_binary}")
|
||||
print(" run: cargo build -p aether-executor")
|
||||
return False
|
||||
|
||||
upstream_listener = await asyncio.start_server(lambda r, w: None, "127.0.0.1", 0)
|
||||
upstream_port = upstream_listener.sockets[0].getsockname()[1]
|
||||
upstream_listener.close()
|
||||
await upstream_listener.wait_closed()
|
||||
|
||||
import uvicorn
|
||||
|
||||
upstream_config = uvicorn.Config(
|
||||
_upstream_app,
|
||||
host="127.0.0.1",
|
||||
port=upstream_port,
|
||||
log_level="warning",
|
||||
)
|
||||
upstream_server = uvicorn.Server(upstream_config)
|
||||
upstream_task = asyncio.create_task(upstream_server.serve())
|
||||
|
||||
relay_listener = await asyncio.start_server(lambda r, w: None, "127.0.0.1", 0)
|
||||
relay_port = relay_listener.sockets[0].getsockname()[1]
|
||||
relay_listener.close()
|
||||
await relay_listener.wait_closed()
|
||||
|
||||
relay_config = uvicorn.Config(
|
||||
_relay_app,
|
||||
host="127.0.0.1",
|
||||
port=relay_port,
|
||||
log_level="warning",
|
||||
)
|
||||
relay_server = uvicorn.Server(relay_config)
|
||||
relay_task = asyncio.create_task(relay_server.serve())
|
||||
|
||||
executor_socket = Path("/tmp/aether-executor-e2e.sock")
|
||||
executor_socket.unlink(missing_ok=True)
|
||||
executor_proc = subprocess.Popen(
|
||||
[
|
||||
str(executor_binary),
|
||||
"--transport",
|
||||
"unix_socket",
|
||||
"--unix-socket",
|
||||
str(executor_socket),
|
||||
],
|
||||
stdout=subprocess.PIPE,
|
||||
stderr=subprocess.STDOUT,
|
||||
text=True,
|
||||
)
|
||||
|
||||
try:
|
||||
print("[1/4] Waiting for upstream server ...")
|
||||
async with httpx.AsyncClient() as client:
|
||||
for _ in range(50):
|
||||
try:
|
||||
resp = await client.post(
|
||||
f"http://127.0.0.1:{upstream_port}/healthz",
|
||||
json={"model": "probe"},
|
||||
timeout=0.2,
|
||||
)
|
||||
if resp.status_code == 200:
|
||||
break
|
||||
except Exception:
|
||||
await asyncio.sleep(0.1)
|
||||
else:
|
||||
print("FAIL: upstream server did not start in time")
|
||||
return False
|
||||
|
||||
for _ in range(50):
|
||||
try:
|
||||
resp = await client.post(
|
||||
f"http://127.0.0.1:{relay_port}/local/relay/probe",
|
||||
content=(0).to_bytes(4, "big"),
|
||||
timeout=0.2,
|
||||
)
|
||||
if resp.status_code in {400, 500}:
|
||||
break
|
||||
except Exception:
|
||||
await asyncio.sleep(0.1)
|
||||
else:
|
||||
print("FAIL: relay server did not start in time")
|
||||
return False
|
||||
|
||||
print("[2/4] Waiting for aether-executor ...")
|
||||
for _ in range(100):
|
||||
if executor_socket.exists():
|
||||
try:
|
||||
async with httpx.AsyncClient(
|
||||
transport=httpx.AsyncHTTPTransport(uds=str(executor_socket)),
|
||||
base_url="http://127.0.0.1:5219",
|
||||
) as client:
|
||||
resp = await client.get("/health", timeout=0.2)
|
||||
if resp.status_code == 200:
|
||||
break
|
||||
except Exception:
|
||||
pass
|
||||
await asyncio.sleep(0.1)
|
||||
else:
|
||||
print("FAIL: executor did not start in time")
|
||||
if executor_proc.stdout is not None:
|
||||
print(executor_proc.stdout.read())
|
||||
return False
|
||||
|
||||
print("[3/4] Sending ExecutionPlan via RustExecutorClient ...")
|
||||
client = RustExecutorClient(
|
||||
transport="unix_socket",
|
||||
socket_path=str(executor_socket),
|
||||
base_url="http://127.0.0.1:5219",
|
||||
request_timeout=5.0,
|
||||
)
|
||||
result = await client.execute_sync_json(
|
||||
ExecutionPlan(
|
||||
request_id="req-e2e-1",
|
||||
candidate_id=None,
|
||||
provider_name="openai",
|
||||
provider_id="prov-1",
|
||||
endpoint_id="ep-1",
|
||||
key_id="key-1",
|
||||
method="POST",
|
||||
url=f"http://127.0.0.1:{upstream_port}/v1/chat/completions",
|
||||
headers={"content-type": "application/json"},
|
||||
body=ExecutionPlanBody(
|
||||
json_body={
|
||||
"model": "gpt-4.1",
|
||||
"messages": [{"role": "user", "content": "hello"}],
|
||||
}
|
||||
),
|
||||
stream=False,
|
||||
provider_api_format="openai:chat",
|
||||
client_api_format="openai:chat",
|
||||
model_name="gpt-4.1",
|
||||
content_type="application/json",
|
||||
timeouts=ExecutionPlanTimeouts(
|
||||
connect_ms=5_000,
|
||||
read_ms=5_000,
|
||||
write_ms=5_000,
|
||||
pool_ms=5_000,
|
||||
total_ms=5_000,
|
||||
),
|
||||
)
|
||||
)
|
||||
|
||||
print("[4/4] Verifying response ...")
|
||||
assert result.status_code == 200
|
||||
assert result.headers.get("x-upstream-test") == "true"
|
||||
assert result.response_json["ok"] is True
|
||||
assert result.response_json["received_model"] == "gpt-4.1"
|
||||
assert result.response_json["path"] == "/v1/chat/completions"
|
||||
|
||||
print("[extra] Verifying tunnel relay sync execution ...")
|
||||
tunnel_result = await client.execute_sync_json(
|
||||
ExecutionPlan(
|
||||
request_id="req-e2e-tunnel-1",
|
||||
candidate_id=None,
|
||||
provider_name="openai",
|
||||
provider_id="prov-1",
|
||||
endpoint_id="ep-1",
|
||||
key_id="key-1",
|
||||
method="POST",
|
||||
url="https://upstream-via-relay.test/v1/chat/completions",
|
||||
headers={"content-type": "application/json"},
|
||||
body=ExecutionPlanBody(
|
||||
json_body={
|
||||
"model": "gpt-4.1",
|
||||
"messages": [{"role": "user", "content": "hello tunnel"}],
|
||||
}
|
||||
),
|
||||
stream=False,
|
||||
provider_api_format="openai:chat",
|
||||
client_api_format="openai:chat",
|
||||
model_name="gpt-4.1",
|
||||
content_type="application/json",
|
||||
proxy=ExecutionProxySnapshot(
|
||||
enabled=True,
|
||||
mode="tunnel",
|
||||
node_id="node-1",
|
||||
label="relay-node",
|
||||
extra={"hub_base_url": f"http://127.0.0.1:{relay_port}"},
|
||||
),
|
||||
timeouts=ExecutionPlanTimeouts(
|
||||
connect_ms=5_000,
|
||||
read_ms=5_000,
|
||||
write_ms=5_000,
|
||||
pool_ms=5_000,
|
||||
total_ms=5_000,
|
||||
),
|
||||
)
|
||||
)
|
||||
assert tunnel_result.status_code == 200
|
||||
assert tunnel_result.headers.get("x-relay-node") == "node-1"
|
||||
assert tunnel_result.response_json["via_tunnel"] is True
|
||||
assert tunnel_result.response_json["meta_url"] == "https://upstream-via-relay.test/v1/chat/completions"
|
||||
|
||||
print("[extra] Verifying HTTP proxy execution ...")
|
||||
proxy_server, proxy_port, proxy_request_line = await _start_proxy_server()
|
||||
try:
|
||||
proxied_result = await client.execute_sync_json(
|
||||
ExecutionPlan(
|
||||
request_id="req-e2e-proxy-1",
|
||||
candidate_id=None,
|
||||
provider_name="openai",
|
||||
provider_id="prov-1",
|
||||
endpoint_id="ep-1",
|
||||
key_id="key-1",
|
||||
method="POST",
|
||||
url=f"http://127.0.0.1:{upstream_port}/v1/chat/completions",
|
||||
headers={"content-type": "application/json"},
|
||||
body=ExecutionPlanBody(
|
||||
json_body={
|
||||
"model": "gpt-4.1",
|
||||
"messages": [{"role": "user", "content": "proxy hello"}],
|
||||
}
|
||||
),
|
||||
stream=False,
|
||||
provider_api_format="openai:chat",
|
||||
client_api_format="openai:chat",
|
||||
model_name="gpt-4.1",
|
||||
content_type="application/json",
|
||||
proxy=ExecutionProxySnapshot(
|
||||
enabled=True,
|
||||
mode="http",
|
||||
label="local-proxy",
|
||||
url=f"http://127.0.0.1:{proxy_port}",
|
||||
),
|
||||
timeouts=ExecutionPlanTimeouts(
|
||||
connect_ms=5_000,
|
||||
read_ms=5_000,
|
||||
write_ms=5_000,
|
||||
pool_ms=5_000,
|
||||
total_ms=5_000,
|
||||
),
|
||||
)
|
||||
)
|
||||
request_line = await asyncio.wait_for(proxy_request_line, timeout=5.0)
|
||||
finally:
|
||||
proxy_server.close()
|
||||
await proxy_server.wait_closed()
|
||||
|
||||
assert proxied_result.status_code == 200
|
||||
assert proxied_result.headers.get("x-proxy-test") == "true"
|
||||
assert proxied_result.response_json["via_proxy"] is True
|
||||
assert request_line.startswith("POST http://127.0.0.1:")
|
||||
assert "/v1/chat/completions HTTP/1.1" in request_line
|
||||
|
||||
print("[extra] Verifying upstream-stream raw byte path ...")
|
||||
stream_result = await client.execute_sync_json(
|
||||
ExecutionPlan(
|
||||
request_id="req-e2e-stream-1",
|
||||
candidate_id=None,
|
||||
provider_name="openai",
|
||||
provider_id="prov-1",
|
||||
endpoint_id="ep-1",
|
||||
key_id="key-1",
|
||||
method="POST",
|
||||
url=f"http://127.0.0.1:{upstream_port}/v1/chat/completions",
|
||||
headers={"content-type": "application/json"},
|
||||
body=ExecutionPlanBody(
|
||||
json_body={
|
||||
"model": "gpt-4.1",
|
||||
"messages": [{"role": "user", "content": "hello stream"}],
|
||||
"stream": True,
|
||||
}
|
||||
),
|
||||
stream=True,
|
||||
provider_api_format="openai:chat",
|
||||
client_api_format="openai:chat",
|
||||
model_name="gpt-4.1",
|
||||
content_type="application/json",
|
||||
timeouts=ExecutionPlanTimeouts(
|
||||
connect_ms=5_000,
|
||||
read_ms=5_000,
|
||||
write_ms=5_000,
|
||||
pool_ms=5_000,
|
||||
total_ms=5_000,
|
||||
),
|
||||
)
|
||||
)
|
||||
assert stream_result.status_code == 200
|
||||
assert stream_result.response_json is None
|
||||
assert stream_result.response_body_bytes is not None
|
||||
assert b"chat.completion.chunk" in stream_result.response_body_bytes
|
||||
assert b"[DONE]" in stream_result.response_body_bytes
|
||||
|
||||
print("[extra] Verifying gzip request body path ...")
|
||||
gzip_result = await client.execute_sync_json(
|
||||
ExecutionPlan(
|
||||
request_id="req-e2e-gzip-1",
|
||||
candidate_id="cand-e2e-gzip-1",
|
||||
provider_name="openai",
|
||||
provider_id="prov-e2e",
|
||||
endpoint_id="ep-e2e",
|
||||
key_id="key-e2e",
|
||||
method="POST",
|
||||
url=f"http://127.0.0.1:{upstream_port}/v1/chat/completions",
|
||||
headers={"content-type": "application/json"},
|
||||
body=ExecutionPlanBody(
|
||||
json_body={
|
||||
"model": "gpt-4.1",
|
||||
"messages": [{"role": "user", "content": "hello gzip"}],
|
||||
}
|
||||
),
|
||||
stream=False,
|
||||
provider_api_format="openai:chat",
|
||||
client_api_format="openai:chat",
|
||||
model_name="gpt-4.1",
|
||||
content_type="application/json",
|
||||
content_encoding="gzip",
|
||||
timeouts=ExecutionPlanTimeouts(connect_ms=5_000, total_ms=30_000),
|
||||
)
|
||||
)
|
||||
assert gzip_result.status_code == 200
|
||||
assert gzip_result.response_json is not None
|
||||
assert gzip_result.response_json["received_content_encoding"] == "gzip"
|
||||
assert gzip_result.response_json["received_model"] == "gpt-4.1"
|
||||
|
||||
print("[extra] Verifying raw body bytes path ...")
|
||||
raw_result = await client.execute_sync_json(
|
||||
ExecutionPlan(
|
||||
request_id="req-e2e-raw-1",
|
||||
candidate_id="cand-e2e-raw-1",
|
||||
provider_name="openai",
|
||||
provider_id="prov-e2e",
|
||||
endpoint_id="ep-e2e",
|
||||
key_id="key-e2e",
|
||||
method="POST",
|
||||
url=f"http://127.0.0.1:{upstream_port}/v1/raw",
|
||||
headers={"content-type": "text/plain"},
|
||||
body=ExecutionPlanBody(
|
||||
body_bytes_b64=base64.b64encode(b"hello raw body").decode("ascii")
|
||||
),
|
||||
stream=False,
|
||||
provider_api_format="openai:chat",
|
||||
client_api_format="openai:chat",
|
||||
model_name="gpt-4.1",
|
||||
content_type="text/plain",
|
||||
timeouts=ExecutionPlanTimeouts(connect_ms=5_000, total_ms=30_000),
|
||||
)
|
||||
)
|
||||
assert raw_result.status_code == 200
|
||||
assert raw_result.response_json is not None
|
||||
assert raw_result.response_json["received_body_text"] == "hello raw body"
|
||||
assert raw_result.response_json["received_content_type"] == "text/plain"
|
||||
|
||||
print("[extra] Verifying sync empty-body delete path ...")
|
||||
delete_result = await client.execute_sync_json(
|
||||
ExecutionPlan(
|
||||
request_id="req-e2e-delete-1",
|
||||
candidate_id=None,
|
||||
provider_name="openai",
|
||||
provider_id="prov-1",
|
||||
endpoint_id="ep-1",
|
||||
key_id="key-1",
|
||||
method="DELETE",
|
||||
url=f"http://127.0.0.1:{upstream_port}/v1/delete",
|
||||
headers={},
|
||||
body=ExecutionPlanBody(),
|
||||
stream=False,
|
||||
provider_api_format="openai:video",
|
||||
client_api_format="openai:video",
|
||||
model_name="sora-2",
|
||||
timeouts=ExecutionPlanTimeouts(
|
||||
connect_ms=5_000,
|
||||
read_ms=5_000,
|
||||
write_ms=5_000,
|
||||
pool_ms=5_000,
|
||||
total_ms=5_000,
|
||||
),
|
||||
)
|
||||
)
|
||||
assert delete_result.status_code == 204
|
||||
assert delete_result.response_json is None
|
||||
assert delete_result.response_body_bytes is None
|
||||
|
||||
print("[extra] Verifying GET/no-body stream path ...")
|
||||
download_stream = await client.execute_stream(
|
||||
ExecutionPlan(
|
||||
request_id="req-e2e-download-1",
|
||||
candidate_id=None,
|
||||
provider_name="openai",
|
||||
provider_id="prov-1",
|
||||
endpoint_id="ep-1",
|
||||
key_id="key-1",
|
||||
method="GET",
|
||||
url=f"http://127.0.0.1:{upstream_port}/v1/download",
|
||||
headers={},
|
||||
body=ExecutionPlanBody(),
|
||||
stream=True,
|
||||
provider_api_format="openai:video",
|
||||
client_api_format="openai:video",
|
||||
model_name="sora-2",
|
||||
timeouts=ExecutionPlanTimeouts(
|
||||
connect_ms=5_000,
|
||||
read_ms=5_000,
|
||||
write_ms=5_000,
|
||||
pool_ms=5_000,
|
||||
total_ms=5_000,
|
||||
),
|
||||
)
|
||||
)
|
||||
try:
|
||||
download_chunks = [chunk async for chunk in download_stream.byte_iterator]
|
||||
finally:
|
||||
await download_stream.response_ctx.__aexit__(None, None, None)
|
||||
|
||||
assert download_stream.status_code == 200
|
||||
assert download_stream.headers.get("x-download-test") == "true"
|
||||
assert b"".join(download_chunks) == b"video-download-bytes"
|
||||
|
||||
print("[extra] Verifying native stream path ...")
|
||||
live_stream = await client.execute_stream(
|
||||
ExecutionPlan(
|
||||
request_id="req-e2e-live-stream-1",
|
||||
candidate_id=None,
|
||||
provider_name="openai",
|
||||
provider_id="prov-1",
|
||||
endpoint_id="ep-1",
|
||||
key_id="key-1",
|
||||
method="POST",
|
||||
url=f"http://127.0.0.1:{upstream_port}/v1/chat/completions",
|
||||
headers={"content-type": "application/json"},
|
||||
body=ExecutionPlanBody(
|
||||
json_body={
|
||||
"model": "gpt-4.1",
|
||||
"messages": [{"role": "user", "content": "hello live stream"}],
|
||||
"stream": True,
|
||||
}
|
||||
),
|
||||
stream=True,
|
||||
provider_api_format="openai:chat",
|
||||
client_api_format="openai:chat",
|
||||
model_name="gpt-4.1",
|
||||
content_type="application/json",
|
||||
timeouts=ExecutionPlanTimeouts(
|
||||
connect_ms=5_000,
|
||||
read_ms=5_000,
|
||||
write_ms=5_000,
|
||||
pool_ms=5_000,
|
||||
total_ms=5_000,
|
||||
),
|
||||
)
|
||||
)
|
||||
try:
|
||||
live_chunks = [chunk async for chunk in live_stream.byte_iterator]
|
||||
finally:
|
||||
await live_stream.response_ctx.__aexit__(None, None, None)
|
||||
|
||||
assert live_stream.status_code == 200
|
||||
assert live_stream.headers.get("x-upstream-test") == "true"
|
||||
live_body = b"".join(live_chunks)
|
||||
assert live_body
|
||||
assert b"chat.completion.chunk" in live_body
|
||||
assert b"[DONE]" in live_body
|
||||
|
||||
print("[extra] Verifying native tunnel stream path ...")
|
||||
tunnel_stream = await client.execute_stream(
|
||||
ExecutionPlan(
|
||||
request_id="req-e2e-tunnel-stream-1",
|
||||
candidate_id=None,
|
||||
provider_name="openai",
|
||||
provider_id="prov-1",
|
||||
endpoint_id="ep-1",
|
||||
key_id="key-1",
|
||||
method="POST",
|
||||
url="https://upstream-via-relay.test/v1/chat/completions",
|
||||
headers={"content-type": "application/json"},
|
||||
body=ExecutionPlanBody(
|
||||
json_body={
|
||||
"model": "gpt-4.1",
|
||||
"messages": [{"role": "user", "content": "hello relay stream"}],
|
||||
"stream": True,
|
||||
}
|
||||
),
|
||||
stream=True,
|
||||
provider_api_format="openai:chat",
|
||||
client_api_format="openai:chat",
|
||||
model_name="gpt-4.1",
|
||||
content_type="application/json",
|
||||
proxy=ExecutionProxySnapshot(
|
||||
enabled=True,
|
||||
mode="tunnel",
|
||||
node_id="node-1",
|
||||
label="relay-node",
|
||||
extra={"hub_base_url": f"http://127.0.0.1:{relay_port}"},
|
||||
),
|
||||
timeouts=ExecutionPlanTimeouts(
|
||||
connect_ms=5_000,
|
||||
read_ms=5_000,
|
||||
write_ms=5_000,
|
||||
pool_ms=5_000,
|
||||
total_ms=5_000,
|
||||
),
|
||||
)
|
||||
)
|
||||
try:
|
||||
tunnel_stream_chunks = [chunk async for chunk in tunnel_stream.byte_iterator]
|
||||
finally:
|
||||
await tunnel_stream.response_ctx.__aexit__(None, None, None)
|
||||
|
||||
tunnel_stream_body = b"".join(tunnel_stream_chunks)
|
||||
assert tunnel_stream.status_code == 200
|
||||
assert tunnel_stream.headers.get("x-relay-node") == "node-1"
|
||||
assert b"relay-stream-1" in tunnel_stream_body
|
||||
assert b"[DONE]" in tunnel_stream_body
|
||||
print("ALL TESTS PASSED")
|
||||
return True
|
||||
|
||||
finally:
|
||||
upstream_server.should_exit = True
|
||||
try:
|
||||
await asyncio.wait_for(upstream_task, timeout=5.0)
|
||||
except Exception:
|
||||
upstream_task.cancel()
|
||||
with contextlib.suppress(Exception):
|
||||
await upstream_task
|
||||
|
||||
relay_server.should_exit = True
|
||||
try:
|
||||
await asyncio.wait_for(relay_task, timeout=5.0)
|
||||
except Exception:
|
||||
relay_task.cancel()
|
||||
with contextlib.suppress(Exception):
|
||||
await relay_task
|
||||
|
||||
executor_proc.send_signal(signal.SIGTERM)
|
||||
try:
|
||||
executor_proc.wait(timeout=5)
|
||||
except subprocess.TimeoutExpired:
|
||||
executor_proc.kill()
|
||||
executor_proc.wait()
|
||||
executor_socket.unlink(missing_ok=True)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
success = asyncio.run(run_test())
|
||||
sys.exit(0 if success else 1)
|
||||
@@ -5,9 +5,11 @@ import sys
|
||||
import types
|
||||
from types import SimpleNamespace
|
||||
from typing import Any, cast
|
||||
from unittest.mock import AsyncMock
|
||||
|
||||
import pytest
|
||||
|
||||
from src.config import config as app_config
|
||||
from src.services.provider_keys.codex_usage_parser import (
|
||||
CodexUsageParseError,
|
||||
parse_codex_wham_usage_response,
|
||||
@@ -17,6 +19,7 @@ from src.services.provider_keys.quota_refresh.antigravity_refresher import (
|
||||
)
|
||||
from src.services.provider_keys.quota_refresh.codex_refresher import refresh_codex_key_quota
|
||||
from src.services.provider_keys.quota_refresh.kiro_refresher import refresh_kiro_key_quota
|
||||
from src.services.request.rust_executor_client import RustExecutorSyncResult
|
||||
|
||||
|
||||
class _FakeDB:
|
||||
@@ -149,6 +152,73 @@ async def test_codex_refresher_http_non_200_returns_error(
|
||||
assert result["status_code"] == 503
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_codex_refresher_prefers_rust_executor(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
from src.services.provider_keys.quota_refresh import codex_refresher as module
|
||||
from src.services.request import rust_executor_client as rust_module
|
||||
|
||||
monkeypatch.setattr(app_config, "executor_backend", "rust")
|
||||
|
||||
key = SimpleNamespace(
|
||||
id="k1", name="K1", api_key="enc", auth_type="api_key", auth_config=None, proxy=None
|
||||
)
|
||||
provider = SimpleNamespace(id="provider-1", proxy=None)
|
||||
endpoint = SimpleNamespace(id="endpoint-1")
|
||||
metadata_updates: dict[str, dict[str, Any]] = {}
|
||||
state_updates: dict[str, dict[str, Any]] = {}
|
||||
captured: dict[str, Any] = {}
|
||||
|
||||
async def _fake_auth_info(_endpoint: Any, _key: Any) -> Any:
|
||||
return None
|
||||
|
||||
_install_module(
|
||||
monkeypatch,
|
||||
"src.services.proxy_node.resolver",
|
||||
{
|
||||
"resolve_effective_proxy": lambda provider_proxy, key_proxy: {"enabled": True, "url": "http://proxy.test:8080"},
|
||||
"build_proxy_client_kwargs": lambda proxy, timeout: {"timeout": timeout},
|
||||
"get_system_proxy_config_async": AsyncMock(return_value=None),
|
||||
"resolve_delegate_config_async": AsyncMock(return_value=None),
|
||||
"build_proxy_url_async": AsyncMock(return_value="http://proxy.test:8080"),
|
||||
"resolve_proxy_info_async": AsyncMock(return_value={"mode": "http", "label": "proxy"}),
|
||||
},
|
||||
)
|
||||
monkeypatch.setattr(module, "get_provider_auth", _fake_auth_info)
|
||||
monkeypatch.setattr(module.crypto_service, "decrypt", lambda _v: "sk-test")
|
||||
monkeypatch.setattr(
|
||||
module, "parse_codex_wham_usage_response", lambda _data: {"used_percent": 12.5}
|
||||
)
|
||||
|
||||
async def _fake_execute_sync_json(self: object, plan: Any) -> RustExecutorSyncResult:
|
||||
captured["plan"] = plan
|
||||
return RustExecutorSyncResult(
|
||||
status_code=200,
|
||||
response_json={"ok": True},
|
||||
headers={"content-type": "application/json"},
|
||||
)
|
||||
|
||||
monkeypatch.setattr(rust_module.RustExecutorClient, "execute_sync_json", _fake_execute_sync_json)
|
||||
|
||||
result = await refresh_codex_key_quota(
|
||||
db=cast(Any, _FakeDB()),
|
||||
provider=cast(Any, provider),
|
||||
key=cast(Any, key),
|
||||
endpoint=cast(Any, endpoint),
|
||||
codex_wham_usage_url="https://example.test/wham/usage",
|
||||
metadata_updates=metadata_updates,
|
||||
state_updates=state_updates,
|
||||
)
|
||||
|
||||
assert result["status"] == "success"
|
||||
assert captured["plan"].method == "GET"
|
||||
assert captured["plan"].url == "https://example.test/wham/usage"
|
||||
assert captured["plan"].proxy is not None
|
||||
assert captured["plan"].proxy.url == "http://proxy.test:8080"
|
||||
assert metadata_updates["k1"]["codex"]["used_percent"] == 12.5
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_codex_refresher_http_401_marks_auth_invalid_without_disabling_key(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
|
||||
@@ -1,11 +1,13 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any
|
||||
from unittest.mock import AsyncMock
|
||||
|
||||
import pytest
|
||||
|
||||
from src.services.provider_ops.service import ProviderOpsService
|
||||
from src.services.provider_ops.types import ConnectorAuthType
|
||||
from src.services.provider_ops.types import ConnectorAuthType, ProviderActionType
|
||||
from src.services.request.rust_executor_client import RustExecutorSyncResult
|
||||
|
||||
|
||||
class _FakeDB:
|
||||
@@ -51,6 +53,43 @@ class _FakeRegistry:
|
||||
return self._architecture
|
||||
|
||||
|
||||
class _SuccessResult:
|
||||
def __init__(self) -> None:
|
||||
self.success = True
|
||||
self.quota = 200.0
|
||||
self.extra = {"window": "day"}
|
||||
|
||||
def to_dict(self) -> dict[str, Any]:
|
||||
return {"success": True, "quota": self.quota}
|
||||
|
||||
|
||||
class _SuccessArchitecture:
|
||||
default_action_configs = {ProviderActionType.QUERY_BALANCE: {"quota_divisor": 100}}
|
||||
|
||||
def get_verify_endpoint(self) -> str:
|
||||
return "/verify"
|
||||
|
||||
async def prepare_verify_config(
|
||||
self,
|
||||
_base_url: str,
|
||||
_config: dict[str, Any],
|
||||
_credentials: dict[str, Any],
|
||||
) -> dict[str, Any]:
|
||||
return {}
|
||||
|
||||
def build_verify_headers(
|
||||
self,
|
||||
_config: dict[str, Any],
|
||||
_credentials: dict[str, Any],
|
||||
) -> dict[str, str]:
|
||||
return {"authorization": "Bearer test"}
|
||||
|
||||
def parse_verify_response(self, status_code: int, data: dict[str, Any]) -> _SuccessResult:
|
||||
assert status_code == 200
|
||||
assert data == {"ok": True}
|
||||
return _SuccessResult()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_verify_auth_returns_failure_when_prepare_verify_config_raises_value_error(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
@@ -72,3 +111,57 @@ async def test_verify_auth_returns_failure_when_prepare_verify_config_raises_val
|
||||
)
|
||||
|
||||
assert result == {"success": False, "message": "invalid refresh token"}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_verify_auth_prefers_rust_executor(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
from src.services.provider_ops import service as module
|
||||
from src.services.request import rust_executor_client as rust_module
|
||||
|
||||
service = ProviderOpsService(_FakeDB())
|
||||
architecture = _SuccessArchitecture()
|
||||
|
||||
monkeypatch.setattr(
|
||||
module,
|
||||
"get_registry",
|
||||
lambda: _FakeRegistry(architecture),
|
||||
)
|
||||
monkeypatch.setattr(module.config, "executor_backend", "rust")
|
||||
monkeypatch.setattr(
|
||||
"src.services.proxy_node.resolver.resolve_ops_proxy_config_async",
|
||||
AsyncMock(return_value=(None, "node-1")),
|
||||
)
|
||||
|
||||
cache_balance = AsyncMock()
|
||||
monkeypatch.setattr(service, "_cache_balance_from_verify", cache_balance)
|
||||
|
||||
captured: dict[str, Any] = {}
|
||||
|
||||
async def _fake_execute_sync_json(self: object, plan: Any) -> RustExecutorSyncResult:
|
||||
captured["plan"] = plan
|
||||
return RustExecutorSyncResult(
|
||||
status_code=200,
|
||||
response_json={"ok": True},
|
||||
headers={"content-type": "application/json"},
|
||||
)
|
||||
|
||||
monkeypatch.setattr(rust_module.RustExecutorClient, "execute_sync_json", _fake_execute_sync_json)
|
||||
|
||||
result = await service.verify_auth(
|
||||
base_url="https://example.com",
|
||||
architecture_id="sub2api",
|
||||
auth_type=ConnectorAuthType.SESSION_LOGIN,
|
||||
config={},
|
||||
credentials={"access_token": "token"},
|
||||
provider_id="provider-1",
|
||||
)
|
||||
|
||||
assert result == {"success": True, "quota": 200.0}
|
||||
assert captured["plan"].method == "GET"
|
||||
assert captured["plan"].url == "https://example.com/verify"
|
||||
assert captured["plan"].proxy is not None
|
||||
assert captured["plan"].proxy.mode == "tunnel"
|
||||
assert captured["plan"].proxy.node_id == "node-1"
|
||||
cache_balance.assert_awaited_once_with("provider-1", 2.0, {"window": "day"})
|
||||
|
||||
173
tests/services/test_video_cancel_rust_executor.py
Normal file
173
tests/services/test_video_cancel_rust_executor.py
Normal file
@@ -0,0 +1,173 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import AsyncMock
|
||||
|
||||
import pytest
|
||||
|
||||
import src.services.task.video.cancel as cancel_mod
|
||||
import src.services.request.rust_executor_client as rust_client_mod
|
||||
from src.core.api_format.conversion.internal_video import VideoStatus
|
||||
from src.services.request.rust_executor_client import RustExecutorSyncResult
|
||||
from src.services.task.video.cancel import VideoTaskCancelService
|
||||
|
||||
|
||||
class _Query:
|
||||
def __init__(self, result: object) -> None:
|
||||
self._result = result
|
||||
|
||||
def filter(self, *args: object, **kwargs: object) -> "_Query":
|
||||
del args, kwargs
|
||||
return self
|
||||
|
||||
def first(self) -> object:
|
||||
return self._result
|
||||
|
||||
|
||||
class _FakeDB:
|
||||
def __init__(self, *results: object) -> None:
|
||||
self._results = list(results)
|
||||
self.committed = False
|
||||
|
||||
def query(self, model: object) -> _Query:
|
||||
del model
|
||||
return _Query(self._results.pop(0))
|
||||
|
||||
def commit(self) -> None:
|
||||
self.committed = True
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_video_cancel_service_uses_rust_for_openai_delete(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
endpoint = SimpleNamespace(
|
||||
id="ep-1",
|
||||
provider_id="prov-1",
|
||||
api_family="openai",
|
||||
endpoint_kind="video",
|
||||
base_url="https://api.openai.com",
|
||||
)
|
||||
key = SimpleNamespace(id="key-1", api_key="encrypted")
|
||||
db = _FakeDB(endpoint, key)
|
||||
service = VideoTaskCancelService(db)
|
||||
task = SimpleNamespace(
|
||||
id="task-1",
|
||||
status=VideoStatus.PROCESSING.value,
|
||||
external_task_id="ext-1",
|
||||
endpoint_id="ep-1",
|
||||
key_id="key-1",
|
||||
request_id="req-1",
|
||||
model="sora-2",
|
||||
completed_at=None,
|
||||
updated_at=None,
|
||||
)
|
||||
|
||||
monkeypatch.setattr(cancel_mod.config, "executor_backend", "rust")
|
||||
monkeypatch.setattr(
|
||||
rust_client_mod.RustExecutorClient,
|
||||
"execute_sync_json",
|
||||
AsyncMock(
|
||||
return_value=RustExecutorSyncResult(
|
||||
status_code=204,
|
||||
headers={},
|
||||
response_json=None,
|
||||
response_body_bytes=None,
|
||||
)
|
||||
),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
"src.core.crypto.crypto_service.decrypt",
|
||||
lambda value: "upstream-key",
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
"src.core.api_format.build_upstream_headers_for_endpoint",
|
||||
lambda *args, **kwargs: {"authorization": "Bearer upstream-key"},
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
"src.services.provider.transport.build_provider_url",
|
||||
lambda endpoint, is_stream=False, key=None: "https://api.openai.com/v1/videos",
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
"src.clients.http_client.HTTPClientPool.get_default_client_async",
|
||||
AsyncMock(side_effect=AssertionError("python fallback should not run")),
|
||||
)
|
||||
monkeypatch.setattr(cancel_mod.UsageService, "finalize_void", lambda *args, **kwargs: True)
|
||||
|
||||
result = await service.cancel_task(
|
||||
task=task,
|
||||
task_id="task-1",
|
||||
original_headers={"x-test": "1"},
|
||||
)
|
||||
|
||||
assert result is None
|
||||
assert db.committed is True
|
||||
assert task.status == VideoStatus.CANCELLED.value
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_video_cancel_service_uses_rust_for_gemini_cancel(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
endpoint = SimpleNamespace(
|
||||
id="ep-1",
|
||||
provider_id="prov-1",
|
||||
api_family="gemini",
|
||||
endpoint_kind="video",
|
||||
base_url="https://generativelanguage.googleapis.com",
|
||||
)
|
||||
key = SimpleNamespace(id="key-1", api_key="encrypted")
|
||||
db = _FakeDB(endpoint, key)
|
||||
service = VideoTaskCancelService(db)
|
||||
task = SimpleNamespace(
|
||||
id="task-1",
|
||||
status=VideoStatus.PROCESSING.value,
|
||||
external_task_id="operations/ext-1",
|
||||
endpoint_id="ep-1",
|
||||
key_id="key-1",
|
||||
request_id="req-1",
|
||||
model="veo-3",
|
||||
completed_at=None,
|
||||
updated_at=None,
|
||||
)
|
||||
|
||||
monkeypatch.setattr(cancel_mod.config, "executor_backend", "rust")
|
||||
execute_sync = AsyncMock(
|
||||
return_value=RustExecutorSyncResult(
|
||||
status_code=200,
|
||||
headers={"content-type": "application/json"},
|
||||
response_json={"ok": True},
|
||||
response_body_bytes=None,
|
||||
)
|
||||
)
|
||||
monkeypatch.setattr(rust_client_mod.RustExecutorClient, "execute_sync_json", execute_sync)
|
||||
monkeypatch.setattr(
|
||||
"src.core.crypto.crypto_service.decrypt",
|
||||
lambda value: "upstream-key",
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
"src.core.api_format.build_upstream_headers_for_endpoint",
|
||||
lambda *args, **kwargs: {"x-goog-api-key": "upstream-key"},
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
"src.services.provider.auth.get_provider_auth",
|
||||
AsyncMock(return_value=SimpleNamespace(auth_header="authorization", auth_value="Bearer token")),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
"src.clients.http_client.HTTPClientPool.get_default_client_async",
|
||||
AsyncMock(side_effect=AssertionError("python fallback should not run")),
|
||||
)
|
||||
monkeypatch.setattr(cancel_mod.UsageService, "finalize_void", lambda *args, **kwargs: True)
|
||||
|
||||
result = await service.cancel_task(
|
||||
task=task,
|
||||
task_id="task-1",
|
||||
original_headers={"x-test": "1"},
|
||||
)
|
||||
|
||||
assert result is None
|
||||
assert db.committed is True
|
||||
assert task.status == VideoStatus.CANCELLED.value
|
||||
plan = execute_sync.await_args.args[0]
|
||||
assert plan.method == "POST"
|
||||
assert plan.body.json_body == {}
|
||||
@@ -6,7 +6,8 @@ import pytest
|
||||
|
||||
from src.core.api_format.conversion.internal_video import InternalVideoPollResult, VideoStatus
|
||||
from src.models.database import VideoTask
|
||||
from src.services.task.video.poller_adapter import VideoTaskPollerAdapter
|
||||
from src.services.request.rust_executor_client import RustExecutorSyncResult
|
||||
from src.services.task.video.poller_adapter import VideoPollContext, VideoTaskPollerAdapter
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@@ -78,3 +79,140 @@ async def test_update_task_after_poll_skips_terminal_cancelled_task(
|
||||
|
||||
finalize.assert_not_awaited()
|
||||
session.commit.assert_not_called()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_video_poller_try_rust_payload_passes_proxy_snapshot(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
from src.services.request import rust_executor_client as rust_mod
|
||||
from src.services.task.video import poller_adapter as mod
|
||||
|
||||
adapter = VideoTaskPollerAdapter()
|
||||
monkeypatch.setattr(mod.config, "executor_backend", "rust")
|
||||
|
||||
proxy_snapshot = object()
|
||||
captured: dict[str, object] = {}
|
||||
|
||||
async def _fake_execute_sync_json(self: object, plan: object) -> RustExecutorSyncResult:
|
||||
captured["plan"] = plan
|
||||
return RustExecutorSyncResult(status_code=200, response_json={"id": "op_1", "done": False})
|
||||
|
||||
monkeypatch.setattr(rust_mod.RustExecutorClient, "execute_sync_json", _fake_execute_sync_json)
|
||||
|
||||
ctx = VideoPollContext(
|
||||
task_id="task-1",
|
||||
external_task_id="vid_1",
|
||||
provider_api_format="openai:video",
|
||||
base_url="https://api.openai.com",
|
||||
upstream_key="upstream",
|
||||
headers={"authorization": "Bearer test"},
|
||||
poll_count=0,
|
||||
retry_count=0,
|
||||
poll_interval_seconds=15,
|
||||
max_poll_count=10,
|
||||
current_status=VideoStatus.PROCESSING.value,
|
||||
proxy_snapshot=proxy_snapshot,
|
||||
)
|
||||
|
||||
payload = await adapter._try_rust_poll_payload(
|
||||
ctx=ctx,
|
||||
url="https://api.openai.com/v1/videos/vid_1",
|
||||
)
|
||||
|
||||
assert payload == {"id": "op_1", "done": False}
|
||||
assert getattr(captured["plan"], "method") == "GET"
|
||||
assert getattr(captured["plan"], "proxy") is proxy_snapshot
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_video_poller_openai_poll_prefers_rust_payload(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
from src.services.task.video import poller_adapter as mod
|
||||
|
||||
adapter = VideoTaskPollerAdapter()
|
||||
|
||||
rust_poll = AsyncMock(return_value={"id": "vid_1", "status": "processing"})
|
||||
monkeypatch.setattr(adapter, "_try_rust_poll_payload", rust_poll)
|
||||
|
||||
normalized = InternalVideoPollResult(status=VideoStatus.PROCESSING, progress_percent=42)
|
||||
normalizer = MagicMock(return_value=normalized)
|
||||
monkeypatch.setattr(adapter._openai_normalizer, "video_poll_to_internal", normalizer)
|
||||
|
||||
get_upstream_client = AsyncMock(side_effect=AssertionError("python upstream client should not be used"))
|
||||
monkeypatch.setattr(mod.HTTPClientPool, "get_upstream_client", get_upstream_client)
|
||||
|
||||
ctx = VideoPollContext(
|
||||
task_id="task-1",
|
||||
external_task_id="vid_1",
|
||||
provider_api_format="openai:video",
|
||||
base_url="https://api.openai.com",
|
||||
upstream_key="upstream",
|
||||
headers={"authorization": "Bearer test"},
|
||||
poll_count=0,
|
||||
retry_count=0,
|
||||
poll_interval_seconds=15,
|
||||
max_poll_count=10,
|
||||
current_status=VideoStatus.PROCESSING.value,
|
||||
)
|
||||
|
||||
result = await adapter._poll_openai_with_context(ctx)
|
||||
|
||||
assert result is normalized
|
||||
normalizer.assert_called_once_with({"id": "vid_1", "status": "processing"})
|
||||
get_upstream_client.assert_not_awaited()
|
||||
rust_poll.assert_awaited_once()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_video_poller_openai_poll_fallback_uses_transport_aware_client(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
from src.services.task.video import poller_adapter as mod
|
||||
|
||||
adapter = VideoTaskPollerAdapter()
|
||||
|
||||
monkeypatch.setattr(adapter, "_try_rust_poll_payload", AsyncMock(return_value=None))
|
||||
|
||||
response = MagicMock()
|
||||
response.status_code = 200
|
||||
response.json.return_value = {"id": "vid_2", "status": "processing"}
|
||||
|
||||
client = MagicMock()
|
||||
client.get = AsyncMock(return_value=response)
|
||||
|
||||
get_upstream_client = AsyncMock(return_value=client)
|
||||
monkeypatch.setattr(mod.HTTPClientPool, "get_upstream_client", get_upstream_client)
|
||||
|
||||
normalized = InternalVideoPollResult(status=VideoStatus.PROCESSING, progress_percent=7)
|
||||
normalizer = MagicMock(return_value=normalized)
|
||||
monkeypatch.setattr(adapter._openai_normalizer, "video_poll_to_internal", normalizer)
|
||||
|
||||
ctx = VideoPollContext(
|
||||
task_id="task-2",
|
||||
external_task_id="vid_2",
|
||||
provider_api_format="openai:video",
|
||||
base_url="https://api.openai.com",
|
||||
upstream_key="upstream",
|
||||
headers={"authorization": "Bearer test"},
|
||||
poll_count=0,
|
||||
retry_count=0,
|
||||
poll_interval_seconds=15,
|
||||
max_poll_count=10,
|
||||
current_status=VideoStatus.PROCESSING.value,
|
||||
proxy_config={"enabled": True, "url": "http://proxy.test:8080"},
|
||||
delegate_config={"node_id": "node-1", "tunnel": True},
|
||||
)
|
||||
|
||||
result = await adapter._poll_openai_with_context(ctx)
|
||||
|
||||
assert result is normalized
|
||||
get_upstream_client.assert_awaited_once_with(
|
||||
{"node_id": "node-1", "tunnel": True},
|
||||
proxy_config={"enabled": True, "url": "http://proxy.test:8080"},
|
||||
)
|
||||
client.get.assert_awaited_once_with(
|
||||
"https://api.openai.com/v1/videos/vid_2",
|
||||
headers={"authorization": "Bearer test"},
|
||||
)
|
||||
|
||||
43
tests/unit/test_executor_config.py
Normal file
43
tests/unit/test_executor_config.py
Normal file
@@ -0,0 +1,43 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import pytest
|
||||
|
||||
from src.config.settings import Config
|
||||
|
||||
|
||||
def test_executor_config_defaults_to_rust_and_unix_socket(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
for key in (
|
||||
"EXECUTOR_BACKEND",
|
||||
"EXECUTOR_TRANSPORT",
|
||||
"EXECUTOR_SOCKET_PATH",
|
||||
"EXECUTOR_BASE_URL",
|
||||
"EXECUTOR_REQUEST_TIMEOUT",
|
||||
"HTTP_REQUEST_TIMEOUT",
|
||||
):
|
||||
monkeypatch.delenv(key, raising=False)
|
||||
|
||||
cfg = Config()
|
||||
|
||||
assert cfg.executor_backend == "rust"
|
||||
assert cfg.executor_transport == "unix_socket"
|
||||
assert cfg.executor_socket_path == "/tmp/aether-executor.sock"
|
||||
assert cfg.executor_base_url == "http://127.0.0.1:5219"
|
||||
assert cfg.executor_request_timeout == cfg.http_request_timeout
|
||||
|
||||
|
||||
def test_executor_config_accepts_env_override(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
monkeypatch.setenv("EXECUTOR_BACKEND", "RUST")
|
||||
monkeypatch.setenv("EXECUTOR_TRANSPORT", "TCP")
|
||||
monkeypatch.setenv("EXECUTOR_SOCKET_PATH", "/var/run/aether.sock")
|
||||
monkeypatch.setenv("EXECUTOR_BASE_URL", "http://127.0.0.1:9311")
|
||||
monkeypatch.setenv("EXECUTOR_REQUEST_TIMEOUT", "12.5")
|
||||
|
||||
cfg = Config()
|
||||
|
||||
assert cfg.executor_backend == "rust"
|
||||
assert cfg.executor_transport == "tcp"
|
||||
assert cfg.executor_socket_path == "/var/run/aether.sock"
|
||||
assert cfg.executor_base_url == "http://127.0.0.1:9311"
|
||||
assert cfg.executor_request_timeout == 12.5
|
||||
344
tests/unit/test_executor_plan.py
Normal file
344
tests/unit/test_executor_plan.py
Normal file
@@ -0,0 +1,344 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import base64
|
||||
|
||||
from src.services.request.executor_plan import PreparedExecutionPlan
|
||||
from src.services.request.executor_plan import (
|
||||
ExecutionPlan,
|
||||
ExecutionPlanBody,
|
||||
ExecutionPlanTimeouts,
|
||||
ExecutionProxySnapshot,
|
||||
build_execution_plan_body,
|
||||
)
|
||||
|
||||
|
||||
def test_execution_plan_to_payload_drops_none_fields() -> None:
|
||||
plan = ExecutionPlan(
|
||||
request_id="req-1",
|
||||
candidate_id=None,
|
||||
provider_name="openai",
|
||||
provider_id="prov-1",
|
||||
endpoint_id="ep-1",
|
||||
key_id="key-1",
|
||||
method="POST",
|
||||
url="https://example.com/v1/chat/completions",
|
||||
headers={"content-type": "application/json"},
|
||||
body=ExecutionPlanBody(json_body={"model": "gpt-4.1"}),
|
||||
stream=False,
|
||||
provider_api_format="openai:chat",
|
||||
client_api_format="openai:chat",
|
||||
model_name="gpt-4.1",
|
||||
proxy=ExecutionProxySnapshot(
|
||||
enabled=True,
|
||||
mode="direct",
|
||||
label="no-proxy",
|
||||
),
|
||||
timeouts=ExecutionPlanTimeouts(
|
||||
connect_ms=10_000,
|
||||
total_ms=300_000,
|
||||
),
|
||||
)
|
||||
|
||||
payload = plan.to_payload()
|
||||
|
||||
assert "candidate_id" not in payload
|
||||
assert payload["body"] == {"json_body": {"model": "gpt-4.1"}}
|
||||
assert payload["proxy"] == {
|
||||
"enabled": True,
|
||||
"mode": "direct",
|
||||
"label": "no-proxy",
|
||||
}
|
||||
assert payload["timeouts"] == {
|
||||
"connect_ms": 10_000,
|
||||
"total_ms": 300_000,
|
||||
}
|
||||
|
||||
|
||||
def test_prepared_execution_plan_remote_eligible_for_non_stream_json() -> None:
|
||||
prepared = PreparedExecutionPlan(
|
||||
contract=ExecutionPlan(
|
||||
request_id="req-1",
|
||||
candidate_id=None,
|
||||
provider_name="openai",
|
||||
provider_id="prov-1",
|
||||
endpoint_id="ep-1",
|
||||
key_id="key-1",
|
||||
method="POST",
|
||||
url="https://example.com/v1/chat/completions",
|
||||
headers={"content-type": "application/json"},
|
||||
body=ExecutionPlanBody(json_body={"model": "gpt-4.1"}),
|
||||
stream=False,
|
||||
provider_api_format="openai:chat",
|
||||
client_api_format="openai:chat",
|
||||
model_name="gpt-4.1",
|
||||
),
|
||||
payload={"model": "gpt-4.1"},
|
||||
headers={"content-type": "application/json"},
|
||||
upstream_is_stream=False,
|
||||
needs_conversion=False,
|
||||
provider_type="openai",
|
||||
request_timeout=300.0,
|
||||
)
|
||||
|
||||
assert prepared.remote_eligible is True
|
||||
|
||||
|
||||
def test_prepared_execution_plan_remote_eligible_allows_upstream_stream() -> None:
|
||||
prepared = PreparedExecutionPlan(
|
||||
contract=ExecutionPlan(
|
||||
request_id="req-1",
|
||||
candidate_id=None,
|
||||
provider_name="openai",
|
||||
provider_id="prov-1",
|
||||
endpoint_id="ep-1",
|
||||
key_id="key-1",
|
||||
method="POST",
|
||||
url="https://example.com/v1/chat/completions",
|
||||
headers={"content-type": "application/json"},
|
||||
body=ExecutionPlanBody(json_body={"model": "gpt-4.1", "stream": True}),
|
||||
stream=True,
|
||||
provider_api_format="openai:chat",
|
||||
client_api_format="openai:chat",
|
||||
model_name="gpt-4.1",
|
||||
),
|
||||
payload={"model": "gpt-4.1", "stream": True},
|
||||
headers={"content-type": "application/json"},
|
||||
upstream_is_stream=True,
|
||||
needs_conversion=False,
|
||||
provider_type="openai",
|
||||
request_timeout=300.0,
|
||||
)
|
||||
|
||||
assert prepared.remote_eligible is True
|
||||
|
||||
|
||||
def test_prepared_execution_plan_remote_eligible_allows_tunnel_delegate_proxy() -> None:
|
||||
prepared = PreparedExecutionPlan(
|
||||
contract=ExecutionPlan(
|
||||
request_id="req-1",
|
||||
candidate_id=None,
|
||||
provider_name="openai",
|
||||
provider_id="prov-1",
|
||||
endpoint_id="ep-1",
|
||||
key_id="key-1",
|
||||
method="POST",
|
||||
url="https://example.com/v1/chat/completions",
|
||||
headers={"content-type": "application/json"},
|
||||
body=ExecutionPlanBody(json_body={"model": "gpt-4.1"}),
|
||||
stream=False,
|
||||
provider_api_format="gemini:chat",
|
||||
client_api_format="openai:chat",
|
||||
model_name="gpt-4.1",
|
||||
proxy=ExecutionProxySnapshot(
|
||||
enabled=True,
|
||||
mode="tunnel",
|
||||
node_id="node-1",
|
||||
label="relay-node",
|
||||
),
|
||||
),
|
||||
payload={"model": "gpt-4.1"},
|
||||
headers={"content-type": "application/json"},
|
||||
upstream_is_stream=False,
|
||||
needs_conversion=True,
|
||||
provider_type="gemini",
|
||||
request_timeout=300.0,
|
||||
delegate_config={"tunnel": True, "node_id": "node-1"},
|
||||
proxy_config={"node_id": "node-1"},
|
||||
)
|
||||
|
||||
assert prepared.remote_eligible is True
|
||||
|
||||
|
||||
def test_prepared_execution_plan_remote_eligible_allows_url_proxy_without_delegate() -> None:
|
||||
prepared = PreparedExecutionPlan(
|
||||
contract=ExecutionPlan(
|
||||
request_id="req-1",
|
||||
candidate_id=None,
|
||||
provider_name="openai",
|
||||
provider_id="prov-1",
|
||||
endpoint_id="ep-1",
|
||||
key_id="key-1",
|
||||
method="POST",
|
||||
url="https://example.com/v1/chat/completions",
|
||||
headers={"content-type": "application/json"},
|
||||
body=ExecutionPlanBody(json_body={"model": "gpt-4.1"}),
|
||||
stream=False,
|
||||
provider_api_format="openai:chat",
|
||||
client_api_format="openai:chat",
|
||||
model_name="gpt-4.1",
|
||||
proxy=ExecutionProxySnapshot(
|
||||
enabled=True,
|
||||
mode="http",
|
||||
label="proxy.internal",
|
||||
url="http://proxy.internal:8080",
|
||||
),
|
||||
),
|
||||
payload={"model": "gpt-4.1"},
|
||||
headers={"content-type": "application/json"},
|
||||
upstream_is_stream=False,
|
||||
needs_conversion=False,
|
||||
provider_type="openai",
|
||||
request_timeout=300.0,
|
||||
proxy_config={"url": "http://proxy.internal:8080"},
|
||||
)
|
||||
|
||||
assert prepared.remote_eligible is True
|
||||
|
||||
|
||||
def test_prepared_execution_plan_remote_eligible_allows_conversion_and_envelope() -> None:
|
||||
prepared = PreparedExecutionPlan(
|
||||
contract=ExecutionPlan(
|
||||
request_id="req-1",
|
||||
candidate_id=None,
|
||||
provider_name="gemini",
|
||||
provider_id="prov-1",
|
||||
endpoint_id="ep-1",
|
||||
key_id="key-1",
|
||||
method="POST",
|
||||
url="https://example.com/v1/chat/completions",
|
||||
headers={"content-type": "application/json"},
|
||||
body=ExecutionPlanBody(json_body={"model": "gpt-4.1"}),
|
||||
stream=False,
|
||||
provider_api_format="gemini:chat",
|
||||
client_api_format="openai:chat",
|
||||
model_name="gpt-4.1",
|
||||
),
|
||||
payload={"model": "gpt-4.1"},
|
||||
headers={"content-type": "application/json"},
|
||||
upstream_is_stream=False,
|
||||
needs_conversion=True,
|
||||
provider_type="gemini",
|
||||
request_timeout=300.0,
|
||||
envelope=object(),
|
||||
)
|
||||
|
||||
assert prepared.remote_eligible is True
|
||||
|
||||
|
||||
def test_prepared_execution_plan_remote_eligible_allows_tls_profile() -> None:
|
||||
prepared = PreparedExecutionPlan(
|
||||
contract=ExecutionPlan(
|
||||
request_id="req-1",
|
||||
candidate_id=None,
|
||||
provider_name="claude",
|
||||
provider_id="prov-1",
|
||||
endpoint_id="ep-1",
|
||||
key_id="key-1",
|
||||
method="POST",
|
||||
url="https://example.com/v1/messages",
|
||||
headers={"content-type": "application/json"},
|
||||
body=ExecutionPlanBody(json_body={"model": "claude-3.7-sonnet"}),
|
||||
stream=False,
|
||||
provider_api_format="claude:chat",
|
||||
client_api_format="claude:chat",
|
||||
model_name="claude-3.7-sonnet",
|
||||
tls_profile="claude_code_nodejs",
|
||||
),
|
||||
payload={"model": "claude-3.7-sonnet"},
|
||||
headers={"content-type": "application/json"},
|
||||
upstream_is_stream=False,
|
||||
needs_conversion=False,
|
||||
provider_type="claude_code",
|
||||
request_timeout=300.0,
|
||||
)
|
||||
|
||||
assert prepared.remote_eligible is True
|
||||
|
||||
|
||||
def test_build_execution_plan_body_encodes_raw_bytes_payload() -> None:
|
||||
body = build_execution_plan_body(
|
||||
b"raw-payload",
|
||||
content_type="text/plain",
|
||||
)
|
||||
|
||||
assert body.json_body is None
|
||||
assert base64.b64decode(body.body_bytes_b64 or "") == b"raw-payload"
|
||||
|
||||
|
||||
def test_prepared_execution_plan_remote_eligible_allows_gzip_json_body() -> None:
|
||||
prepared = PreparedExecutionPlan(
|
||||
contract=ExecutionPlan(
|
||||
request_id="req-1",
|
||||
candidate_id=None,
|
||||
provider_name="openai",
|
||||
provider_id="prov-1",
|
||||
endpoint_id="ep-1",
|
||||
key_id="key-1",
|
||||
method="POST",
|
||||
url="https://example.com/v1/chat/completions",
|
||||
headers={"content-type": "application/json"},
|
||||
body=ExecutionPlanBody(json_body={"model": "gpt-4.1"}),
|
||||
stream=False,
|
||||
provider_api_format="openai:chat",
|
||||
client_api_format="openai:chat",
|
||||
model_name="gpt-4.1",
|
||||
content_encoding="gzip",
|
||||
),
|
||||
payload={"model": "gpt-4.1"},
|
||||
headers={"content-type": "application/json"},
|
||||
upstream_is_stream=False,
|
||||
needs_conversion=False,
|
||||
provider_type="openai",
|
||||
request_timeout=300.0,
|
||||
)
|
||||
|
||||
assert prepared.remote_eligible is True
|
||||
|
||||
|
||||
def test_prepared_execution_plan_remote_eligible_allows_raw_body_with_passthrough_encoding() -> None:
|
||||
prepared = PreparedExecutionPlan(
|
||||
contract=ExecutionPlan(
|
||||
request_id="req-1",
|
||||
candidate_id=None,
|
||||
provider_name="openai",
|
||||
provider_id="prov-1",
|
||||
endpoint_id="ep-1",
|
||||
key_id="key-1",
|
||||
method="POST",
|
||||
url="https://example.com/v1/chat/completions",
|
||||
headers={"content-type": "text/plain"},
|
||||
body=ExecutionPlanBody(body_bytes_b64=base64.b64encode(b"raw").decode("ascii")),
|
||||
stream=False,
|
||||
provider_api_format="openai:chat",
|
||||
client_api_format="openai:chat",
|
||||
model_name="gpt-4.1",
|
||||
content_encoding="br",
|
||||
),
|
||||
payload={"model": "gpt-4.1"},
|
||||
headers={"content-type": "text/plain"},
|
||||
upstream_is_stream=False,
|
||||
needs_conversion=False,
|
||||
provider_type="openai",
|
||||
request_timeout=300.0,
|
||||
)
|
||||
|
||||
assert prepared.remote_eligible is True
|
||||
|
||||
|
||||
def test_prepared_execution_plan_remote_eligible_allows_empty_get_body() -> None:
|
||||
prepared = PreparedExecutionPlan(
|
||||
contract=ExecutionPlan(
|
||||
request_id="req-1",
|
||||
candidate_id=None,
|
||||
provider_name="openai",
|
||||
provider_id="prov-1",
|
||||
endpoint_id="ep-1",
|
||||
key_id="key-1",
|
||||
method="GET",
|
||||
url="https://example.com/v1/videos/video-1/content",
|
||||
headers={},
|
||||
body=ExecutionPlanBody(),
|
||||
stream=True,
|
||||
provider_api_format="openai:video",
|
||||
client_api_format="openai:video",
|
||||
model_name="sora-2",
|
||||
),
|
||||
payload={},
|
||||
headers={},
|
||||
upstream_is_stream=True,
|
||||
needs_conversion=False,
|
||||
provider_type="openai",
|
||||
request_timeout=300.0,
|
||||
)
|
||||
|
||||
assert prepared.remote_eligible is True
|
||||
Reference in New Issue
Block a user