mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-02 17:30:23 +08:00
refactor: 拆分 gateway 单体为独立 crate,新增 systemd 部署方案
将 gateway 内部的 model-fetch、provider-transport、scheduler-core、 usage-runtime、video-tasks-core 模块提取为独立 crate;重构 gateway 内部模块结构(state/router/cache/data/query 等);移除大量遗留模块 文件;新增 systemd 二进制部署骨架及相关文档;更新前端 usage 相关 API 和组件。
This commit is contained in:
24
apps/aether-gateway/src/tests/control/admin.rs
Normal file
24
apps/aether-gateway/src/tests/control/admin.rs
Normal file
@@ -0,0 +1,24 @@
|
||||
mod adaptive;
|
||||
mod api_keys;
|
||||
mod billing;
|
||||
mod endpoints;
|
||||
mod gemini_files;
|
||||
mod health_access;
|
||||
mod ldap;
|
||||
mod models;
|
||||
mod monitoring;
|
||||
mod oauth;
|
||||
mod payments;
|
||||
mod pool;
|
||||
mod provider_ops;
|
||||
mod provider_query;
|
||||
mod provider_strategy;
|
||||
mod providers;
|
||||
mod proxy_nodes;
|
||||
mod security;
|
||||
mod stats;
|
||||
mod system;
|
||||
mod usage;
|
||||
mod users;
|
||||
mod video_tasks;
|
||||
mod wallets;
|
||||
449
apps/aether-gateway/src/tests/control/admin/adaptive.rs
Normal file
449
apps/aether-gateway/src/tests/control/admin/adaptive.rs
Normal file
@@ -0,0 +1,449 @@
|
||||
use std::sync::{Arc, Mutex};
|
||||
|
||||
use aether_data::repository::provider_catalog::{
|
||||
InMemoryProviderCatalogReadRepository, ProviderCatalogReadRepository, StoredProviderCatalogKey,
|
||||
};
|
||||
use axum::body::Body;
|
||||
use axum::routing::any;
|
||||
use axum::{extract::Request, Router};
|
||||
use http::StatusCode;
|
||||
use serde_json::json;
|
||||
|
||||
use super::super::{build_router_with_state, sample_key, sample_provider, start_server, AppState};
|
||||
use crate::constants::{
|
||||
GATEWAY_HEADER, TRUSTED_ADMIN_SESSION_ID_HEADER, TRUSTED_ADMIN_USER_ID_HEADER,
|
||||
TRUSTED_ADMIN_USER_ROLE_HEADER,
|
||||
};
|
||||
use crate::data::GatewayDataState;
|
||||
|
||||
fn adaptive_repository_with_keys(
|
||||
keys: Vec<StoredProviderCatalogKey>,
|
||||
) -> Arc<InMemoryProviderCatalogReadRepository> {
|
||||
Arc::new(InMemoryProviderCatalogReadRepository::seed(
|
||||
vec![sample_provider("provider-openai", "openai", 10)],
|
||||
vec![],
|
||||
keys,
|
||||
))
|
||||
}
|
||||
|
||||
async fn assert_adaptive_route_returns_local_503(
|
||||
data_state: crate::data::GatewayDataState,
|
||||
method: http::Method,
|
||||
path: &str,
|
||||
body: Option<serde_json::Value>,
|
||||
expected_message: &str,
|
||||
) {
|
||||
let upstream_hits = Arc::new(Mutex::new(0usize));
|
||||
let upstream_hits_clone = Arc::clone(&upstream_hits);
|
||||
let upstream = Router::new().route(
|
||||
"/{*path}",
|
||||
any(move |_request: Request| {
|
||||
let upstream_hits_inner = Arc::clone(&upstream_hits_clone);
|
||||
async move {
|
||||
*upstream_hits_inner.lock().expect("mutex should lock") += 1;
|
||||
(StatusCode::OK, Body::from("unexpected upstream hit"))
|
||||
}
|
||||
}),
|
||||
);
|
||||
|
||||
let (upstream_url, upstream_handle) = start_server(upstream).await;
|
||||
let gateway = build_router_with_state(
|
||||
AppState::new()
|
||||
.expect("gateway should build")
|
||||
.with_data_state_for_tests(data_state),
|
||||
);
|
||||
let (gateway_url, gateway_handle) = start_server(gateway).await;
|
||||
|
||||
let mut request = reqwest::Client::new()
|
||||
.request(method, format!("{gateway_url}{path}"))
|
||||
.header(crate::constants::GATEWAY_HEADER, "rust-phase3b")
|
||||
.header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
|
||||
.header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
|
||||
.header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123");
|
||||
if let Some(body) = body {
|
||||
request = request.json(&body);
|
||||
}
|
||||
let response = request.send().await.expect("request should succeed");
|
||||
|
||||
assert_eq!(response.status(), StatusCode::SERVICE_UNAVAILABLE);
|
||||
let payload: serde_json::Value = response.json().await.expect("json body should parse");
|
||||
assert_eq!(payload["error"]["message"], expected_message);
|
||||
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
|
||||
|
||||
gateway_handle.abort();
|
||||
upstream_handle.abort();
|
||||
}
|
||||
|
||||
const ADAPTIVE_DATA_UNAVAILABLE_MESSAGE: &str = "Admin adaptive data unavailable";
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_handles_admin_adaptive_keys_locally_with_trusted_admin_principal() {
|
||||
let upstream_hits = Arc::new(Mutex::new(0usize));
|
||||
let upstream_hits_clone = Arc::clone(&upstream_hits);
|
||||
let upstream = Router::new().route(
|
||||
"/{*path}",
|
||||
any(move |_request: Request| {
|
||||
let upstream_hits_inner = Arc::clone(&upstream_hits_clone);
|
||||
async move {
|
||||
*upstream_hits_inner.lock().expect("mutex should lock") += 1;
|
||||
(StatusCode::OK, Body::from("unexpected upstream hit"))
|
||||
}
|
||||
}),
|
||||
);
|
||||
|
||||
let mut adaptive_key = sample_key("key-adaptive", "provider-openai", "openai:chat", "sk-test");
|
||||
adaptive_key.rpm_limit = None;
|
||||
adaptive_key.learned_rpm_limit = Some(17);
|
||||
adaptive_key.concurrent_429_count = Some(2);
|
||||
adaptive_key.rpm_429_count = Some(1);
|
||||
|
||||
let mut fixed_key = sample_key("key-fixed", "provider-openai", "openai:chat", "sk-fixed");
|
||||
fixed_key.rpm_limit = Some(9);
|
||||
fixed_key.learned_rpm_limit = Some(11);
|
||||
|
||||
let repository = adaptive_repository_with_keys(vec![adaptive_key, fixed_key]);
|
||||
|
||||
let (upstream_url, upstream_handle) = start_server(upstream).await;
|
||||
let gateway = build_router_with_state(
|
||||
AppState::new()
|
||||
.expect("gateway should build")
|
||||
.with_data_state_for_tests(
|
||||
GatewayDataState::with_provider_catalog_repository_for_tests(repository),
|
||||
),
|
||||
);
|
||||
let (gateway_url, gateway_handle) = start_server(gateway).await;
|
||||
|
||||
let response = reqwest::Client::new()
|
||||
.get(format!("{gateway_url}/api/admin/adaptive/keys"))
|
||||
.header(crate::constants::GATEWAY_HEADER, "rust-phase3b")
|
||||
.header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
|
||||
.header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
|
||||
.header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
|
||||
.send()
|
||||
.await
|
||||
.expect("request should succeed");
|
||||
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
let payload: serde_json::Value = response.json().await.expect("json body should parse");
|
||||
let items = payload.as_array().expect("items should be an array");
|
||||
assert_eq!(items.len(), 1);
|
||||
assert_eq!(items[0]["id"], "key-adaptive");
|
||||
assert_eq!(items[0]["is_adaptive"], true);
|
||||
assert_eq!(items[0]["effective_limit"], 17);
|
||||
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
|
||||
|
||||
gateway_handle.abort();
|
||||
upstream_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_handles_admin_adaptive_summary_locally_with_trusted_admin_principal() {
|
||||
let upstream_hits = Arc::new(Mutex::new(0usize));
|
||||
let upstream_hits_clone = Arc::clone(&upstream_hits);
|
||||
let upstream = Router::new().route(
|
||||
"/{*path}",
|
||||
any(move |_request: Request| {
|
||||
let upstream_hits_inner = Arc::clone(&upstream_hits_clone);
|
||||
async move {
|
||||
*upstream_hits_inner.lock().expect("mutex should lock") += 1;
|
||||
(StatusCode::OK, Body::from("unexpected upstream hit"))
|
||||
}
|
||||
}),
|
||||
);
|
||||
|
||||
let mut first_key = sample_key("key-a", "provider-openai", "openai:chat", "sk-a");
|
||||
first_key.rpm_limit = None;
|
||||
first_key.concurrent_429_count = Some(2);
|
||||
first_key.rpm_429_count = Some(3);
|
||||
first_key.adjustment_history = Some(json!([
|
||||
{"timestamp": "2026-03-28T15:00:00Z", "delta": 2},
|
||||
{"timestamp": "2026-03-28T14:00:00Z", "delta": -1}
|
||||
]));
|
||||
|
||||
let mut second_key = sample_key("key-b", "provider-openai", "openai:chat", "sk-b");
|
||||
second_key.rpm_limit = None;
|
||||
second_key.concurrent_429_count = Some(1);
|
||||
second_key.rpm_429_count = Some(1);
|
||||
second_key.adjustment_history = Some(json!([
|
||||
{"timestamp": "2026-03-28T16:00:00Z", "delta": 1}
|
||||
]));
|
||||
|
||||
let repository = adaptive_repository_with_keys(vec![first_key, second_key]);
|
||||
|
||||
let (upstream_url, upstream_handle) = start_server(upstream).await;
|
||||
let gateway = build_router_with_state(
|
||||
AppState::new()
|
||||
.expect("gateway should build")
|
||||
.with_data_state_for_tests(
|
||||
GatewayDataState::with_provider_catalog_repository_for_tests(repository),
|
||||
),
|
||||
);
|
||||
let (gateway_url, gateway_handle) = start_server(gateway).await;
|
||||
|
||||
let response = reqwest::Client::new()
|
||||
.get(format!("{gateway_url}/api/admin/adaptive/summary"))
|
||||
.header(crate::constants::GATEWAY_HEADER, "rust-phase3b")
|
||||
.header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
|
||||
.header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
|
||||
.header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
|
||||
.send()
|
||||
.await
|
||||
.expect("request should succeed");
|
||||
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
let payload: serde_json::Value = response.json().await.expect("json body should parse");
|
||||
assert_eq!(payload["total_adaptive_keys"], 2);
|
||||
assert_eq!(payload["total_concurrent_429_errors"], 3);
|
||||
assert_eq!(payload["total_rpm_429_errors"], 4);
|
||||
assert_eq!(payload["total_adjustments"], 3);
|
||||
assert_eq!(payload["recent_adjustments"][0]["key_id"], "key-b");
|
||||
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
|
||||
|
||||
gateway_handle.abort();
|
||||
upstream_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_handles_admin_adaptive_keys_returns_service_unavailable_when_disabled() {
|
||||
assert_adaptive_route_returns_local_503(
|
||||
crate::data::GatewayDataState::disabled(),
|
||||
http::Method::GET,
|
||||
"/api/admin/adaptive/keys",
|
||||
None,
|
||||
ADAPTIVE_DATA_UNAVAILABLE_MESSAGE,
|
||||
)
|
||||
.await;
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_handles_admin_adaptive_mode_returns_service_unavailable_without_writer() {
|
||||
let mut adaptive_key = sample_key("key-adaptive", "provider-openai", "openai:chat", "sk-test");
|
||||
adaptive_key.rpm_limit = None;
|
||||
let repository = adaptive_repository_with_keys(vec![adaptive_key]);
|
||||
|
||||
assert_adaptive_route_returns_local_503(
|
||||
crate::data::GatewayDataState::with_provider_catalog_reader_for_tests(
|
||||
repository,
|
||||
),
|
||||
http::Method::PATCH,
|
||||
"/api/admin/adaptive/keys/key-adaptive/mode",
|
||||
Some(json!({ "enabled": false, "fixed_limit": 20 })),
|
||||
ADAPTIVE_DATA_UNAVAILABLE_MESSAGE,
|
||||
)
|
||||
.await;
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_handles_admin_adaptive_stats_locally_with_trusted_admin_principal() {
|
||||
let upstream_hits = Arc::new(Mutex::new(0usize));
|
||||
let upstream_hits_clone = Arc::clone(&upstream_hits);
|
||||
let upstream = Router::new().route(
|
||||
"/{*path}",
|
||||
any(move |_request: Request| {
|
||||
let upstream_hits_inner = Arc::clone(&upstream_hits_clone);
|
||||
async move {
|
||||
*upstream_hits_inner.lock().expect("mutex should lock") += 1;
|
||||
(StatusCode::OK, Body::from("unexpected upstream hit"))
|
||||
}
|
||||
}),
|
||||
);
|
||||
|
||||
let mut key = sample_key("key-openai", "provider-openai", "openai:chat", "sk-test");
|
||||
key.rpm_limit = None;
|
||||
key.learned_rpm_limit = Some(12);
|
||||
key.concurrent_429_count = Some(2);
|
||||
key.rpm_429_count = Some(1);
|
||||
key.last_429_at_unix_secs = Some(1_711_111_111);
|
||||
key.last_429_type = Some("rpm".to_string());
|
||||
key.adjustment_history = Some(json!([
|
||||
{"timestamp": "2026-03-28T14:00:00Z", "delta": 2}
|
||||
]));
|
||||
key.status_snapshot = Some(json!({
|
||||
"learning_confidence": 0.9,
|
||||
"enforcement_active": true,
|
||||
"observation_count": 8,
|
||||
"header_observation_count": 3,
|
||||
"latest_upstream_limit": 15
|
||||
}));
|
||||
let repository = adaptive_repository_with_keys(vec![key]);
|
||||
|
||||
let (upstream_url, upstream_handle) = start_server(upstream).await;
|
||||
let gateway = build_router_with_state(
|
||||
AppState::new()
|
||||
.expect("gateway should build")
|
||||
.with_data_state_for_tests(
|
||||
GatewayDataState::with_provider_catalog_repository_for_tests(repository),
|
||||
),
|
||||
);
|
||||
let (gateway_url, gateway_handle) = start_server(gateway).await;
|
||||
|
||||
let response = reqwest::Client::new()
|
||||
.get(format!(
|
||||
"{gateway_url}/api/admin/adaptive/keys/key-openai/stats"
|
||||
))
|
||||
.header(crate::constants::GATEWAY_HEADER, "rust-phase3b")
|
||||
.header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
|
||||
.header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
|
||||
.header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
|
||||
.send()
|
||||
.await
|
||||
.expect("request should succeed");
|
||||
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
let payload: serde_json::Value = response.json().await.expect("json body should parse");
|
||||
assert_eq!(payload["adaptive_mode"], true);
|
||||
assert_eq!(payload["effective_limit"], 12);
|
||||
assert_eq!(payload["learning_confidence"], 0.9);
|
||||
assert_eq!(payload["latest_upstream_limit"], 15);
|
||||
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
|
||||
|
||||
gateway_handle.abort();
|
||||
upstream_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_toggles_admin_adaptive_mode_locally_with_trusted_admin_principal() {
|
||||
let repository = adaptive_repository_with_keys({
|
||||
let mut key = sample_key("key-openai", "provider-openai", "openai:chat", "sk-test");
|
||||
key.rpm_limit = None;
|
||||
key.learned_rpm_limit = Some(18);
|
||||
vec![key]
|
||||
});
|
||||
|
||||
let gateway = build_router_with_state(
|
||||
AppState::new()
|
||||
.expect("gateway should build")
|
||||
.with_data_state_for_tests(
|
||||
GatewayDataState::with_provider_catalog_repository_for_tests(Arc::clone(
|
||||
&repository,
|
||||
)),
|
||||
),
|
||||
);
|
||||
let (gateway_url, gateway_handle) = start_server(gateway).await;
|
||||
|
||||
let response = reqwest::Client::new()
|
||||
.patch(format!(
|
||||
"{gateway_url}/api/admin/adaptive/keys/key-openai/mode"
|
||||
))
|
||||
.header(crate::constants::GATEWAY_HEADER, "rust-phase3b")
|
||||
.header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
|
||||
.header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
|
||||
.header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
|
||||
.json(&json!({ "enabled": false, "fixed_limit": 12 }))
|
||||
.send()
|
||||
.await
|
||||
.expect("request should succeed");
|
||||
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
let payload: serde_json::Value = response.json().await.expect("json body should parse");
|
||||
assert_eq!(payload["is_adaptive"], false);
|
||||
assert_eq!(payload["rpm_limit"], 12);
|
||||
|
||||
let updated = repository
|
||||
.list_keys_by_ids(&["key-openai".to_string()])
|
||||
.await
|
||||
.expect("repository query should succeed");
|
||||
assert_eq!(updated[0].rpm_limit, Some(12));
|
||||
|
||||
gateway_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_sets_admin_adaptive_limit_locally_with_trusted_admin_principal() {
|
||||
let repository = adaptive_repository_with_keys({
|
||||
let mut key = sample_key("key-openai", "provider-openai", "openai:chat", "sk-test");
|
||||
key.rpm_limit = None;
|
||||
vec![key]
|
||||
});
|
||||
|
||||
let gateway = build_router_with_state(
|
||||
AppState::new()
|
||||
.expect("gateway should build")
|
||||
.with_data_state_for_tests(
|
||||
GatewayDataState::with_provider_catalog_repository_for_tests(Arc::clone(
|
||||
&repository,
|
||||
)),
|
||||
),
|
||||
);
|
||||
let (gateway_url, gateway_handle) = start_server(gateway).await;
|
||||
|
||||
let response = reqwest::Client::new()
|
||||
.patch(format!(
|
||||
"{gateway_url}/api/admin/adaptive/keys/key-openai/limit?limit=9"
|
||||
))
|
||||
.header(crate::constants::GATEWAY_HEADER, "rust-phase3b")
|
||||
.header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
|
||||
.header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
|
||||
.header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
|
||||
.send()
|
||||
.await
|
||||
.expect("request should succeed");
|
||||
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
let payload: serde_json::Value = response.json().await.expect("json body should parse");
|
||||
assert_eq!(payload["rpm_limit"], 9);
|
||||
assert_eq!(payload["is_adaptive"], false);
|
||||
|
||||
let updated = repository
|
||||
.list_keys_by_ids(&["key-openai".to_string()])
|
||||
.await
|
||||
.expect("repository query should succeed");
|
||||
assert_eq!(updated[0].rpm_limit, Some(9));
|
||||
|
||||
gateway_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_resets_admin_adaptive_learning_locally_with_trusted_admin_principal() {
|
||||
let repository = adaptive_repository_with_keys({
|
||||
let mut key = sample_key("key-openai", "provider-openai", "openai:chat", "sk-test");
|
||||
key.rpm_limit = None;
|
||||
key.learned_rpm_limit = Some(18);
|
||||
key.concurrent_429_count = Some(2);
|
||||
key.rpm_429_count = Some(1);
|
||||
key.last_429_at_unix_secs = Some(1_711_111_111);
|
||||
key.last_429_type = Some("rpm".to_string());
|
||||
key.adjustment_history = Some(json!([
|
||||
{"timestamp": "2026-03-28T14:00:00Z", "delta": 2}
|
||||
]));
|
||||
vec![key]
|
||||
});
|
||||
|
||||
let gateway = build_router_with_state(
|
||||
AppState::new()
|
||||
.expect("gateway should build")
|
||||
.with_data_state_for_tests(
|
||||
GatewayDataState::with_provider_catalog_repository_for_tests(Arc::clone(
|
||||
&repository,
|
||||
)),
|
||||
),
|
||||
);
|
||||
let (gateway_url, gateway_handle) = start_server(gateway).await;
|
||||
|
||||
let response = reqwest::Client::new()
|
||||
.delete(format!(
|
||||
"{gateway_url}/api/admin/adaptive/keys/key-openai/learning"
|
||||
))
|
||||
.header(crate::constants::GATEWAY_HEADER, "rust-phase3b")
|
||||
.header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
|
||||
.header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
|
||||
.header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
|
||||
.send()
|
||||
.await
|
||||
.expect("request should succeed");
|
||||
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
let payload: serde_json::Value = response.json().await.expect("json body should parse");
|
||||
assert_eq!(payload["message"], "学习状态已重置");
|
||||
|
||||
let updated = repository
|
||||
.list_keys_by_ids(&["key-openai".to_string()])
|
||||
.await
|
||||
.expect("repository query should succeed");
|
||||
assert_eq!(updated[0].learned_rpm_limit, None);
|
||||
assert_eq!(updated[0].concurrent_429_count, None);
|
||||
assert_eq!(updated[0].rpm_429_count, None);
|
||||
assert_eq!(updated[0].adjustment_history, None);
|
||||
|
||||
gateway_handle.abort();
|
||||
}
|
||||
566
apps/aether-gateway/src/tests/control/admin/api_keys.rs
Normal file
566
apps/aether-gateway/src/tests/control/admin/api_keys.rs
Normal file
@@ -0,0 +1,566 @@
|
||||
use std::sync::{Arc, Mutex};
|
||||
|
||||
use aether_crypto::{encrypt_python_fernet_plaintext, DEVELOPMENT_ENCRYPTION_KEY};
|
||||
use aether_data::repository::auth::{
|
||||
InMemoryAuthApiKeySnapshotRepository, StoredAuthApiKeyExportRecord, StoredAuthApiKeySnapshot,
|
||||
};
|
||||
use aether_data::repository::usage::{InMemoryUsageReadRepository, StoredRequestUsageAudit};
|
||||
use aether_data::repository::wallet::{InMemoryWalletRepository, StoredWalletSnapshot};
|
||||
use axum::body::Body;
|
||||
use axum::routing::any;
|
||||
use axum::{extract::Request, Router};
|
||||
use http::StatusCode;
|
||||
use serde_json::json;
|
||||
|
||||
use super::super::{build_router_with_state, start_server, AppState};
|
||||
use crate::constants::{
|
||||
GATEWAY_HEADER, TRUSTED_ADMIN_SESSION_ID_HEADER, TRUSTED_ADMIN_USER_ID_HEADER,
|
||||
TRUSTED_ADMIN_USER_ROLE_HEADER,
|
||||
};
|
||||
use crate::data::GatewayDataState;
|
||||
|
||||
fn admin_request(builder: reqwest::RequestBuilder) -> reqwest::RequestBuilder {
|
||||
builder
|
||||
.header(crate::constants::GATEWAY_HEADER, "rust-phase3b")
|
||||
.header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
|
||||
.header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
|
||||
.header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
|
||||
}
|
||||
|
||||
async fn start_api_keys_upstream(
|
||||
path: &'static str,
|
||||
) -> (String, Arc<Mutex<usize>>, tokio::task::JoinHandle<()>) {
|
||||
let upstream_hits = Arc::new(Mutex::new(0usize));
|
||||
let upstream_hits_clone = Arc::clone(&upstream_hits);
|
||||
let upstream = Router::new().route(
|
||||
path,
|
||||
any(move |_request: Request| {
|
||||
let upstream_hits_inner = Arc::clone(&upstream_hits_clone);
|
||||
async move {
|
||||
*upstream_hits_inner.lock().expect("mutex should lock") += 1;
|
||||
(StatusCode::OK, Body::from("unexpected upstream hit"))
|
||||
}
|
||||
}),
|
||||
);
|
||||
|
||||
let (upstream_url, upstream_handle) = start_server(upstream).await;
|
||||
(upstream_url, upstream_hits, upstream_handle)
|
||||
}
|
||||
|
||||
fn sample_standalone_api_key_snapshot(
|
||||
api_key_id: &str,
|
||||
user_id: &str,
|
||||
is_active: bool,
|
||||
) -> StoredAuthApiKeySnapshot {
|
||||
StoredAuthApiKeySnapshot::new(
|
||||
user_id.to_string(),
|
||||
"standalone-owner".to_string(),
|
||||
Some("owner@example.com".to_string()),
|
||||
"admin".to_string(),
|
||||
"local".to_string(),
|
||||
true,
|
||||
false,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
api_key_id.to_string(),
|
||||
Some(format!("key-{api_key_id}")),
|
||||
is_active,
|
||||
false,
|
||||
true,
|
||||
Some(120),
|
||||
Some(5),
|
||||
Some(4_102_444_800),
|
||||
Some(json!(["openai"])),
|
||||
Some(json!(["openai:chat"])),
|
||||
Some(json!(["gpt-4.1"])),
|
||||
)
|
||||
.expect("snapshot should build")
|
||||
}
|
||||
|
||||
fn sample_standalone_export_record(
|
||||
api_key_id: &str,
|
||||
user_id: &str,
|
||||
plaintext_key: &str,
|
||||
is_active: bool,
|
||||
) -> StoredAuthApiKeyExportRecord {
|
||||
StoredAuthApiKeyExportRecord::new(
|
||||
user_id.to_string(),
|
||||
api_key_id.to_string(),
|
||||
format!("hash-{api_key_id}"),
|
||||
Some(
|
||||
encrypt_python_fernet_plaintext(DEVELOPMENT_ENCRYPTION_KEY, plaintext_key)
|
||||
.expect("key should encrypt"),
|
||||
),
|
||||
Some(format!("key-{api_key_id}")),
|
||||
Some(json!(["openai"])),
|
||||
Some(json!(["openai:chat"])),
|
||||
Some(json!(["gpt-4.1"])),
|
||||
Some(120),
|
||||
Some(5),
|
||||
None,
|
||||
is_active,
|
||||
Some(4_102_444_800),
|
||||
false,
|
||||
7,
|
||||
1.25,
|
||||
true,
|
||||
)
|
||||
.expect("export record should build")
|
||||
}
|
||||
|
||||
fn sample_standalone_wallet(api_key_id: &str) -> StoredWalletSnapshot {
|
||||
StoredWalletSnapshot::new(
|
||||
format!("wallet-{api_key_id}"),
|
||||
None,
|
||||
Some(api_key_id.to_string()),
|
||||
18.5,
|
||||
1.5,
|
||||
"unlimited".to_string(),
|
||||
"USD".to_string(),
|
||||
"active".to_string(),
|
||||
30.0,
|
||||
2.0,
|
||||
0.0,
|
||||
0.0,
|
||||
1_710_000_000,
|
||||
)
|
||||
.expect("wallet should build")
|
||||
}
|
||||
|
||||
fn sample_usage_row(
|
||||
id: &str,
|
||||
request_id: &str,
|
||||
api_key_id: &str,
|
||||
total_tokens: i32,
|
||||
) -> StoredRequestUsageAudit {
|
||||
StoredRequestUsageAudit::new(
|
||||
id.to_string(),
|
||||
request_id.to_string(),
|
||||
Some("user-1".to_string()),
|
||||
Some(api_key_id.to_string()),
|
||||
Some("user-user-1".to_string()),
|
||||
Some(format!("key-{api_key_id}")),
|
||||
"openai".to_string(),
|
||||
"gpt-4.1".to_string(),
|
||||
Some("gpt-4.1".to_string()),
|
||||
Some("provider-1".to_string()),
|
||||
Some("endpoint-1".to_string()),
|
||||
Some("provider-key-1".to_string()),
|
||||
Some("chat".to_string()),
|
||||
Some("openai:chat".to_string()),
|
||||
Some("openai".to_string()),
|
||||
Some("chat".to_string()),
|
||||
Some("openai:chat".to_string()),
|
||||
Some("openai".to_string()),
|
||||
Some("chat".to_string()),
|
||||
false,
|
||||
false,
|
||||
total_tokens / 2,
|
||||
total_tokens - (total_tokens / 2),
|
||||
total_tokens,
|
||||
0.42,
|
||||
0.42,
|
||||
Some(200),
|
||||
None,
|
||||
None,
|
||||
Some(240),
|
||||
Some(80),
|
||||
"completed".to_string(),
|
||||
"settled".to_string(),
|
||||
1_711_000_000,
|
||||
1_711_000_001,
|
||||
Some(1_711_000_002),
|
||||
)
|
||||
.expect("usage row should build")
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_handles_admin_api_keys_list_locally_with_trusted_admin_principal() {
|
||||
let (upstream_url, upstream_hits, upstream_handle) =
|
||||
start_api_keys_upstream("/api/admin/api-keys").await;
|
||||
let auth_repository = Arc::new(
|
||||
InMemoryAuthApiKeySnapshotRepository::seed(vec![
|
||||
(
|
||||
None,
|
||||
sample_standalone_api_key_snapshot("key-1", "user-1", true),
|
||||
),
|
||||
(
|
||||
None,
|
||||
sample_standalone_api_key_snapshot("key-2", "user-2", false),
|
||||
),
|
||||
])
|
||||
.with_export_records([
|
||||
sample_standalone_export_record("key-1", "user-1", "sk-key-1-plaintext", true),
|
||||
sample_standalone_export_record("key-2", "user-2", "sk-key-2-plaintext", false),
|
||||
]),
|
||||
);
|
||||
let usage_repository = Arc::new(InMemoryUsageReadRepository::seed(vec![
|
||||
sample_usage_row("usage-1", "req-1", "key-1", 90),
|
||||
sample_usage_row("usage-2", "req-2", "key-2", 120),
|
||||
]));
|
||||
let gateway = build_router_with_state(
|
||||
AppState::new()
|
||||
.expect("gateway should build")
|
||||
.with_data_state_for_tests(
|
||||
crate::data::GatewayDataState::with_auth_wallet_and_usage_for_tests(
|
||||
auth_repository,
|
||||
Arc::new(InMemoryWalletRepository::seed(
|
||||
Vec::<StoredWalletSnapshot>::new(),
|
||||
)),
|
||||
usage_repository,
|
||||
)
|
||||
.with_encryption_key_for_tests(DEVELOPMENT_ENCRYPTION_KEY),
|
||||
),
|
||||
);
|
||||
let (gateway_url, gateway_handle) = start_server(gateway).await;
|
||||
|
||||
let response = admin_request(reqwest::Client::new().get(format!(
|
||||
"{gateway_url}/api/admin/api-keys?is_active=true&limit=10&skip=0"
|
||||
)))
|
||||
.send()
|
||||
.await
|
||||
.expect("request should succeed");
|
||||
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
let payload: serde_json::Value = response.json().await.expect("json body should parse");
|
||||
assert_eq!(payload["total"], json!(1));
|
||||
assert_eq!(payload["limit"], json!(10));
|
||||
assert_eq!(payload["skip"], json!(0));
|
||||
assert_eq!(payload["api_keys"][0]["id"], json!("key-1"));
|
||||
assert_eq!(payload["api_keys"][0]["is_standalone"], json!(true));
|
||||
assert_eq!(
|
||||
payload["api_keys"][0]["key_display"],
|
||||
json!("sk-key-1-p...text")
|
||||
);
|
||||
assert_eq!(payload["api_keys"][0]["total_requests"], json!(7));
|
||||
assert_eq!(payload["api_keys"][0]["total_tokens"], json!(90));
|
||||
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
|
||||
|
||||
gateway_handle.abort();
|
||||
upstream_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_handles_admin_api_keys_detail_locally_with_trusted_admin_principal() {
|
||||
let (upstream_url, upstream_hits, upstream_handle) =
|
||||
start_api_keys_upstream("/api/admin/api-keys/key-1").await;
|
||||
let auth_repository = Arc::new(
|
||||
InMemoryAuthApiKeySnapshotRepository::seed(vec![(
|
||||
None,
|
||||
sample_standalone_api_key_snapshot("key-1", "user-1", true),
|
||||
)])
|
||||
.with_export_records([sample_standalone_export_record(
|
||||
"key-1",
|
||||
"user-1",
|
||||
"sk-key-1-plaintext",
|
||||
true,
|
||||
)]),
|
||||
);
|
||||
let wallet_repository = Arc::new(InMemoryWalletRepository::seed(vec![
|
||||
sample_standalone_wallet("key-1"),
|
||||
]));
|
||||
let usage_repository = Arc::new(InMemoryUsageReadRepository::seed(vec![sample_usage_row(
|
||||
"usage-1", "req-1", "key-1", 150,
|
||||
)]));
|
||||
let gateway = build_router_with_state(
|
||||
AppState::new()
|
||||
.expect("gateway should build")
|
||||
.with_data_state_for_tests(
|
||||
crate::data::GatewayDataState::with_auth_wallet_and_usage_for_tests(
|
||||
auth_repository,
|
||||
wallet_repository,
|
||||
usage_repository,
|
||||
)
|
||||
.with_encryption_key_for_tests(DEVELOPMENT_ENCRYPTION_KEY),
|
||||
),
|
||||
);
|
||||
let (gateway_url, gateway_handle) = start_server(gateway).await;
|
||||
|
||||
let response = admin_request(
|
||||
reqwest::Client::new().get(format!("{gateway_url}/api/admin/api-keys/key-1")),
|
||||
)
|
||||
.send()
|
||||
.await
|
||||
.expect("request should succeed");
|
||||
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
let payload: serde_json::Value = response.json().await.expect("json body should parse");
|
||||
assert_eq!(payload["id"], json!("key-1"));
|
||||
assert_eq!(payload["user_id"], json!("user-1"));
|
||||
assert_eq!(payload["wallet"]["id"], json!("wallet-key-1"));
|
||||
assert_eq!(payload["wallet"]["unlimited"], json!(true));
|
||||
assert_eq!(payload["wallet"]["balance"], json!(20.0));
|
||||
assert_eq!(payload["key_display"], json!("sk-key-1-p...text"));
|
||||
assert_eq!(payload["total_tokens"], json!(150));
|
||||
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
|
||||
|
||||
gateway_handle.abort();
|
||||
upstream_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_handles_admin_api_keys_full_key_locally_with_trusted_admin_principal() {
|
||||
let (upstream_url, upstream_hits, upstream_handle) =
|
||||
start_api_keys_upstream("/api/admin/api-keys/key-1").await;
|
||||
let auth_repository = Arc::new(
|
||||
InMemoryAuthApiKeySnapshotRepository::seed(vec![(
|
||||
None,
|
||||
sample_standalone_api_key_snapshot("key-1", "user-1", true),
|
||||
)])
|
||||
.with_export_records([sample_standalone_export_record(
|
||||
"key-1",
|
||||
"user-1",
|
||||
"sk-key-1-plaintext",
|
||||
true,
|
||||
)]),
|
||||
);
|
||||
let wallet_repository = Arc::new(InMemoryWalletRepository::seed(vec![
|
||||
sample_standalone_wallet("key-1"),
|
||||
]));
|
||||
let gateway = build_router_with_state(
|
||||
AppState::new()
|
||||
.expect("gateway should build")
|
||||
.with_data_state_for_tests(
|
||||
crate::data::GatewayDataState::with_auth_and_wallet_for_tests(
|
||||
auth_repository,
|
||||
wallet_repository,
|
||||
)
|
||||
.with_encryption_key_for_tests(DEVELOPMENT_ENCRYPTION_KEY),
|
||||
),
|
||||
);
|
||||
let (gateway_url, gateway_handle) = start_server(gateway).await;
|
||||
|
||||
let response = admin_request(reqwest::Client::new().get(format!(
|
||||
"{gateway_url}/api/admin/api-keys/key-1?include_key=true"
|
||||
)))
|
||||
.send()
|
||||
.await
|
||||
.expect("request should succeed");
|
||||
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
let payload: serde_json::Value = response.json().await.expect("json body should parse");
|
||||
assert_eq!(payload, json!({ "key": "sk-key-1-plaintext" }));
|
||||
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
|
||||
|
||||
gateway_handle.abort();
|
||||
upstream_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_handles_admin_api_keys_create_locally_with_trusted_admin_principal() {
|
||||
let (upstream_url, upstream_hits, upstream_handle) =
|
||||
start_api_keys_upstream("/api/admin/api-keys").await;
|
||||
let repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed(Vec::<(
|
||||
Option<String>,
|
||||
StoredAuthApiKeySnapshot,
|
||||
)>::new()));
|
||||
let gateway = build_router_with_state(
|
||||
AppState::new()
|
||||
.expect("gateway should build")
|
||||
.with_data_state_for_tests(
|
||||
crate::data::GatewayDataState::with_auth_api_key_repository_for_tests(
|
||||
Arc::clone(&repository),
|
||||
)
|
||||
.with_encryption_key_for_tests(DEVELOPMENT_ENCRYPTION_KEY),
|
||||
),
|
||||
);
|
||||
let (gateway_url, gateway_handle) = start_server(gateway).await;
|
||||
|
||||
let response =
|
||||
admin_request(reqwest::Client::new().post(format!("{gateway_url}/api/admin/api-keys")))
|
||||
.json(&json!({
|
||||
"name": "standalone-key",
|
||||
"rate_limit": 180,
|
||||
"allowed_providers": ["openai"],
|
||||
"allowed_api_formats": ["openai:chat"],
|
||||
"allowed_models": ["gpt-4.1"],
|
||||
}))
|
||||
.send()
|
||||
.await
|
||||
.expect("request should succeed");
|
||||
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
let payload: serde_json::Value = response.json().await.expect("json body should parse");
|
||||
assert_eq!(payload["name"], json!("standalone-key"));
|
||||
assert_eq!(payload["is_standalone"], json!(true));
|
||||
assert_eq!(payload["rate_limit"], json!(180));
|
||||
assert_eq!(payload["allowed_providers"], json!(["openai"]));
|
||||
assert_eq!(payload["allowed_api_formats"], json!(["openai:chat"]));
|
||||
assert_eq!(payload["allowed_models"], json!(["gpt-4.1"]));
|
||||
assert_eq!(payload["wallet"], serde_json::Value::Null);
|
||||
let plaintext = payload["key"]
|
||||
.as_str()
|
||||
.expect("plaintext key should exist")
|
||||
.to_string();
|
||||
assert!(plaintext.starts_with("sk-"));
|
||||
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
|
||||
|
||||
let list_response =
|
||||
admin_request(reqwest::Client::new().get(format!("{gateway_url}/api/admin/api-keys")))
|
||||
.send()
|
||||
.await
|
||||
.expect("list request should succeed");
|
||||
assert_eq!(list_response.status(), StatusCode::OK);
|
||||
let list_payload: serde_json::Value =
|
||||
list_response.json().await.expect("list json should parse");
|
||||
assert_eq!(list_payload["total"], json!(1));
|
||||
assert_eq!(list_payload["api_keys"][0]["name"], json!("standalone-key"));
|
||||
|
||||
gateway_handle.abort();
|
||||
upstream_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_handles_admin_api_keys_update_locally_with_trusted_admin_principal() {
|
||||
let (upstream_url, upstream_hits, upstream_handle) =
|
||||
start_api_keys_upstream("/api/admin/api-keys/key-123").await;
|
||||
let repository = Arc::new(
|
||||
InMemoryAuthApiKeySnapshotRepository::seed(vec![(
|
||||
Some("hash-key-123".to_string()),
|
||||
sample_standalone_api_key_snapshot("key-123", "admin-user-123", true),
|
||||
)])
|
||||
.with_export_records([sample_standalone_export_record(
|
||||
"key-123",
|
||||
"admin-user-123",
|
||||
"sk-key-123-plaintext",
|
||||
true,
|
||||
)]),
|
||||
);
|
||||
let gateway = build_router_with_state(
|
||||
AppState::new()
|
||||
.expect("gateway should build")
|
||||
.with_data_state_for_tests(
|
||||
crate::data::GatewayDataState::with_auth_api_key_repository_for_tests(
|
||||
Arc::clone(&repository),
|
||||
)
|
||||
.with_encryption_key_for_tests(DEVELOPMENT_ENCRYPTION_KEY),
|
||||
),
|
||||
);
|
||||
let (gateway_url, gateway_handle) = start_server(gateway).await;
|
||||
|
||||
let response = admin_request(
|
||||
reqwest::Client::new().put(format!("{gateway_url}/api/admin/api-keys/key-123")),
|
||||
)
|
||||
.json(&json!({
|
||||
"name": "renamed-key",
|
||||
"rate_limit": 240,
|
||||
"allowed_providers": ["gemini"],
|
||||
"allowed_api_formats": ["gemini:chat"],
|
||||
"allowed_models": ["gemini-2.5-pro"],
|
||||
}))
|
||||
.send()
|
||||
.await
|
||||
.expect("request should succeed");
|
||||
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
let payload: serde_json::Value = response.json().await.expect("json body should parse");
|
||||
assert_eq!(payload["id"], json!("key-123"));
|
||||
assert_eq!(payload["name"], json!("renamed-key"));
|
||||
assert_eq!(payload["rate_limit"], json!(240));
|
||||
assert_eq!(payload["allowed_providers"], json!(["gemini"]));
|
||||
assert_eq!(payload["allowed_api_formats"], json!(["gemini:chat"]));
|
||||
assert_eq!(payload["allowed_models"], json!(["gemini-2.5-pro"]));
|
||||
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
|
||||
|
||||
gateway_handle.abort();
|
||||
upstream_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_handles_admin_api_keys_toggle_locally_with_trusted_admin_principal() {
|
||||
let (upstream_url, upstream_hits, upstream_handle) =
|
||||
start_api_keys_upstream("/api/admin/api-keys/key-123").await;
|
||||
let repository = Arc::new(
|
||||
InMemoryAuthApiKeySnapshotRepository::seed(vec![(
|
||||
Some("hash-key-123".to_string()),
|
||||
sample_standalone_api_key_snapshot("key-123", "admin-user-123", true),
|
||||
)])
|
||||
.with_export_records([sample_standalone_export_record(
|
||||
"key-123",
|
||||
"admin-user-123",
|
||||
"sk-key-123-plaintext",
|
||||
true,
|
||||
)]),
|
||||
);
|
||||
let gateway = build_router_with_state(
|
||||
AppState::new()
|
||||
.expect("gateway should build")
|
||||
.with_data_state_for_tests(
|
||||
crate::data::GatewayDataState::with_auth_api_key_repository_for_tests(
|
||||
Arc::clone(&repository),
|
||||
)
|
||||
.with_encryption_key_for_tests(DEVELOPMENT_ENCRYPTION_KEY),
|
||||
),
|
||||
);
|
||||
let (gateway_url, gateway_handle) = start_server(gateway).await;
|
||||
|
||||
let response = admin_request(
|
||||
reqwest::Client::new().patch(format!("{gateway_url}/api/admin/api-keys/key-123")),
|
||||
)
|
||||
.send()
|
||||
.await
|
||||
.expect("request should succeed");
|
||||
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
let payload: serde_json::Value = response.json().await.expect("json body should parse");
|
||||
assert_eq!(payload["id"], json!("key-123"));
|
||||
assert_eq!(payload["is_active"], json!(false));
|
||||
assert_eq!(payload["message"], json!("API密钥已禁用"));
|
||||
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
|
||||
|
||||
gateway_handle.abort();
|
||||
upstream_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_handles_admin_api_keys_delete_locally_with_trusted_admin_principal() {
|
||||
let (upstream_url, upstream_hits, upstream_handle) =
|
||||
start_api_keys_upstream("/api/admin/api-keys/key-123").await;
|
||||
let repository = Arc::new(
|
||||
InMemoryAuthApiKeySnapshotRepository::seed(vec![(
|
||||
Some("hash-key-123".to_string()),
|
||||
sample_standalone_api_key_snapshot("key-123", "admin-user-123", true),
|
||||
)])
|
||||
.with_export_records([sample_standalone_export_record(
|
||||
"key-123",
|
||||
"admin-user-123",
|
||||
"sk-key-123-plaintext",
|
||||
true,
|
||||
)]),
|
||||
);
|
||||
let gateway = build_router_with_state(
|
||||
AppState::new()
|
||||
.expect("gateway should build")
|
||||
.with_data_state_for_tests(
|
||||
crate::data::GatewayDataState::with_auth_api_key_repository_for_tests(
|
||||
Arc::clone(&repository),
|
||||
)
|
||||
.with_encryption_key_for_tests(DEVELOPMENT_ENCRYPTION_KEY),
|
||||
),
|
||||
);
|
||||
let (gateway_url, gateway_handle) = start_server(gateway).await;
|
||||
|
||||
let response = admin_request(
|
||||
reqwest::Client::new().delete(format!("{gateway_url}/api/admin/api-keys/key-123")),
|
||||
)
|
||||
.send()
|
||||
.await
|
||||
.expect("request should succeed");
|
||||
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
let payload: serde_json::Value = response.json().await.expect("json body should parse");
|
||||
assert_eq!(payload["message"], json!("API密钥已删除"));
|
||||
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
|
||||
|
||||
let detail_response = admin_request(
|
||||
reqwest::Client::new().get(format!("{gateway_url}/api/admin/api-keys/key-123")),
|
||||
)
|
||||
.send()
|
||||
.await
|
||||
.expect("detail request should succeed");
|
||||
assert_eq!(detail_response.status(), StatusCode::NOT_FOUND);
|
||||
|
||||
gateway_handle.abort();
|
||||
upstream_handle.abort();
|
||||
}
|
||||
663
apps/aether-gateway/src/tests/control/admin/billing.rs
Normal file
663
apps/aether-gateway/src/tests/control/admin/billing.rs
Normal file
@@ -0,0 +1,663 @@
|
||||
use std::sync::{Arc, Mutex};
|
||||
|
||||
use aether_data::repository::users::InMemoryUserReadRepository;
|
||||
use axum::body::Body;
|
||||
use axum::routing::any;
|
||||
use axum::{extract::Request, Router};
|
||||
use http::StatusCode;
|
||||
use serde_json::json;
|
||||
|
||||
use super::super::{build_router_with_state, start_server, AppState};
|
||||
use crate::constants::{
|
||||
GATEWAY_HEADER, TRUSTED_ADMIN_SESSION_ID_HEADER, TRUSTED_ADMIN_USER_ID_HEADER,
|
||||
TRUSTED_ADMIN_USER_ROLE_HEADER,
|
||||
};
|
||||
use crate::data::GatewayDataState;
|
||||
|
||||
async fn assert_admin_billing_get_returns_local(
|
||||
path: &str,
|
||||
) -> (serde_json::Value, Arc<Mutex<usize>>) {
|
||||
let upstream_hits = Arc::new(Mutex::new(0usize));
|
||||
let upstream_hits_clone = Arc::clone(&upstream_hits);
|
||||
let upstream = Router::new().route(
|
||||
path,
|
||||
any(move |_request: Request| {
|
||||
let upstream_hits_inner = Arc::clone(&upstream_hits_clone);
|
||||
async move {
|
||||
*upstream_hits_inner.lock().expect("mutex should lock") += 1;
|
||||
(StatusCode::OK, Body::from("unexpected upstream hit"))
|
||||
}
|
||||
}),
|
||||
);
|
||||
|
||||
let (upstream_url, upstream_handle) = start_server(upstream).await;
|
||||
let gateway = build_router_with_state(AppState::new().expect("gateway should build"));
|
||||
let (gateway_url, gateway_handle) = start_server(gateway).await;
|
||||
|
||||
let response = reqwest::Client::new()
|
||||
.get(format!("{gateway_url}{path}"))
|
||||
.header(crate::constants::GATEWAY_HEADER, "rust-phase3b")
|
||||
.header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
|
||||
.header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
|
||||
.header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
|
||||
.send()
|
||||
.await
|
||||
.expect("request should succeed");
|
||||
|
||||
let status = response.status();
|
||||
let payload: serde_json::Value = response.json().await.expect("json body should parse");
|
||||
gateway_handle.abort();
|
||||
upstream_handle.abort();
|
||||
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
|
||||
assert_eq!(status, StatusCode::OK);
|
||||
(payload, upstream_hits)
|
||||
}
|
||||
|
||||
async fn send_admin_billing_request(
|
||||
gateway_url: &str,
|
||||
method: http::Method,
|
||||
path: &str,
|
||||
body: Option<serde_json::Value>,
|
||||
) -> reqwest::Response {
|
||||
let client = reqwest::Client::new();
|
||||
let mut request = client
|
||||
.request(method, format!("{gateway_url}{path}"))
|
||||
.header(crate::constants::GATEWAY_HEADER, "rust-phase3b")
|
||||
.header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
|
||||
.header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
|
||||
.header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123");
|
||||
if let Some(body) = body {
|
||||
request = request.json(&body);
|
||||
}
|
||||
request.send().await.expect("request should succeed")
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_handles_admin_billing_presets_locally_with_trusted_admin_principal() {
|
||||
let upstream_hits = Arc::new(Mutex::new(0usize));
|
||||
let upstream_hits_clone = Arc::clone(&upstream_hits);
|
||||
let upstream = Router::new().route(
|
||||
"/api/admin/billing/presets",
|
||||
any(move |_request: Request| {
|
||||
let upstream_hits_inner = Arc::clone(&upstream_hits_clone);
|
||||
async move {
|
||||
*upstream_hits_inner.lock().expect("mutex should lock") += 1;
|
||||
(StatusCode::OK, Body::from("unexpected upstream hit"))
|
||||
}
|
||||
}),
|
||||
);
|
||||
|
||||
let (upstream_url, upstream_handle) = start_server(upstream).await;
|
||||
let gateway = build_router_with_state(AppState::new().expect("gateway should build"));
|
||||
let (gateway_url, gateway_handle) = start_server(gateway).await;
|
||||
|
||||
let response = reqwest::Client::new()
|
||||
.get(format!("{gateway_url}/api/admin/billing/presets"))
|
||||
.header(crate::constants::GATEWAY_HEADER, "rust-phase3b")
|
||||
.header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
|
||||
.header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
|
||||
.header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
|
||||
.send()
|
||||
.await
|
||||
.expect("request should succeed");
|
||||
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
let payload: serde_json::Value = response.json().await.expect("json body should parse");
|
||||
assert_eq!(payload["items"].as_array().map(Vec::len), Some(1));
|
||||
assert_eq!(payload["items"][0]["name"], "aether-core");
|
||||
assert_eq!(payload["items"][0]["version"], "1.0");
|
||||
assert_eq!(payload["items"][0]["collector_count"], 16);
|
||||
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
|
||||
|
||||
gateway_handle.abort();
|
||||
upstream_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_handles_admin_billing_apply_preset_locally_with_trusted_admin_principal() {
|
||||
let upstream_hits = Arc::new(Mutex::new(0usize));
|
||||
let upstream_hits_clone = Arc::clone(&upstream_hits);
|
||||
let upstream = Router::new().route(
|
||||
"/api/admin/billing/presets/apply",
|
||||
any(move |_request: Request| {
|
||||
let upstream_hits_inner = Arc::clone(&upstream_hits_clone);
|
||||
async move {
|
||||
*upstream_hits_inner.lock().expect("mutex should lock") += 1;
|
||||
(StatusCode::OK, Body::from("unexpected upstream hit"))
|
||||
}
|
||||
}),
|
||||
);
|
||||
let (upstream_url, upstream_handle) = start_server(upstream).await;
|
||||
let gateway = build_router_with_state(
|
||||
AppState::new()
|
||||
.expect("gateway should build")
|
||||
.with_admin_billing_collectors_for_tests(Vec::<
|
||||
crate::AdminBillingCollectorRecord,
|
||||
>::new()),
|
||||
);
|
||||
let (gateway_url, gateway_handle) = start_server(gateway).await;
|
||||
|
||||
let merge_response = send_admin_billing_request(
|
||||
&gateway_url,
|
||||
http::Method::POST,
|
||||
"/api/admin/billing/presets/apply",
|
||||
Some(json!({"preset":"aether-core","mode":"merge"})),
|
||||
)
|
||||
.await;
|
||||
assert_eq!(merge_response.status(), StatusCode::OK);
|
||||
let merge_payload: serde_json::Value =
|
||||
merge_response.json().await.expect("json body should parse");
|
||||
assert_eq!(merge_payload["ok"], json!(true));
|
||||
assert_eq!(merge_payload["preset"], json!("aether-core"));
|
||||
assert_eq!(merge_payload["mode"], json!("merge"));
|
||||
assert_eq!(merge_payload["created"], json!(16));
|
||||
assert_eq!(merge_payload["updated"], json!(0));
|
||||
assert_eq!(merge_payload["skipped"], json!(0));
|
||||
assert_eq!(merge_payload["errors"], json!([]));
|
||||
|
||||
let merge_again_response = send_admin_billing_request(
|
||||
&gateway_url,
|
||||
http::Method::POST,
|
||||
"/api/admin/billing/presets/apply",
|
||||
Some(json!({"preset":"default","mode":"merge"})),
|
||||
)
|
||||
.await;
|
||||
assert_eq!(merge_again_response.status(), StatusCode::OK);
|
||||
let merge_again_payload: serde_json::Value = merge_again_response
|
||||
.json()
|
||||
.await
|
||||
.expect("json body should parse");
|
||||
assert_eq!(merge_again_payload["ok"], json!(true));
|
||||
assert_eq!(merge_again_payload["preset"], json!("aether-core"));
|
||||
assert_eq!(merge_again_payload["created"], json!(0));
|
||||
assert_eq!(merge_again_payload["updated"], json!(0));
|
||||
assert_eq!(merge_again_payload["skipped"], json!(16));
|
||||
|
||||
let overwrite_response = send_admin_billing_request(
|
||||
&gateway_url,
|
||||
http::Method::POST,
|
||||
"/api/admin/billing/presets/apply",
|
||||
Some(json!({"preset":"aether-core","mode":"overwrite"})),
|
||||
)
|
||||
.await;
|
||||
assert_eq!(overwrite_response.status(), StatusCode::OK);
|
||||
let overwrite_payload: serde_json::Value = overwrite_response
|
||||
.json()
|
||||
.await
|
||||
.expect("json body should parse");
|
||||
assert_eq!(overwrite_payload["ok"], json!(true));
|
||||
assert_eq!(overwrite_payload["preset"], json!("aether-core"));
|
||||
assert_eq!(overwrite_payload["mode"], json!("overwrite"));
|
||||
assert_eq!(overwrite_payload["created"], json!(0));
|
||||
assert_eq!(overwrite_payload["updated"], json!(16));
|
||||
assert_eq!(overwrite_payload["skipped"], json!(0));
|
||||
assert_eq!(overwrite_payload["errors"], json!([]));
|
||||
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
|
||||
|
||||
gateway_handle.abort();
|
||||
upstream_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_returns_conflict_for_admin_billing_apply_preset_when_backend_unavailable() {
|
||||
let upstream_hits = Arc::new(Mutex::new(0usize));
|
||||
let upstream_hits_clone = Arc::clone(&upstream_hits);
|
||||
let upstream = Router::new().route(
|
||||
"/api/admin/billing/presets/apply",
|
||||
any(move |_request: Request| {
|
||||
let upstream_hits_inner = Arc::clone(&upstream_hits_clone);
|
||||
async move {
|
||||
*upstream_hits_inner.lock().expect("mutex should lock") += 1;
|
||||
(StatusCode::OK, Body::from("unexpected upstream hit"))
|
||||
}
|
||||
}),
|
||||
);
|
||||
let (upstream_url, upstream_handle) = start_server(upstream).await;
|
||||
let mut state = AppState::new()
|
||||
.expect("gateway should build")
|
||||
.with_data_state_for_tests(
|
||||
crate::data::GatewayDataState::with_user_reader_for_tests(Arc::new(
|
||||
InMemoryUserReadRepository::seed_auth_users(Vec::new()),
|
||||
)),
|
||||
);
|
||||
state.admin_billing_collector_store = None;
|
||||
let gateway = build_router_with_state(state);
|
||||
let (gateway_url, gateway_handle) = start_server(gateway).await;
|
||||
|
||||
let response = send_admin_billing_request(
|
||||
&gateway_url,
|
||||
http::Method::POST,
|
||||
"/api/admin/billing/presets/apply",
|
||||
Some(json!({"preset":"aether-core","mode":"merge"})),
|
||||
)
|
||||
.await;
|
||||
assert_eq!(response.status(), StatusCode::CONFLICT);
|
||||
let payload: serde_json::Value = response.json().await.expect("json body should parse");
|
||||
assert_eq!(payload["error_code"], "read_only_mode");
|
||||
assert_eq!(payload["detail"], "当前为只读模式,无法应用计费预设");
|
||||
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
|
||||
|
||||
gateway_handle.abort();
|
||||
upstream_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_handles_admin_billing_rule_routes_locally_with_trusted_admin_principal() {
|
||||
let upstream_hits = Arc::new(Mutex::new(0usize));
|
||||
let upstream_hits_clone = Arc::clone(&upstream_hits);
|
||||
let upstream_hits_clone_detail = Arc::clone(&upstream_hits);
|
||||
let upstream = Router::new()
|
||||
.route(
|
||||
"/api/admin/billing/rules",
|
||||
any(move |_request: Request| {
|
||||
let upstream_hits_inner = Arc::clone(&upstream_hits_clone);
|
||||
async move {
|
||||
*upstream_hits_inner.lock().expect("mutex should lock") += 1;
|
||||
(StatusCode::OK, Body::from("unexpected upstream hit"))
|
||||
}
|
||||
}),
|
||||
)
|
||||
.route(
|
||||
"/api/admin/billing/rules/rule-1",
|
||||
any(move |_request: Request| {
|
||||
let upstream_hits_inner = Arc::clone(&upstream_hits_clone_detail);
|
||||
async move {
|
||||
*upstream_hits_inner.lock().expect("mutex should lock") += 1;
|
||||
(StatusCode::OK, Body::from("unexpected upstream hit"))
|
||||
}
|
||||
}),
|
||||
);
|
||||
let (upstream_url, upstream_handle) = start_server(upstream).await;
|
||||
let gateway = build_router_with_state(
|
||||
AppState::new()
|
||||
.expect("gateway should build")
|
||||
.with_admin_billing_rules_for_tests([
|
||||
crate::AdminBillingRuleRecord {
|
||||
id: "rule-1".to_string(),
|
||||
name: "chat-input".to_string(),
|
||||
task_type: "chat".to_string(),
|
||||
global_model_id: Some("global-gpt-5".to_string()),
|
||||
model_id: Some("model-1".to_string()),
|
||||
expression: "input_tokens * 0.000001".to_string(),
|
||||
variables: json!({ "base": 1.5 }),
|
||||
dimension_mappings: json!({
|
||||
"input_tokens": { "source": "dimension", "dimension": "input_tokens" }
|
||||
}),
|
||||
is_enabled: true,
|
||||
created_at_unix_secs: 1_710_000_000,
|
||||
updated_at_unix_secs: 1_710_000_100,
|
||||
},
|
||||
crate::AdminBillingRuleRecord {
|
||||
id: "rule-2".to_string(),
|
||||
name: "image-output".to_string(),
|
||||
task_type: "image".to_string(),
|
||||
global_model_id: None,
|
||||
model_id: Some("model-2".to_string()),
|
||||
expression: "images * 0.01".to_string(),
|
||||
variables: json!({}),
|
||||
dimension_mappings: json!({
|
||||
"images": { "source": "dimension", "dimension": "image_count" }
|
||||
}),
|
||||
is_enabled: false,
|
||||
created_at_unix_secs: 1_710_000_000,
|
||||
updated_at_unix_secs: 1_710_000_050,
|
||||
},
|
||||
]),
|
||||
);
|
||||
let (gateway_url, gateway_handle) = start_server(gateway).await;
|
||||
|
||||
let list_response = reqwest::Client::new()
|
||||
.get(format!(
|
||||
"{gateway_url}/api/admin/billing/rules?task_type=chat&is_enabled=true&page=1&page_size=1"
|
||||
))
|
||||
.header(crate::constants::GATEWAY_HEADER, "rust-phase3b")
|
||||
.header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
|
||||
.header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
|
||||
.header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
|
||||
.send()
|
||||
.await
|
||||
.expect("request should succeed");
|
||||
assert_eq!(list_response.status(), StatusCode::OK);
|
||||
let list_payload: serde_json::Value =
|
||||
list_response.json().await.expect("json body should parse");
|
||||
assert_eq!(list_payload["items"].as_array().map(Vec::len), Some(1));
|
||||
assert_eq!(list_payload["items"][0]["id"], "rule-1");
|
||||
assert_eq!(list_payload["items"][0]["name"], "chat-input");
|
||||
assert_eq!(list_payload["items"][0]["task_type"], "chat");
|
||||
assert_eq!(list_payload["items"][0]["is_enabled"], true);
|
||||
assert_eq!(list_payload["total"], 1);
|
||||
assert_eq!(list_payload["page"], 1);
|
||||
assert_eq!(list_payload["page_size"], 1);
|
||||
assert_eq!(list_payload["pages"], 1);
|
||||
|
||||
let detail_response = reqwest::Client::new()
|
||||
.get(format!("{gateway_url}/api/admin/billing/rules/rule-1"))
|
||||
.header(crate::constants::GATEWAY_HEADER, "rust-phase3b")
|
||||
.header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
|
||||
.header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
|
||||
.header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
|
||||
.send()
|
||||
.await
|
||||
.expect("request should succeed");
|
||||
assert_eq!(detail_response.status(), StatusCode::OK);
|
||||
let detail_payload: serde_json::Value = detail_response
|
||||
.json()
|
||||
.await
|
||||
.expect("json body should parse");
|
||||
assert_eq!(detail_payload["id"], "rule-1");
|
||||
assert_eq!(detail_payload["name"], "chat-input");
|
||||
assert_eq!(detail_payload["task_type"], "chat");
|
||||
assert_eq!(detail_payload["global_model_id"], "global-gpt-5");
|
||||
assert_eq!(detail_payload["model_id"], "model-1");
|
||||
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
|
||||
gateway_handle.abort();
|
||||
upstream_handle.abort();
|
||||
|
||||
let upstream =
|
||||
Router::new().route(
|
||||
"/api/admin/billing/rules",
|
||||
any(|_request: Request| async move {
|
||||
(StatusCode::OK, Body::from("unexpected upstream hit"))
|
||||
}),
|
||||
);
|
||||
let (upstream_url, upstream_handle) = start_server(upstream).await;
|
||||
let gateway = build_router_with_state(AppState::new().expect("gateway should build"));
|
||||
let (gateway_url, gateway_handle) = start_server(gateway).await;
|
||||
|
||||
let create_response = send_admin_billing_request(
|
||||
&gateway_url,
|
||||
http::Method::POST,
|
||||
"/api/admin/billing/rules",
|
||||
Some(json!({
|
||||
"name": "rule",
|
||||
"task_type": "chat",
|
||||
"model_id": "model-1",
|
||||
"expression": "input_tokens * 0.000001",
|
||||
"variables": { "base": 1.5 },
|
||||
"dimension_mappings": {
|
||||
"input_tokens": { "source": "dimension", "dimension": "input_tokens" }
|
||||
},
|
||||
"is_enabled": true
|
||||
})),
|
||||
)
|
||||
.await;
|
||||
assert_eq!(create_response.status(), StatusCode::OK);
|
||||
let create_payload: serde_json::Value = create_response
|
||||
.json()
|
||||
.await
|
||||
.expect("json body should parse");
|
||||
let rule_id = create_payload["id"]
|
||||
.as_str()
|
||||
.expect("rule id should exist")
|
||||
.to_string();
|
||||
assert_eq!(create_payload["name"], json!("rule"));
|
||||
assert_eq!(create_payload["model_id"], json!("model-1"));
|
||||
|
||||
let update_response = send_admin_billing_request(
|
||||
&gateway_url,
|
||||
http::Method::PUT,
|
||||
&format!("/api/admin/billing/rules/{rule_id}"),
|
||||
Some(json!({
|
||||
"name": "rule-updated",
|
||||
"task_type": "video",
|
||||
"global_model_id": "global-model-1",
|
||||
"expression": "max(output_tokens, 1)",
|
||||
"variables": { "factor": 2 },
|
||||
"dimension_mappings": {
|
||||
"output_tokens": { "source": "dimension", "dimension": "output_tokens" }
|
||||
},
|
||||
"is_enabled": false
|
||||
})),
|
||||
)
|
||||
.await;
|
||||
assert_eq!(update_response.status(), StatusCode::OK);
|
||||
let update_payload: serde_json::Value = update_response
|
||||
.json()
|
||||
.await
|
||||
.expect("json body should parse");
|
||||
assert_eq!(update_payload["id"], json!(rule_id));
|
||||
assert_eq!(update_payload["name"], json!("rule-updated"));
|
||||
assert_eq!(update_payload["task_type"], json!("video"));
|
||||
assert_eq!(update_payload["global_model_id"], json!("global-model-1"));
|
||||
assert_eq!(update_payload["model_id"], serde_json::Value::Null);
|
||||
assert_eq!(update_payload["is_enabled"], json!(false));
|
||||
|
||||
gateway_handle.abort();
|
||||
upstream_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_returns_conflict_for_admin_billing_rule_create_when_backend_unavailable() {
|
||||
let upstream_hits = Arc::new(Mutex::new(0usize));
|
||||
let upstream_hits_clone = Arc::clone(&upstream_hits);
|
||||
let upstream = Router::new().route(
|
||||
"/api/admin/billing/rules",
|
||||
any(move |_request: Request| {
|
||||
let upstream_hits_inner = Arc::clone(&upstream_hits_clone);
|
||||
async move {
|
||||
*upstream_hits_inner.lock().expect("mutex should lock") += 1;
|
||||
(StatusCode::OK, Body::from("unexpected upstream hit"))
|
||||
}
|
||||
}),
|
||||
);
|
||||
let (upstream_url, upstream_handle) = start_server(upstream).await;
|
||||
let mut state = AppState::new()
|
||||
.expect("gateway should build")
|
||||
.with_data_state_for_tests(
|
||||
crate::data::GatewayDataState::with_user_reader_for_tests(Arc::new(
|
||||
InMemoryUserReadRepository::seed_auth_users(Vec::new()),
|
||||
)),
|
||||
);
|
||||
state.admin_billing_rule_store = None;
|
||||
let gateway = build_router_with_state(state);
|
||||
let (gateway_url, gateway_handle) = start_server(gateway).await;
|
||||
|
||||
let response = send_admin_billing_request(
|
||||
&gateway_url,
|
||||
http::Method::POST,
|
||||
"/api/admin/billing/rules",
|
||||
Some(json!({
|
||||
"name": "rule",
|
||||
"task_type": "chat",
|
||||
"model_id": "model-1",
|
||||
"expression": "input_tokens * 0.000001",
|
||||
"variables": {},
|
||||
"dimension_mappings": {},
|
||||
"is_enabled": true
|
||||
})),
|
||||
)
|
||||
.await;
|
||||
assert_eq!(response.status(), StatusCode::CONFLICT);
|
||||
let payload: serde_json::Value = response.json().await.expect("json body should parse");
|
||||
assert_eq!(payload["error_code"], "read_only_mode");
|
||||
assert_eq!(payload["detail"], "当前为只读模式,无法创建计费规则");
|
||||
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
|
||||
|
||||
gateway_handle.abort();
|
||||
upstream_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_handles_admin_billing_collector_routes_locally_with_trusted_admin_principal() {
|
||||
let upstream_hits = Arc::new(Mutex::new(0usize));
|
||||
let upstream_hits_clone = Arc::clone(&upstream_hits);
|
||||
let upstream_hits_clone_list = Arc::clone(&upstream_hits);
|
||||
let upstream = Router::new()
|
||||
.route(
|
||||
"/api/admin/billing/collectors",
|
||||
any(move |_request: Request| {
|
||||
let upstream_hits_inner = Arc::clone(&upstream_hits_clone_list);
|
||||
async move {
|
||||
*upstream_hits_inner.lock().expect("mutex should lock") += 1;
|
||||
(StatusCode::OK, Body::from("unexpected upstream hit"))
|
||||
}
|
||||
}),
|
||||
)
|
||||
.route(
|
||||
"/api/admin/billing/collectors/collector-1",
|
||||
any(move |_request: Request| {
|
||||
let upstream_hits_inner = Arc::clone(&upstream_hits_clone);
|
||||
async move {
|
||||
*upstream_hits_inner.lock().expect("mutex should lock") += 1;
|
||||
(StatusCode::OK, Body::from("unexpected upstream hit"))
|
||||
}
|
||||
}),
|
||||
);
|
||||
let (upstream_url, upstream_handle) = start_server(upstream).await;
|
||||
let gateway = build_router_with_state(
|
||||
AppState::new()
|
||||
.expect("gateway should build")
|
||||
.with_admin_billing_collectors_for_tests([
|
||||
crate::AdminBillingCollectorRecord {
|
||||
id: "collector-1".to_string(),
|
||||
api_format: "OPENAI".to_string(),
|
||||
task_type: "chat".to_string(),
|
||||
dimension_name: "tokens".to_string(),
|
||||
source_type: "request".to_string(),
|
||||
source_path: Some("usage.input_tokens".to_string()),
|
||||
value_type: "float".to_string(),
|
||||
transform_expression: None,
|
||||
default_value: None,
|
||||
priority: 10,
|
||||
is_enabled: true,
|
||||
created_at_unix_secs: 1_710_000_000,
|
||||
updated_at_unix_secs: 1_710_000_100,
|
||||
},
|
||||
crate::AdminBillingCollectorRecord {
|
||||
id: "collector-2".to_string(),
|
||||
api_format: "OPENAI".to_string(),
|
||||
task_type: "chat".to_string(),
|
||||
dimension_name: "latency".to_string(),
|
||||
source_type: "computed".to_string(),
|
||||
source_path: None,
|
||||
value_type: "int".to_string(),
|
||||
transform_expression: Some("max(latency_ms, 1)".to_string()),
|
||||
default_value: Some("1".to_string()),
|
||||
priority: 5,
|
||||
is_enabled: false,
|
||||
created_at_unix_secs: 1_710_000_000,
|
||||
updated_at_unix_secs: 1_710_000_050,
|
||||
},
|
||||
]),
|
||||
);
|
||||
let (gateway_url, gateway_handle) = start_server(gateway).await;
|
||||
let list_response = reqwest::Client::new()
|
||||
.get(format!(
|
||||
"{gateway_url}/api/admin/billing/collectors?api_format=OPENAI&task_type=chat&is_enabled=true&page=1&page_size=10"
|
||||
))
|
||||
.header(crate::constants::GATEWAY_HEADER, "rust-phase3b")
|
||||
.header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
|
||||
.header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
|
||||
.header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
|
||||
.send()
|
||||
.await
|
||||
.expect("request should succeed");
|
||||
assert_eq!(list_response.status(), StatusCode::OK);
|
||||
let list_payload: serde_json::Value =
|
||||
list_response.json().await.expect("json body should parse");
|
||||
assert_eq!(list_payload["items"].as_array().map(Vec::len), Some(1));
|
||||
assert_eq!(list_payload["items"][0]["id"], "collector-1");
|
||||
assert_eq!(list_payload["items"][0]["dimension_name"], "tokens");
|
||||
assert_eq!(list_payload["total"], 1);
|
||||
assert_eq!(list_payload["page"], 1);
|
||||
assert_eq!(list_payload["page_size"], 10);
|
||||
assert_eq!(list_payload["pages"], 1);
|
||||
|
||||
let detail_response = reqwest::Client::new()
|
||||
.get(format!(
|
||||
"{gateway_url}/api/admin/billing/collectors/collector-1"
|
||||
))
|
||||
.header(crate::constants::GATEWAY_HEADER, "rust-phase3b")
|
||||
.header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
|
||||
.header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
|
||||
.header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
|
||||
.send()
|
||||
.await
|
||||
.expect("request should succeed");
|
||||
assert_eq!(detail_response.status(), StatusCode::OK);
|
||||
let detail_payload: serde_json::Value = detail_response
|
||||
.json()
|
||||
.await
|
||||
.expect("json body should parse");
|
||||
assert_eq!(detail_payload["id"], "collector-1");
|
||||
assert_eq!(detail_payload["api_format"], "OPENAI");
|
||||
assert_eq!(detail_payload["source_path"], "usage.input_tokens");
|
||||
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
|
||||
gateway_handle.abort();
|
||||
upstream_handle.abort();
|
||||
|
||||
let upstream =
|
||||
Router::new().route(
|
||||
"/api/admin/billing/collectors",
|
||||
any(|_request: Request| async move {
|
||||
(StatusCode::OK, Body::from("unexpected upstream hit"))
|
||||
}),
|
||||
);
|
||||
let (upstream_url, upstream_handle) = start_server(upstream).await;
|
||||
let gateway = build_router_with_state(AppState::new().expect("gateway should build"));
|
||||
let (gateway_url, gateway_handle) = start_server(gateway).await;
|
||||
|
||||
let create_response = send_admin_billing_request(
|
||||
&gateway_url,
|
||||
http::Method::POST,
|
||||
"/api/admin/billing/collectors",
|
||||
Some(json!({
|
||||
"api_format": "openai",
|
||||
"task_type": "chat",
|
||||
"dimension_name": "tokens",
|
||||
"source_type": "request",
|
||||
"source_path": "usage.input_tokens",
|
||||
"value_type": "float",
|
||||
"priority": 10,
|
||||
"is_enabled": true
|
||||
})),
|
||||
)
|
||||
.await;
|
||||
assert_eq!(create_response.status(), StatusCode::OK);
|
||||
let create_payload: serde_json::Value = create_response
|
||||
.json()
|
||||
.await
|
||||
.expect("json body should parse");
|
||||
let collector_id = create_payload["id"]
|
||||
.as_str()
|
||||
.expect("collector id should exist")
|
||||
.to_string();
|
||||
assert_eq!(create_payload["api_format"], json!("OPENAI"));
|
||||
assert_eq!(create_payload["source_path"], json!("usage.input_tokens"));
|
||||
|
||||
let update_response = send_admin_billing_request(
|
||||
&gateway_url,
|
||||
http::Method::PUT,
|
||||
&format!("/api/admin/billing/collectors/{collector_id}"),
|
||||
Some(json!({
|
||||
"api_format": "openai",
|
||||
"task_type": "chat",
|
||||
"dimension_name": "tokens",
|
||||
"source_type": "computed",
|
||||
"source_path": null,
|
||||
"value_type": "int",
|
||||
"transform_expression": "max(total_tokens, 1)",
|
||||
"default_value": "1",
|
||||
"priority": 12,
|
||||
"is_enabled": false
|
||||
})),
|
||||
)
|
||||
.await;
|
||||
assert_eq!(update_response.status(), StatusCode::OK);
|
||||
let update_payload: serde_json::Value = update_response
|
||||
.json()
|
||||
.await
|
||||
.expect("json body should parse");
|
||||
assert_eq!(update_payload["id"], json!(collector_id));
|
||||
assert_eq!(update_payload["source_type"], json!("computed"));
|
||||
assert_eq!(update_payload["source_path"], serde_json::Value::Null);
|
||||
assert_eq!(
|
||||
update_payload["transform_expression"],
|
||||
json!("max(total_tokens, 1)")
|
||||
);
|
||||
assert_eq!(update_payload["default_value"], json!("1"));
|
||||
assert_eq!(update_payload["is_enabled"], json!(false));
|
||||
|
||||
gateway_handle.abort();
|
||||
upstream_handle.abort();
|
||||
}
|
||||
3
apps/aether-gateway/src/tests/control/admin/endpoints.rs
Normal file
3
apps/aether-gateway/src/tests/control/admin/endpoints.rs
Normal file
@@ -0,0 +1,3 @@
|
||||
mod keys;
|
||||
mod quota;
|
||||
mod routes;
|
||||
891
apps/aether-gateway/src/tests/control/admin/endpoints/keys.rs
Normal file
891
apps/aether-gateway/src/tests/control/admin/endpoints/keys.rs
Normal file
@@ -0,0 +1,891 @@
|
||||
use std::sync::{Arc, Mutex};
|
||||
|
||||
use aether_crypto::{
|
||||
decrypt_python_fernet_ciphertext, encrypt_python_fernet_plaintext, DEVELOPMENT_ENCRYPTION_KEY,
|
||||
};
|
||||
use aether_data::repository::provider_catalog::{
|
||||
InMemoryProviderCatalogReadRepository, ProviderCatalogReadRepository,
|
||||
};
|
||||
use axum::body::Body;
|
||||
use axum::routing::any;
|
||||
use axum::{extract::Request, Router};
|
||||
use http::StatusCode;
|
||||
use serde_json::json;
|
||||
|
||||
use super::super::super::{
|
||||
build_router_with_state, sample_endpoint, sample_key, sample_provider, start_server, AppState,
|
||||
};
|
||||
use crate::constants::{
|
||||
GATEWAY_HEADER, TRUSTED_ADMIN_SESSION_ID_HEADER, TRUSTED_ADMIN_USER_ID_HEADER,
|
||||
TRUSTED_ADMIN_USER_ROLE_HEADER,
|
||||
};
|
||||
use crate::data::GatewayDataState;
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_handles_admin_provider_keys_locally_with_trusted_admin_principal() {
|
||||
let upstream_hits = Arc::new(Mutex::new(0usize));
|
||||
let upstream_hits_clone = Arc::clone(&upstream_hits);
|
||||
let upstream = Router::new().route(
|
||||
"/api/admin/endpoints/providers/provider-openai/keys",
|
||||
any(move |_request: Request| {
|
||||
let upstream_hits_inner = Arc::clone(&upstream_hits_clone);
|
||||
async move {
|
||||
*upstream_hits_inner.lock().expect("mutex should lock") += 1;
|
||||
(StatusCode::OK, Body::from("unexpected upstream hit"))
|
||||
}
|
||||
}),
|
||||
);
|
||||
|
||||
let mut key_a = sample_key(
|
||||
"key-openai-a",
|
||||
"provider-openai",
|
||||
"openai:chat",
|
||||
"sk-test-a",
|
||||
);
|
||||
key_a.internal_priority = 10;
|
||||
key_a.request_count = Some(12);
|
||||
key_a.success_count = Some(9);
|
||||
key_a.error_count = Some(3);
|
||||
key_a.total_response_time_ms = Some(1800);
|
||||
key_a.created_at_unix_secs = Some(1_711_000_000);
|
||||
key_a.updated_at_unix_secs = Some(1_711_000_100);
|
||||
key_a.last_used_at_unix_secs = Some(1_711_000_120);
|
||||
key_a.note = Some("primary key".to_string());
|
||||
key_a.status_snapshot = Some(json!({
|
||||
"oauth": {"code": "none", "requires_reauth": false, "expiring_soon": false},
|
||||
"account": {"code": "ok", "blocked": false, "recoverable": false},
|
||||
"quota": {"code": "unknown", "exhausted": false}
|
||||
}));
|
||||
|
||||
let mut key_b = sample_key(
|
||||
"key-openai-b",
|
||||
"provider-openai",
|
||||
"openai:chat",
|
||||
"sk-test-b",
|
||||
);
|
||||
key_b.internal_priority = 20;
|
||||
key_b.request_count = Some(4);
|
||||
key_b.success_count = Some(4);
|
||||
key_b.error_count = Some(0);
|
||||
key_b.total_response_time_ms = Some(400);
|
||||
key_b.created_at_unix_secs = Some(1_711_100_000);
|
||||
key_b.updated_at_unix_secs = Some(1_711_100_100);
|
||||
|
||||
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
|
||||
vec![sample_provider("provider-openai", "openai", 10)],
|
||||
vec![],
|
||||
vec![key_a, key_b],
|
||||
));
|
||||
|
||||
let (upstream_url, upstream_handle) = start_server(upstream).await;
|
||||
let gateway = build_router_with_state(
|
||||
AppState::new()
|
||||
.expect("gateway should build")
|
||||
.with_data_state_for_tests(GatewayDataState::with_provider_catalog_reader_for_tests(
|
||||
provider_catalog_repository,
|
||||
)),
|
||||
);
|
||||
let (gateway_url, gateway_handle) = start_server(gateway).await;
|
||||
|
||||
let response = reqwest::Client::new()
|
||||
.get(format!(
|
||||
"{gateway_url}/api/admin/endpoints/providers/provider-openai/keys?skip=0&limit=50"
|
||||
))
|
||||
.header(crate::constants::GATEWAY_HEADER, "rust-phase3b")
|
||||
.header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
|
||||
.header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
|
||||
.header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
|
||||
.send()
|
||||
.await
|
||||
.expect("request should succeed");
|
||||
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
let payload: serde_json::Value = response.json().await.expect("json body should parse");
|
||||
let items = payload.as_array().expect("payload should be an array");
|
||||
assert_eq!(items.len(), 2);
|
||||
assert_eq!(items[0]["id"], "key-openai-a");
|
||||
assert_eq!(items[0]["internal_priority"], 10);
|
||||
assert_eq!(items[0]["request_count"], 12);
|
||||
assert_eq!(items[0]["success_count"], 9);
|
||||
assert_eq!(items[0]["error_count"], 3);
|
||||
assert_eq!(items[0]["note"], "primary key");
|
||||
assert_eq!(items[0]["api_key_masked"], "sk-test-a***");
|
||||
assert_eq!(items[1]["id"], "key-openai-b");
|
||||
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
|
||||
|
||||
gateway_handle.abort();
|
||||
upstream_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_creates_admin_provider_key_locally_with_trusted_admin_principal() {
|
||||
let upstream_hits = Arc::new(Mutex::new(0usize));
|
||||
let upstream_hits_clone = Arc::clone(&upstream_hits);
|
||||
let upstream = Router::new().route(
|
||||
"/api/admin/endpoints/providers/provider-openai/keys",
|
||||
any(move |_request: Request| {
|
||||
let upstream_hits_inner = Arc::clone(&upstream_hits_clone);
|
||||
async move {
|
||||
*upstream_hits_inner.lock().expect("mutex should lock") += 1;
|
||||
(StatusCode::OK, Body::from("unexpected upstream hit"))
|
||||
}
|
||||
}),
|
||||
);
|
||||
|
||||
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
|
||||
vec![sample_provider("provider-openai", "openai", 10)],
|
||||
vec![],
|
||||
vec![],
|
||||
));
|
||||
|
||||
let (upstream_url, upstream_handle) = start_server(upstream).await;
|
||||
let gateway = build_router_with_state(
|
||||
AppState::new()
|
||||
.expect("gateway should build")
|
||||
.with_data_state_for_tests(
|
||||
GatewayDataState::with_provider_catalog_repository_for_tests(
|
||||
provider_catalog_repository.clone(),
|
||||
)
|
||||
.with_encryption_key_for_tests(DEVELOPMENT_ENCRYPTION_KEY),
|
||||
),
|
||||
);
|
||||
let (gateway_url, gateway_handle) = start_server(gateway).await;
|
||||
|
||||
let response = reqwest::Client::new()
|
||||
.post(format!(
|
||||
"{gateway_url}/api/admin/endpoints/providers/provider-openai/keys"
|
||||
))
|
||||
.header(crate::constants::GATEWAY_HEADER, "rust-phase3b")
|
||||
.header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
|
||||
.header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
|
||||
.header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
|
||||
.json(&json!({
|
||||
"api_formats": ["openai:chat"],
|
||||
"api_key": "sk-created-openai",
|
||||
"name": "created key",
|
||||
"internal_priority": 15,
|
||||
"capabilities": {"cache_1h": true},
|
||||
"note": "created from rust"
|
||||
}))
|
||||
.send()
|
||||
.await
|
||||
.expect("request should succeed");
|
||||
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
let payload: serde_json::Value = response.json().await.expect("json body should parse");
|
||||
assert_eq!(payload["provider_id"], "provider-openai");
|
||||
assert_eq!(payload["name"], "created key");
|
||||
assert_eq!(payload["internal_priority"], 15);
|
||||
assert_eq!(payload["api_formats"], json!(["openai:chat"]));
|
||||
assert_eq!(payload["api_key_masked"], "sk-creat***enai");
|
||||
assert_eq!(payload["request_count"], 0);
|
||||
assert_eq!(payload["success_count"], 0);
|
||||
assert_eq!(payload["error_count"], 0);
|
||||
assert_eq!(payload["note"], "created from rust");
|
||||
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
|
||||
|
||||
let keys = provider_catalog_repository
|
||||
.list_keys_by_provider_ids(&["provider-openai".to_string()])
|
||||
.await
|
||||
.expect("keys should read");
|
||||
assert_eq!(keys.len(), 1);
|
||||
assert_eq!(keys[0].name, "created key");
|
||||
assert_eq!(keys[0].auth_type, "api_key");
|
||||
assert_eq!(keys[0].internal_priority, 15);
|
||||
|
||||
gateway_handle.abort();
|
||||
upstream_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_reveals_admin_provider_key_locally_with_trusted_admin_principal() {
|
||||
let upstream_hits = Arc::new(Mutex::new(0usize));
|
||||
let upstream_hits_clone = Arc::clone(&upstream_hits);
|
||||
let upstream = Router::new().route(
|
||||
"/api/admin/endpoints/keys/key-openai-a/reveal",
|
||||
any(move |_request: Request| {
|
||||
let upstream_hits_inner = Arc::clone(&upstream_hits_clone);
|
||||
async move {
|
||||
*upstream_hits_inner.lock().expect("mutex should lock") += 1;
|
||||
(StatusCode::OK, Body::from("unexpected upstream hit"))
|
||||
}
|
||||
}),
|
||||
);
|
||||
|
||||
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
|
||||
vec![sample_provider("provider-openai", "openai", 10)],
|
||||
vec![],
|
||||
vec![sample_key(
|
||||
"key-openai-a",
|
||||
"provider-openai",
|
||||
"openai:chat",
|
||||
"sk-test-a",
|
||||
)],
|
||||
));
|
||||
|
||||
let (upstream_url, upstream_handle) = start_server(upstream).await;
|
||||
let gateway = build_router_with_state(
|
||||
AppState::new()
|
||||
.expect("gateway should build")
|
||||
.with_data_state_for_tests(
|
||||
GatewayDataState::with_provider_catalog_repository_for_tests(
|
||||
provider_catalog_repository,
|
||||
)
|
||||
.with_encryption_key_for_tests(DEVELOPMENT_ENCRYPTION_KEY),
|
||||
),
|
||||
);
|
||||
let (gateway_url, gateway_handle) = start_server(gateway).await;
|
||||
|
||||
let response = reqwest::Client::new()
|
||||
.get(format!(
|
||||
"{gateway_url}/api/admin/endpoints/keys/key-openai-a/reveal"
|
||||
))
|
||||
.header(crate::constants::GATEWAY_HEADER, "rust-phase3b")
|
||||
.header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
|
||||
.header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
|
||||
.header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
|
||||
.send()
|
||||
.await
|
||||
.expect("request should succeed");
|
||||
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
let payload: serde_json::Value = response.json().await.expect("json body should parse");
|
||||
assert_eq!(payload["auth_type"], "api_key");
|
||||
assert_eq!(payload["api_key"], "sk-test-a");
|
||||
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
|
||||
|
||||
gateway_handle.abort();
|
||||
upstream_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_exports_admin_provider_key_locally_with_trusted_admin_principal() {
|
||||
let upstream_hits = Arc::new(Mutex::new(0usize));
|
||||
let upstream_hits_clone = Arc::clone(&upstream_hits);
|
||||
let upstream = Router::new().route(
|
||||
"/api/admin/endpoints/keys/key-kiro-a/export",
|
||||
any(move |_request: Request| {
|
||||
let upstream_hits_inner = Arc::clone(&upstream_hits_clone);
|
||||
async move {
|
||||
*upstream_hits_inner.lock().expect("mutex should lock") += 1;
|
||||
(StatusCode::OK, Body::from("unexpected upstream hit"))
|
||||
}
|
||||
}),
|
||||
);
|
||||
|
||||
let mut key = sample_key(
|
||||
"key-kiro-a",
|
||||
"provider-kiro",
|
||||
"claude:cli",
|
||||
"oauth-access-token",
|
||||
);
|
||||
key.auth_type = "oauth".to_string();
|
||||
key.encrypted_auth_config = Some(
|
||||
encrypt_python_fernet_plaintext(
|
||||
DEVELOPMENT_ENCRYPTION_KEY,
|
||||
r#"{"provider_type":"kiro","auth_method":"idc","refresh_token":"rt-kiro-123"}"#,
|
||||
)
|
||||
.expect("auth config ciphertext should build"),
|
||||
);
|
||||
key.upstream_metadata = Some(json!({"kiro": {"email": "alice@example.com"}}));
|
||||
|
||||
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
|
||||
vec![sample_provider("provider-kiro", "kiro", 10)],
|
||||
vec![],
|
||||
vec![key],
|
||||
));
|
||||
|
||||
let (upstream_url, upstream_handle) = start_server(upstream).await;
|
||||
let gateway = build_router_with_state(
|
||||
AppState::new()
|
||||
.expect("gateway should build")
|
||||
.with_data_state_for_tests(
|
||||
GatewayDataState::with_provider_catalog_reader_for_tests(
|
||||
provider_catalog_repository,
|
||||
)
|
||||
.with_encryption_key_for_tests(DEVELOPMENT_ENCRYPTION_KEY),
|
||||
),
|
||||
);
|
||||
let (gateway_url, gateway_handle) = start_server(gateway).await;
|
||||
|
||||
let response = reqwest::Client::new()
|
||||
.get(format!(
|
||||
"{gateway_url}/api/admin/endpoints/keys/key-kiro-a/export"
|
||||
))
|
||||
.header(crate::constants::GATEWAY_HEADER, "rust-phase3b")
|
||||
.header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
|
||||
.header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
|
||||
.header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
|
||||
.send()
|
||||
.await
|
||||
.expect("request should succeed");
|
||||
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
let payload: serde_json::Value = response.json().await.expect("json body should parse");
|
||||
assert_eq!(payload["provider_type"], "kiro");
|
||||
assert_eq!(payload["auth_method"], "idc");
|
||||
assert_eq!(payload["refresh_token"], "rt-kiro-123");
|
||||
assert_eq!(payload["email"], "alice@example.com");
|
||||
assert_eq!(payload["name"], "default");
|
||||
assert!(payload.get("exported_at").is_some());
|
||||
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
|
||||
|
||||
gateway_handle.abort();
|
||||
upstream_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_clears_admin_provider_key_oauth_invalid_locally_with_trusted_admin_principal() {
|
||||
let upstream_hits = Arc::new(Mutex::new(0usize));
|
||||
let upstream_hits_clone = Arc::clone(&upstream_hits);
|
||||
let upstream = Router::new().route(
|
||||
"/api/admin/endpoints/keys/key-openai-a/clear-oauth-invalid",
|
||||
any(move |_request: Request| {
|
||||
let upstream_hits_inner = Arc::clone(&upstream_hits_clone);
|
||||
async move {
|
||||
*upstream_hits_inner.lock().expect("mutex should lock") += 1;
|
||||
(StatusCode::OK, Body::from("unexpected upstream hit"))
|
||||
}
|
||||
}),
|
||||
);
|
||||
|
||||
let mut key = sample_key(
|
||||
"key-openai-a",
|
||||
"provider-openai",
|
||||
"openai:chat",
|
||||
"sk-test-a",
|
||||
);
|
||||
key.oauth_invalid_at_unix_secs = Some(1_710_000_000);
|
||||
key.oauth_invalid_reason = Some("token expired".to_string());
|
||||
|
||||
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
|
||||
vec![sample_provider("provider-openai", "openai", 10)],
|
||||
vec![],
|
||||
vec![key],
|
||||
));
|
||||
|
||||
let (upstream_url, upstream_handle) = start_server(upstream).await;
|
||||
let gateway = build_router_with_state(
|
||||
AppState::new()
|
||||
.expect("gateway should build")
|
||||
.with_data_state_for_tests(
|
||||
GatewayDataState::with_provider_catalog_repository_for_tests(
|
||||
provider_catalog_repository.clone(),
|
||||
)
|
||||
.with_encryption_key_for_tests(DEVELOPMENT_ENCRYPTION_KEY),
|
||||
),
|
||||
);
|
||||
let (gateway_url, gateway_handle) = start_server(gateway).await;
|
||||
|
||||
let response = reqwest::Client::new()
|
||||
.post(format!(
|
||||
"{gateway_url}/api/admin/endpoints/keys/key-openai-a/clear-oauth-invalid"
|
||||
))
|
||||
.header(crate::constants::GATEWAY_HEADER, "rust-phase3b")
|
||||
.header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
|
||||
.header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
|
||||
.header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
|
||||
.send()
|
||||
.await
|
||||
.expect("request should succeed");
|
||||
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
let payload: serde_json::Value = response.json().await.expect("json body should parse");
|
||||
assert_eq!(payload["message"], "已清除 OAuth 失效标记");
|
||||
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
|
||||
|
||||
let reloaded = provider_catalog_repository
|
||||
.list_keys_by_ids(&["key-openai-a".to_string()])
|
||||
.await
|
||||
.expect("keys should read");
|
||||
assert_eq!(reloaded.len(), 1);
|
||||
assert_eq!(reloaded[0].oauth_invalid_at_unix_secs, None);
|
||||
assert_eq!(reloaded[0].oauth_invalid_reason, None);
|
||||
|
||||
gateway_handle.abort();
|
||||
upstream_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_noops_admin_provider_key_oauth_invalid_clear_when_marker_absent() {
|
||||
let upstream_hits = Arc::new(Mutex::new(0usize));
|
||||
let upstream_hits_clone = Arc::clone(&upstream_hits);
|
||||
let upstream = Router::new().route(
|
||||
"/api/admin/endpoints/keys/key-openai-a/clear-oauth-invalid",
|
||||
any(move |_request: Request| {
|
||||
let upstream_hits_inner = Arc::clone(&upstream_hits_clone);
|
||||
async move {
|
||||
*upstream_hits_inner.lock().expect("mutex should lock") += 1;
|
||||
(StatusCode::OK, Body::from("unexpected upstream hit"))
|
||||
}
|
||||
}),
|
||||
);
|
||||
|
||||
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
|
||||
vec![sample_provider("provider-openai", "openai", 10)],
|
||||
vec![],
|
||||
vec![sample_key(
|
||||
"key-openai-a",
|
||||
"provider-openai",
|
||||
"openai:chat",
|
||||
"sk-test-a",
|
||||
)],
|
||||
));
|
||||
|
||||
let (upstream_url, upstream_handle) = start_server(upstream).await;
|
||||
let gateway = build_router_with_state(
|
||||
AppState::new()
|
||||
.expect("gateway should build")
|
||||
.with_data_state_for_tests(
|
||||
GatewayDataState::with_provider_catalog_repository_for_tests(
|
||||
provider_catalog_repository,
|
||||
)
|
||||
.with_encryption_key_for_tests(DEVELOPMENT_ENCRYPTION_KEY),
|
||||
),
|
||||
);
|
||||
let (gateway_url, gateway_handle) = start_server(gateway).await;
|
||||
|
||||
let response = reqwest::Client::new()
|
||||
.post(format!(
|
||||
"{gateway_url}/api/admin/endpoints/keys/key-openai-a/clear-oauth-invalid"
|
||||
))
|
||||
.header(crate::constants::GATEWAY_HEADER, "rust-phase3b")
|
||||
.header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
|
||||
.header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
|
||||
.header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
|
||||
.send()
|
||||
.await
|
||||
.expect("request should succeed");
|
||||
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
let payload: serde_json::Value = response.json().await.expect("json body should parse");
|
||||
assert_eq!(payload["message"], "该 Key 当前无失效标记,无需清除");
|
||||
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
|
||||
|
||||
gateway_handle.abort();
|
||||
upstream_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_updates_admin_provider_key_locally_with_trusted_admin_principal() {
|
||||
let upstream_hits = Arc::new(Mutex::new(0usize));
|
||||
let upstream_hits_clone = Arc::clone(&upstream_hits);
|
||||
let upstream = Router::new().route(
|
||||
"/api/admin/endpoints/keys/key-openai-a",
|
||||
any(move |_request: Request| {
|
||||
let upstream_hits_inner = Arc::clone(&upstream_hits_clone);
|
||||
async move {
|
||||
*upstream_hits_inner.lock().expect("mutex should lock") += 1;
|
||||
(StatusCode::OK, Body::from("unexpected upstream hit"))
|
||||
}
|
||||
}),
|
||||
);
|
||||
|
||||
let mut key = sample_key(
|
||||
"key-openai-a",
|
||||
"provider-openai",
|
||||
"openai:chat",
|
||||
"sk-test-a",
|
||||
);
|
||||
key.learned_rpm_limit = Some(88);
|
||||
key.allowed_models = Some(json!(["gpt-4.1"]));
|
||||
key.fingerprint = Some(json!({"user_agent": "old-ua"}));
|
||||
|
||||
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
|
||||
vec![sample_provider("provider-openai", "openai", 10)],
|
||||
vec![],
|
||||
vec![key],
|
||||
));
|
||||
|
||||
let (upstream_url, upstream_handle) = start_server(upstream).await;
|
||||
let gateway = build_router_with_state(
|
||||
AppState::new()
|
||||
.expect("gateway should build")
|
||||
.with_data_state_for_tests(
|
||||
GatewayDataState::with_provider_catalog_repository_for_tests(
|
||||
provider_catalog_repository.clone(),
|
||||
)
|
||||
.with_encryption_key_for_tests(DEVELOPMENT_ENCRYPTION_KEY),
|
||||
),
|
||||
);
|
||||
let (gateway_url, gateway_handle) = start_server(gateway).await;
|
||||
|
||||
let response = reqwest::Client::new()
|
||||
.put(format!(
|
||||
"{gateway_url}/api/admin/endpoints/keys/key-openai-a"
|
||||
))
|
||||
.header(crate::constants::GATEWAY_HEADER, "rust-phase3b")
|
||||
.header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
|
||||
.header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
|
||||
.header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
|
||||
.json(&json!({
|
||||
"api_key": "sk-updated-openai",
|
||||
"name": "updated key",
|
||||
"internal_priority": 15,
|
||||
"rpm_limit": null,
|
||||
"allowed_models": [],
|
||||
"note": "updated from rust",
|
||||
"is_active": false,
|
||||
"fingerprint": {"user_agent": "new-ua"}
|
||||
}))
|
||||
.send()
|
||||
.await
|
||||
.expect("request should succeed");
|
||||
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
let payload: serde_json::Value = response.json().await.expect("json body should parse");
|
||||
assert_eq!(payload["name"], "updated key");
|
||||
assert_eq!(payload["internal_priority"], 15);
|
||||
assert_eq!(payload["rpm_limit"], serde_json::Value::Null);
|
||||
assert_eq!(payload["learned_rpm_limit"], serde_json::Value::Null);
|
||||
assert_eq!(payload["allowed_models"], json!([]));
|
||||
assert_eq!(payload["note"], "updated from rust");
|
||||
assert_eq!(payload["is_active"], false);
|
||||
assert_eq!(payload["fingerprint"]["user_agent"], "new-ua");
|
||||
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
|
||||
|
||||
let reloaded = provider_catalog_repository
|
||||
.list_keys_by_ids(&["key-openai-a".to_string()])
|
||||
.await
|
||||
.expect("keys should read");
|
||||
assert_eq!(reloaded.len(), 1);
|
||||
assert_eq!(reloaded[0].name, "updated key");
|
||||
assert_eq!(reloaded[0].internal_priority, 15);
|
||||
assert_eq!(reloaded[0].rpm_limit, None);
|
||||
assert_eq!(reloaded[0].learned_rpm_limit, None);
|
||||
assert_eq!(reloaded[0].allowed_models, None);
|
||||
assert_eq!(reloaded[0].note.as_deref(), Some("updated from rust"));
|
||||
assert!(!reloaded[0].is_active);
|
||||
let decrypted = decrypt_python_fernet_ciphertext(
|
||||
DEVELOPMENT_ENCRYPTION_KEY,
|
||||
&reloaded[0].encrypted_api_key,
|
||||
)
|
||||
.expect("ciphertext should decrypt");
|
||||
assert_eq!(decrypted, "sk-updated-openai");
|
||||
|
||||
gateway_handle.abort();
|
||||
upstream_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_rejects_admin_provider_key_update_when_api_key_duplicates_existing_key() {
|
||||
let upstream_hits = Arc::new(Mutex::new(0usize));
|
||||
let upstream_hits_clone = Arc::clone(&upstream_hits);
|
||||
let upstream = Router::new().route(
|
||||
"/api/admin/endpoints/keys/key-openai-a",
|
||||
any(move |_request: Request| {
|
||||
let upstream_hits_inner = Arc::clone(&upstream_hits_clone);
|
||||
async move {
|
||||
*upstream_hits_inner.lock().expect("mutex should lock") += 1;
|
||||
(StatusCode::OK, Body::from("unexpected upstream hit"))
|
||||
}
|
||||
}),
|
||||
);
|
||||
|
||||
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
|
||||
vec![sample_provider("provider-openai", "openai", 10)],
|
||||
vec![],
|
||||
vec![
|
||||
sample_key(
|
||||
"key-openai-a",
|
||||
"provider-openai",
|
||||
"openai:chat",
|
||||
"sk-test-a",
|
||||
),
|
||||
sample_key(
|
||||
"key-openai-b",
|
||||
"provider-openai",
|
||||
"openai:chat",
|
||||
"sk-test-b",
|
||||
),
|
||||
],
|
||||
));
|
||||
|
||||
let (upstream_url, upstream_handle) = start_server(upstream).await;
|
||||
let gateway = build_router_with_state(
|
||||
AppState::new()
|
||||
.expect("gateway should build")
|
||||
.with_data_state_for_tests(
|
||||
GatewayDataState::with_provider_catalog_repository_for_tests(
|
||||
provider_catalog_repository,
|
||||
)
|
||||
.with_encryption_key_for_tests(DEVELOPMENT_ENCRYPTION_KEY),
|
||||
),
|
||||
);
|
||||
let (gateway_url, gateway_handle) = start_server(gateway).await;
|
||||
|
||||
let response = reqwest::Client::new()
|
||||
.put(format!(
|
||||
"{gateway_url}/api/admin/endpoints/keys/key-openai-a"
|
||||
))
|
||||
.header(crate::constants::GATEWAY_HEADER, "rust-phase3b")
|
||||
.header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
|
||||
.header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
|
||||
.header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
|
||||
.json(&json!({
|
||||
"api_key": "sk-test-b"
|
||||
}))
|
||||
.send()
|
||||
.await
|
||||
.expect("request should succeed");
|
||||
|
||||
assert_eq!(response.status(), StatusCode::BAD_REQUEST);
|
||||
let payload: serde_json::Value = response.json().await.expect("json body should parse");
|
||||
assert!(payload["detail"]
|
||||
.as_str()
|
||||
.expect("detail should be string")
|
||||
.contains("该 API Key 已存在于当前 Provider 中"));
|
||||
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
|
||||
|
||||
gateway_handle.abort();
|
||||
upstream_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_deletes_admin_provider_key_locally_with_trusted_admin_principal() {
|
||||
let upstream_hits = Arc::new(Mutex::new(0usize));
|
||||
let upstream_hits_clone = Arc::clone(&upstream_hits);
|
||||
let upstream = Router::new().route(
|
||||
"/api/admin/endpoints/keys/key-openai-a",
|
||||
any(move |_request: Request| {
|
||||
let upstream_hits_inner = Arc::clone(&upstream_hits_clone);
|
||||
async move {
|
||||
*upstream_hits_inner.lock().expect("mutex should lock") += 1;
|
||||
(StatusCode::OK, Body::from("unexpected upstream hit"))
|
||||
}
|
||||
}),
|
||||
);
|
||||
|
||||
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
|
||||
vec![sample_provider("provider-openai", "openai", 10)],
|
||||
vec![],
|
||||
vec![sample_key(
|
||||
"key-openai-a",
|
||||
"provider-openai",
|
||||
"openai:chat",
|
||||
"sk-test-a",
|
||||
)],
|
||||
));
|
||||
|
||||
let (upstream_url, upstream_handle) = start_server(upstream).await;
|
||||
let gateway = build_router_with_state(
|
||||
AppState::new()
|
||||
.expect("gateway should build")
|
||||
.with_data_state_for_tests(
|
||||
GatewayDataState::with_provider_catalog_repository_for_tests(
|
||||
provider_catalog_repository.clone(),
|
||||
)
|
||||
.with_encryption_key_for_tests(DEVELOPMENT_ENCRYPTION_KEY),
|
||||
),
|
||||
);
|
||||
let (gateway_url, gateway_handle) = start_server(gateway).await;
|
||||
|
||||
let response = reqwest::Client::new()
|
||||
.delete(format!(
|
||||
"{gateway_url}/api/admin/endpoints/keys/key-openai-a"
|
||||
))
|
||||
.header(crate::constants::GATEWAY_HEADER, "rust-phase3b")
|
||||
.header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
|
||||
.header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
|
||||
.header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
|
||||
.send()
|
||||
.await
|
||||
.expect("request should succeed");
|
||||
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
let payload: serde_json::Value = response.json().await.expect("json body should parse");
|
||||
assert_eq!(payload["message"], "Key key-openai-a 已删除");
|
||||
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
|
||||
|
||||
let reloaded = provider_catalog_repository
|
||||
.list_keys_by_ids(&["key-openai-a".to_string()])
|
||||
.await
|
||||
.expect("keys should read");
|
||||
assert!(reloaded.is_empty());
|
||||
|
||||
gateway_handle.abort();
|
||||
upstream_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_batch_deletes_admin_provider_keys_locally_with_trusted_admin_principal() {
|
||||
let upstream_hits = Arc::new(Mutex::new(0usize));
|
||||
let upstream_hits_clone = Arc::clone(&upstream_hits);
|
||||
let upstream = Router::new().route(
|
||||
"/api/admin/endpoints/keys/batch-delete",
|
||||
any(move |_request: Request| {
|
||||
let upstream_hits_inner = Arc::clone(&upstream_hits_clone);
|
||||
async move {
|
||||
*upstream_hits_inner.lock().expect("mutex should lock") += 1;
|
||||
(StatusCode::OK, Body::from("unexpected upstream hit"))
|
||||
}
|
||||
}),
|
||||
);
|
||||
|
||||
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
|
||||
vec![sample_provider("provider-openai", "openai", 10)],
|
||||
vec![],
|
||||
vec![
|
||||
sample_key(
|
||||
"key-openai-a",
|
||||
"provider-openai",
|
||||
"openai:chat",
|
||||
"sk-test-a",
|
||||
),
|
||||
sample_key(
|
||||
"key-openai-b",
|
||||
"provider-openai",
|
||||
"openai:chat",
|
||||
"sk-test-b",
|
||||
),
|
||||
],
|
||||
));
|
||||
|
||||
let (upstream_url, upstream_handle) = start_server(upstream).await;
|
||||
let gateway = build_router_with_state(
|
||||
AppState::new()
|
||||
.expect("gateway should build")
|
||||
.with_data_state_for_tests(
|
||||
GatewayDataState::with_provider_catalog_repository_for_tests(
|
||||
provider_catalog_repository.clone(),
|
||||
)
|
||||
.with_encryption_key_for_tests(DEVELOPMENT_ENCRYPTION_KEY),
|
||||
),
|
||||
);
|
||||
let (gateway_url, gateway_handle) = start_server(gateway).await;
|
||||
|
||||
let response = reqwest::Client::new()
|
||||
.post(format!(
|
||||
"{gateway_url}/api/admin/endpoints/keys/batch-delete"
|
||||
))
|
||||
.header(crate::constants::GATEWAY_HEADER, "rust-phase3b")
|
||||
.header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
|
||||
.header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
|
||||
.header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
|
||||
.json(&json!({
|
||||
"ids": ["key-openai-a", "key-missing"]
|
||||
}))
|
||||
.send()
|
||||
.await
|
||||
.expect("request should succeed");
|
||||
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
let payload: serde_json::Value = response.json().await.expect("json body should parse");
|
||||
assert_eq!(payload["success_count"], 1);
|
||||
assert_eq!(payload["failed_count"], 1);
|
||||
assert_eq!(payload["failed"][0]["id"], "key-missing");
|
||||
assert_eq!(payload["failed"][0]["error"], "not found");
|
||||
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
|
||||
|
||||
let reloaded = provider_catalog_repository
|
||||
.list_keys_by_provider_ids(&["provider-openai".to_string()])
|
||||
.await
|
||||
.expect("keys should read");
|
||||
assert_eq!(reloaded.len(), 1);
|
||||
assert_eq!(reloaded[0].id, "key-openai-b");
|
||||
|
||||
gateway_handle.abort();
|
||||
upstream_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_handles_admin_keys_grouped_by_format_locally_with_trusted_admin_principal() {
|
||||
let upstream_hits = Arc::new(Mutex::new(0usize));
|
||||
let upstream_hits_clone = Arc::clone(&upstream_hits);
|
||||
let upstream = Router::new().route(
|
||||
"/api/admin/endpoints/keys/grouped-by-format",
|
||||
any(move |_request: Request| {
|
||||
let upstream_hits_inner = Arc::clone(&upstream_hits_clone);
|
||||
async move {
|
||||
*upstream_hits_inner.lock().expect("mutex should lock") += 1;
|
||||
(StatusCode::OK, Body::from("unexpected upstream hit"))
|
||||
}
|
||||
}),
|
||||
);
|
||||
|
||||
let mut key_a = sample_key(
|
||||
"key-openai-a",
|
||||
"provider-openai",
|
||||
"openai:chat",
|
||||
"sk-test-a",
|
||||
);
|
||||
key_a.internal_priority = 10;
|
||||
key_a.request_count = Some(12);
|
||||
key_a.success_count = Some(9);
|
||||
key_a.created_at_unix_secs = Some(1_711_000_000);
|
||||
key_a.updated_at_unix_secs = Some(1_711_000_100);
|
||||
key_a.capabilities = Some(json!({"cache_1h": true, "gemini_files": false}));
|
||||
key_a.global_priority_by_format = Some(json!({"openai:chat": 3}));
|
||||
key_a.health_by_format = Some(json!({"openai:chat": {"health_score": 0.8}}));
|
||||
key_a.circuit_breaker_by_format = Some(json!({"openai:chat": {"open": false}}));
|
||||
|
||||
let mut key_b = sample_key("key-claude-a", "provider-claude", "claude:chat", "sk-ant-a");
|
||||
key_b.internal_priority = 20;
|
||||
key_b.request_count = Some(2);
|
||||
key_b.success_count = Some(1);
|
||||
key_b.created_at_unix_secs = Some(1_711_100_000);
|
||||
key_b.updated_at_unix_secs = Some(1_711_100_100);
|
||||
|
||||
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
|
||||
vec![
|
||||
sample_provider("provider-openai", "openai", 10),
|
||||
sample_provider("provider-claude", "claude", 20)
|
||||
.with_transport_fields(false, false, true, None, None, None, None, None, None),
|
||||
],
|
||||
vec![
|
||||
sample_endpoint(
|
||||
"endpoint-openai-chat",
|
||||
"provider-openai",
|
||||
"openai:chat",
|
||||
"https://api.openai.example",
|
||||
),
|
||||
sample_endpoint(
|
||||
"endpoint-claude-chat",
|
||||
"provider-claude",
|
||||
"claude:chat",
|
||||
"https://api.claude.example",
|
||||
),
|
||||
],
|
||||
vec![key_a, key_b],
|
||||
));
|
||||
|
||||
let (upstream_url, upstream_handle) = start_server(upstream).await;
|
||||
let gateway = build_router_with_state(
|
||||
AppState::new()
|
||||
.expect("gateway should build")
|
||||
.with_data_state_for_tests(
|
||||
GatewayDataState::with_provider_catalog_reader_for_tests(
|
||||
provider_catalog_repository,
|
||||
)
|
||||
.with_encryption_key_for_tests(DEVELOPMENT_ENCRYPTION_KEY),
|
||||
),
|
||||
);
|
||||
let (gateway_url, gateway_handle) = start_server(gateway).await;
|
||||
|
||||
let response = reqwest::Client::new()
|
||||
.get(format!(
|
||||
"{gateway_url}/api/admin/endpoints/keys/grouped-by-format"
|
||||
))
|
||||
.header(crate::constants::GATEWAY_HEADER, "rust-phase3b")
|
||||
.header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
|
||||
.header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
|
||||
.header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
|
||||
.send()
|
||||
.await
|
||||
.expect("request should succeed");
|
||||
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
let payload: serde_json::Value = response.json().await.expect("json body should parse");
|
||||
assert_eq!(payload["openai:chat"][0]["id"], "key-openai-a");
|
||||
assert_eq!(payload["openai:chat"][0]["provider_name"], "openai");
|
||||
assert_eq!(
|
||||
payload["openai:chat"][0]["endpoint_base_url"],
|
||||
"https://api.openai.example"
|
||||
);
|
||||
assert_eq!(payload["openai:chat"][0]["capabilities"], json!(["1h缓存"]));
|
||||
assert_eq!(payload["claude:chat"][0]["provider_active"], false);
|
||||
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
|
||||
|
||||
gateway_handle.abort();
|
||||
upstream_handle.abort();
|
||||
}
|
||||
644
apps/aether-gateway/src/tests/control/admin/endpoints/quota.rs
Normal file
644
apps/aether-gateway/src/tests/control/admin/endpoints/quota.rs
Normal file
@@ -0,0 +1,644 @@
|
||||
use std::collections::BTreeMap;
|
||||
use std::sync::{Arc, Mutex};
|
||||
|
||||
use aether_crypto::{encrypt_python_fernet_plaintext, DEVELOPMENT_ENCRYPTION_KEY};
|
||||
use aether_data::repository::provider_catalog::{
|
||||
InMemoryProviderCatalogReadRepository, ProviderCatalogReadRepository, StoredProviderCatalogKey,
|
||||
StoredProviderCatalogProvider,
|
||||
};
|
||||
use axum::body::{to_bytes, Body};
|
||||
use axum::routing::any;
|
||||
use axum::{extract::Request, Json, Router};
|
||||
use http::StatusCode;
|
||||
use serde_json::json;
|
||||
|
||||
use super::super::super::{
|
||||
build_router_with_state, build_state_with_execution_runtime_override, sample_endpoint,
|
||||
sample_key, start_server,
|
||||
};
|
||||
use crate::constants::{
|
||||
GATEWAY_HEADER, TRUSTED_ADMIN_SESSION_ID_HEADER, TRUSTED_ADMIN_USER_ID_HEADER,
|
||||
TRUSTED_ADMIN_USER_ROLE_HEADER,
|
||||
};
|
||||
use crate::data::GatewayDataState;
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_refreshes_admin_provider_quota_locally_for_codex_with_trusted_admin_principal() {
|
||||
#[derive(Debug, Clone)]
|
||||
struct SeenExecutionRuntimeRequest {
|
||||
url: String,
|
||||
authorization: String,
|
||||
provider_api_format: String,
|
||||
}
|
||||
|
||||
let upstream_hits = Arc::new(Mutex::new(0usize));
|
||||
let upstream_hits_clone = Arc::clone(&upstream_hits);
|
||||
let upstream = Router::new().route(
|
||||
"/api/admin/endpoints/providers/provider-codex/refresh-quota",
|
||||
any(move |_request: Request| {
|
||||
let upstream_hits_inner = Arc::clone(&upstream_hits_clone);
|
||||
async move {
|
||||
*upstream_hits_inner.lock().expect("mutex should lock") += 1;
|
||||
(StatusCode::OK, Body::from("unexpected upstream hit"))
|
||||
}
|
||||
}),
|
||||
);
|
||||
|
||||
let seen_execution_runtime = Arc::new(Mutex::new(None::<SeenExecutionRuntimeRequest>));
|
||||
let seen_execution_runtime_clone = Arc::clone(&seen_execution_runtime);
|
||||
let execution_runtime = Router::new().route(
|
||||
"/v1/execute/sync",
|
||||
any(move |request: Request| {
|
||||
let seen_execution_runtime_inner = Arc::clone(&seen_execution_runtime_clone);
|
||||
async move {
|
||||
let plan: aether_contracts::ExecutionPlan = serde_json::from_slice(
|
||||
&to_bytes(request.into_body(), usize::MAX)
|
||||
.await
|
||||
.expect("body should read"),
|
||||
)
|
||||
.expect("plan should parse");
|
||||
*seen_execution_runtime_inner
|
||||
.lock()
|
||||
.expect("mutex should lock") = Some(SeenExecutionRuntimeRequest {
|
||||
url: plan.url.clone(),
|
||||
authorization: plan
|
||||
.headers
|
||||
.get("authorization")
|
||||
.cloned()
|
||||
.unwrap_or_default(),
|
||||
provider_api_format: plan.provider_api_format.clone(),
|
||||
});
|
||||
let result = aether_contracts::ExecutionResult {
|
||||
request_id: plan.request_id,
|
||||
candidate_id: None,
|
||||
status_code: 200,
|
||||
headers: BTreeMap::new(),
|
||||
body: Some(aether_contracts::ResponseBody {
|
||||
json_body: Some(json!({
|
||||
"plan_type": "plus",
|
||||
"rate_limit": {
|
||||
"primary_window": {
|
||||
"used_percent": 12.5,
|
||||
"reset_after_seconds": 18000,
|
||||
"reset_at": 1_900_000_000u64,
|
||||
"window_minutes": 300
|
||||
},
|
||||
"secondary_window": {
|
||||
"used_percent": 55.0,
|
||||
"reset_after_seconds": 604800,
|
||||
"reset_at": 1_900_500_000u64,
|
||||
"window_minutes": 10080
|
||||
}
|
||||
},
|
||||
"credits": {
|
||||
"has_credits": true,
|
||||
"balance": 42.0,
|
||||
"unlimited": false
|
||||
}
|
||||
})),
|
||||
body_bytes_b64: None,
|
||||
}),
|
||||
telemetry: None,
|
||||
error: None,
|
||||
};
|
||||
(StatusCode::OK, Json(result))
|
||||
}
|
||||
}),
|
||||
);
|
||||
|
||||
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
|
||||
vec![StoredProviderCatalogProvider::new(
|
||||
"provider-codex".to_string(),
|
||||
"codex".to_string(),
|
||||
Some("https://example.com".to_string()),
|
||||
"codex".to_string(),
|
||||
)
|
||||
.expect("provider should build")],
|
||||
vec![sample_endpoint(
|
||||
"endpoint-codex-cli",
|
||||
"provider-codex",
|
||||
"openai:cli",
|
||||
"https://chatgpt.com/backend-api",
|
||||
)],
|
||||
vec![sample_key(
|
||||
"key-codex-a",
|
||||
"provider-codex",
|
||||
"openai:cli",
|
||||
"sk-codex-123",
|
||||
)],
|
||||
));
|
||||
|
||||
let (upstream_url, upstream_handle) = start_server(upstream).await;
|
||||
let (execution_runtime_url, execution_runtime_handle) = start_server(execution_runtime).await;
|
||||
let gateway = build_router_with_state(
|
||||
build_state_with_execution_runtime_override(execution_runtime_url.clone())
|
||||
.with_data_state_for_tests(
|
||||
GatewayDataState::with_provider_catalog_repository_for_tests(
|
||||
provider_catalog_repository.clone(),
|
||||
)
|
||||
.with_encryption_key_for_tests(DEVELOPMENT_ENCRYPTION_KEY),
|
||||
),
|
||||
);
|
||||
let (gateway_url, gateway_handle) = start_server(gateway).await;
|
||||
|
||||
let response = reqwest::Client::new()
|
||||
.post(format!(
|
||||
"{gateway_url}/api/admin/endpoints/providers/provider-codex/refresh-quota"
|
||||
))
|
||||
.header(crate::constants::GATEWAY_HEADER, "rust-phase3b")
|
||||
.header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
|
||||
.header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
|
||||
.header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
|
||||
.send()
|
||||
.await
|
||||
.expect("request should succeed");
|
||||
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
let payload: serde_json::Value = response.json().await.expect("json body should parse");
|
||||
assert_eq!(payload["success"], 1);
|
||||
assert_eq!(payload["failed"], 0);
|
||||
assert_eq!(payload["total"], 1);
|
||||
assert_eq!(payload["results"][0]["status"], "success");
|
||||
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
|
||||
|
||||
let seen_execution_runtime_request = seen_execution_runtime
|
||||
.lock()
|
||||
.expect("mutex should lock")
|
||||
.clone()
|
||||
.expect("execution runtime request should be captured");
|
||||
assert_eq!(
|
||||
seen_execution_runtime_request.url,
|
||||
"https://chatgpt.com/backend-api/wham/usage"
|
||||
);
|
||||
assert_eq!(
|
||||
seen_execution_runtime_request.authorization,
|
||||
"Bearer sk-codex-123"
|
||||
);
|
||||
assert_eq!(
|
||||
seen_execution_runtime_request.provider_api_format,
|
||||
"openai:cli"
|
||||
);
|
||||
|
||||
let reloaded = provider_catalog_repository
|
||||
.list_keys_by_ids(&["key-codex-a".to_string()])
|
||||
.await
|
||||
.expect("keys should read");
|
||||
assert_eq!(reloaded.len(), 1);
|
||||
assert_eq!(reloaded[0].oauth_invalid_reason, None);
|
||||
assert_eq!(
|
||||
reloaded[0]
|
||||
.upstream_metadata
|
||||
.as_ref()
|
||||
.and_then(|value| value.get("codex"))
|
||||
.and_then(|value| value.get("plan_type")),
|
||||
Some(&json!("plus"))
|
||||
);
|
||||
assert_eq!(
|
||||
reloaded[0]
|
||||
.upstream_metadata
|
||||
.as_ref()
|
||||
.and_then(|value| value.get("codex"))
|
||||
.and_then(|value| value.get("primary_used_percent")),
|
||||
Some(&json!(55.0))
|
||||
);
|
||||
assert_eq!(
|
||||
reloaded[0]
|
||||
.upstream_metadata
|
||||
.as_ref()
|
||||
.and_then(|value| value.get("codex"))
|
||||
.and_then(|value| value.get("secondary_used_percent")),
|
||||
Some(&json!(12.5))
|
||||
);
|
||||
|
||||
gateway_handle.abort();
|
||||
execution_runtime_handle.abort();
|
||||
upstream_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_refreshes_admin_provider_quota_locally_for_kiro_with_trusted_admin_principal() {
|
||||
#[derive(Debug, Clone)]
|
||||
struct SeenExecutionRuntimeRequest {
|
||||
url: String,
|
||||
authorization: String,
|
||||
provider_api_format: String,
|
||||
}
|
||||
|
||||
let upstream_hits = Arc::new(Mutex::new(0usize));
|
||||
let upstream_hits_clone = Arc::clone(&upstream_hits);
|
||||
let upstream = Router::new().route(
|
||||
"/api/admin/endpoints/providers/provider-kiro/refresh-quota",
|
||||
any(move |_request: Request| {
|
||||
let upstream_hits_inner = Arc::clone(&upstream_hits_clone);
|
||||
async move {
|
||||
*upstream_hits_inner.lock().expect("mutex should lock") += 1;
|
||||
(StatusCode::OK, Body::from("unexpected upstream hit"))
|
||||
}
|
||||
}),
|
||||
);
|
||||
|
||||
let seen_execution_runtime = Arc::new(Mutex::new(None::<SeenExecutionRuntimeRequest>));
|
||||
let seen_execution_runtime_clone = Arc::clone(&seen_execution_runtime);
|
||||
let execution_runtime = Router::new().route(
|
||||
"/v1/execute/sync",
|
||||
any(move |request: Request| {
|
||||
let seen_execution_runtime_inner = Arc::clone(&seen_execution_runtime_clone);
|
||||
async move {
|
||||
let plan: aether_contracts::ExecutionPlan = serde_json::from_slice(
|
||||
&to_bytes(request.into_body(), usize::MAX)
|
||||
.await
|
||||
.expect("body should read"),
|
||||
)
|
||||
.expect("plan should parse");
|
||||
*seen_execution_runtime_inner
|
||||
.lock()
|
||||
.expect("mutex should lock") = Some(SeenExecutionRuntimeRequest {
|
||||
url: plan.url.clone(),
|
||||
authorization: plan
|
||||
.headers
|
||||
.get("authorization")
|
||||
.cloned()
|
||||
.unwrap_or_default(),
|
||||
provider_api_format: plan.provider_api_format.clone(),
|
||||
});
|
||||
let result = aether_contracts::ExecutionResult {
|
||||
request_id: plan.request_id,
|
||||
candidate_id: None,
|
||||
status_code: 200,
|
||||
headers: BTreeMap::new(),
|
||||
body: Some(aether_contracts::ResponseBody {
|
||||
json_body: Some(json!({
|
||||
"subscriptionInfo": {
|
||||
"subscriptionTitle": "KIRO PRO+"
|
||||
},
|
||||
"usageBreakdownList": [{
|
||||
"currentUsageWithPrecision": 5.0,
|
||||
"usageLimitWithPrecision": 20.0,
|
||||
"nextDateReset": 1_900_000_000u64
|
||||
}],
|
||||
"desktopUserInfo": {
|
||||
"email": "dev@example.com"
|
||||
}
|
||||
})),
|
||||
body_bytes_b64: None,
|
||||
}),
|
||||
telemetry: None,
|
||||
error: None,
|
||||
};
|
||||
(StatusCode::OK, Json(result))
|
||||
}
|
||||
}),
|
||||
);
|
||||
|
||||
let encrypted_auth_config = encrypt_python_fernet_plaintext(
|
||||
DEVELOPMENT_ENCRYPTION_KEY,
|
||||
r#"{
|
||||
"access_token":"kiro-access-token",
|
||||
"api_region":"us-west-2",
|
||||
"machine_id":"123e4567-e89b-12d3-a456-426614174000",
|
||||
"kiro_version":"1.2.3"
|
||||
}"#,
|
||||
)
|
||||
.expect("auth config ciphertext should build");
|
||||
let encrypted_api_key =
|
||||
encrypt_python_fernet_plaintext(DEVELOPMENT_ENCRYPTION_KEY, "__placeholder__")
|
||||
.expect("api key ciphertext should build");
|
||||
let key = StoredProviderCatalogKey::new(
|
||||
"key-kiro-a".to_string(),
|
||||
"provider-kiro".to_string(),
|
||||
"default".to_string(),
|
||||
"bearer".to_string(),
|
||||
None,
|
||||
true,
|
||||
)
|
||||
.expect("key should build")
|
||||
.with_transport_fields(
|
||||
Some(json!(["claude:cli"])),
|
||||
encrypted_api_key,
|
||||
Some(encrypted_auth_config),
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
)
|
||||
.expect("key transport should build");
|
||||
|
||||
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
|
||||
vec![StoredProviderCatalogProvider::new(
|
||||
"provider-kiro".to_string(),
|
||||
"kiro".to_string(),
|
||||
Some("https://example.com".to_string()),
|
||||
"kiro".to_string(),
|
||||
)
|
||||
.expect("provider should build")],
|
||||
vec![sample_endpoint(
|
||||
"endpoint-kiro-cli",
|
||||
"provider-kiro",
|
||||
"claude:cli",
|
||||
"https://q.us-west-2.amazonaws.com",
|
||||
)],
|
||||
vec![key],
|
||||
));
|
||||
|
||||
let (upstream_url, upstream_handle) = start_server(upstream).await;
|
||||
let (execution_runtime_url, execution_runtime_handle) = start_server(execution_runtime).await;
|
||||
let gateway = build_router_with_state(
|
||||
build_state_with_execution_runtime_override(execution_runtime_url.clone())
|
||||
.with_data_state_for_tests(
|
||||
GatewayDataState::with_provider_catalog_repository_for_tests(
|
||||
provider_catalog_repository.clone(),
|
||||
)
|
||||
.with_encryption_key_for_tests(DEVELOPMENT_ENCRYPTION_KEY),
|
||||
),
|
||||
);
|
||||
let (gateway_url, gateway_handle) = start_server(gateway).await;
|
||||
|
||||
let response = reqwest::Client::new()
|
||||
.post(format!(
|
||||
"{gateway_url}/api/admin/endpoints/providers/provider-kiro/refresh-quota"
|
||||
))
|
||||
.header(crate::constants::GATEWAY_HEADER, "rust-phase3b")
|
||||
.header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
|
||||
.header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
|
||||
.header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
|
||||
.send()
|
||||
.await
|
||||
.expect("request should succeed");
|
||||
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
let payload: serde_json::Value = response.json().await.expect("json body should parse");
|
||||
assert_eq!(payload["success"], 1);
|
||||
assert_eq!(payload["failed"], 0);
|
||||
assert_eq!(payload["total"], 1);
|
||||
assert_eq!(payload["results"][0]["status"], "success");
|
||||
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
|
||||
|
||||
let seen_execution_runtime_request = seen_execution_runtime
|
||||
.lock()
|
||||
.expect("mutex should lock")
|
||||
.clone()
|
||||
.expect("execution runtime request should be captured");
|
||||
assert!(seen_execution_runtime_request
|
||||
.url
|
||||
.starts_with("https://q.us-west-2.amazonaws.com/getUsageLimits?"),);
|
||||
assert_eq!(
|
||||
seen_execution_runtime_request.authorization,
|
||||
"Bearer kiro-access-token"
|
||||
);
|
||||
assert_eq!(
|
||||
seen_execution_runtime_request.provider_api_format,
|
||||
"kiro:usage"
|
||||
);
|
||||
|
||||
let reloaded = provider_catalog_repository
|
||||
.list_keys_by_ids(&["key-kiro-a".to_string()])
|
||||
.await
|
||||
.expect("keys should read");
|
||||
assert_eq!(reloaded.len(), 1);
|
||||
assert_eq!(reloaded[0].oauth_invalid_reason, None);
|
||||
assert_eq!(
|
||||
reloaded[0]
|
||||
.upstream_metadata
|
||||
.as_ref()
|
||||
.and_then(|value| value.get("kiro"))
|
||||
.and_then(|value| value.get("subscription_title")),
|
||||
Some(&json!("KIRO PRO+"))
|
||||
);
|
||||
assert_eq!(
|
||||
reloaded[0]
|
||||
.upstream_metadata
|
||||
.as_ref()
|
||||
.and_then(|value| value.get("kiro"))
|
||||
.and_then(|value| value.get("remaining")),
|
||||
Some(&json!(15.0))
|
||||
);
|
||||
assert_eq!(
|
||||
reloaded[0]
|
||||
.upstream_metadata
|
||||
.as_ref()
|
||||
.and_then(|value| value.get("kiro"))
|
||||
.and_then(|value| value.get("email")),
|
||||
Some(&json!("dev@example.com"))
|
||||
);
|
||||
|
||||
gateway_handle.abort();
|
||||
execution_runtime_handle.abort();
|
||||
upstream_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_refreshes_admin_provider_quota_locally_for_antigravity_with_trusted_admin_principal(
|
||||
) {
|
||||
#[derive(Debug, Clone)]
|
||||
struct SeenExecutionRuntimeRequest {
|
||||
url: String,
|
||||
authorization: String,
|
||||
provider_api_format: String,
|
||||
request_body: Option<serde_json::Value>,
|
||||
}
|
||||
|
||||
let upstream_hits = Arc::new(Mutex::new(0usize));
|
||||
let upstream_hits_clone = Arc::clone(&upstream_hits);
|
||||
let upstream = Router::new().route(
|
||||
"/api/admin/endpoints/providers/provider-antigravity/refresh-quota",
|
||||
any(move |_request: Request| {
|
||||
let upstream_hits_inner = Arc::clone(&upstream_hits_clone);
|
||||
async move {
|
||||
*upstream_hits_inner.lock().expect("mutex should lock") += 1;
|
||||
(StatusCode::OK, Body::from("unexpected upstream hit"))
|
||||
}
|
||||
}),
|
||||
);
|
||||
|
||||
let seen_execution_runtime = Arc::new(Mutex::new(None::<SeenExecutionRuntimeRequest>));
|
||||
let seen_execution_runtime_clone = Arc::clone(&seen_execution_runtime);
|
||||
let execution_runtime = Router::new().route(
|
||||
"/v1/execute/sync",
|
||||
any(move |request: Request| {
|
||||
let seen_execution_runtime_inner = Arc::clone(&seen_execution_runtime_clone);
|
||||
async move {
|
||||
let plan: aether_contracts::ExecutionPlan = serde_json::from_slice(
|
||||
&to_bytes(request.into_body(), usize::MAX)
|
||||
.await
|
||||
.expect("body should read"),
|
||||
)
|
||||
.expect("plan should parse");
|
||||
*seen_execution_runtime_inner
|
||||
.lock()
|
||||
.expect("mutex should lock") = Some(SeenExecutionRuntimeRequest {
|
||||
url: plan.url.clone(),
|
||||
authorization: plan
|
||||
.headers
|
||||
.get("authorization")
|
||||
.cloned()
|
||||
.unwrap_or_default(),
|
||||
provider_api_format: plan.provider_api_format.clone(),
|
||||
request_body: plan.body.json_body.clone(),
|
||||
});
|
||||
let result = aether_contracts::ExecutionResult {
|
||||
request_id: plan.request_id,
|
||||
candidate_id: None,
|
||||
status_code: 200,
|
||||
headers: BTreeMap::new(),
|
||||
body: Some(aether_contracts::ResponseBody {
|
||||
json_body: Some(json!({
|
||||
"models": {
|
||||
"claude-sonnet-4": {
|
||||
"displayName": "Claude Sonnet 4",
|
||||
"quotaInfo": {
|
||||
"remainingFraction": 0.25,
|
||||
"resetTime": "2026-03-27T00:00:00Z"
|
||||
}
|
||||
},
|
||||
"gemini-2.5-pro": {
|
||||
"displayName": "Gemini 2.5 Pro"
|
||||
}
|
||||
}
|
||||
})),
|
||||
body_bytes_b64: None,
|
||||
}),
|
||||
telemetry: None,
|
||||
error: None,
|
||||
};
|
||||
(StatusCode::OK, Json(result))
|
||||
}
|
||||
}),
|
||||
);
|
||||
|
||||
let encrypted_auth_config = encrypt_python_fernet_plaintext(
|
||||
DEVELOPMENT_ENCRYPTION_KEY,
|
||||
r#"{
|
||||
"project_id":"project-ant-123",
|
||||
"client_version":"1.18.4",
|
||||
"session_id":"session-ant-1"
|
||||
}"#,
|
||||
)
|
||||
.expect("auth config ciphertext should build");
|
||||
let key = StoredProviderCatalogKey::new(
|
||||
"key-antigravity-a".to_string(),
|
||||
"provider-antigravity".to_string(),
|
||||
"default".to_string(),
|
||||
"oauth".to_string(),
|
||||
None,
|
||||
true,
|
||||
)
|
||||
.expect("key should build")
|
||||
.with_transport_fields(
|
||||
Some(json!(["gemini:chat"])),
|
||||
encrypt_python_fernet_plaintext(DEVELOPMENT_ENCRYPTION_KEY, "ya29.ant-token")
|
||||
.expect("api key ciphertext should build"),
|
||||
Some(encrypted_auth_config),
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
)
|
||||
.expect("key transport should build");
|
||||
|
||||
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
|
||||
vec![StoredProviderCatalogProvider::new(
|
||||
"provider-antigravity".to_string(),
|
||||
"antigravity".to_string(),
|
||||
Some("https://example.com".to_string()),
|
||||
"antigravity".to_string(),
|
||||
)
|
||||
.expect("provider should build")],
|
||||
vec![sample_endpoint(
|
||||
"endpoint-antigravity-chat",
|
||||
"provider-antigravity",
|
||||
"gemini:chat",
|
||||
"https://daily-cloudcode-pa.googleapis.com",
|
||||
)],
|
||||
vec![key],
|
||||
));
|
||||
|
||||
let (upstream_url, upstream_handle) = start_server(upstream).await;
|
||||
let (execution_runtime_url, execution_runtime_handle) = start_server(execution_runtime).await;
|
||||
let gateway = build_router_with_state(
|
||||
build_state_with_execution_runtime_override(execution_runtime_url.clone())
|
||||
.with_data_state_for_tests(
|
||||
GatewayDataState::with_provider_catalog_repository_for_tests(
|
||||
provider_catalog_repository.clone(),
|
||||
)
|
||||
.with_encryption_key_for_tests(DEVELOPMENT_ENCRYPTION_KEY),
|
||||
),
|
||||
);
|
||||
let (gateway_url, gateway_handle) = start_server(gateway).await;
|
||||
|
||||
let response = reqwest::Client::new()
|
||||
.post(format!(
|
||||
"{gateway_url}/api/admin/endpoints/providers/provider-antigravity/refresh-quota"
|
||||
))
|
||||
.header(crate::constants::GATEWAY_HEADER, "rust-phase3b")
|
||||
.header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
|
||||
.header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
|
||||
.header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
|
||||
.send()
|
||||
.await
|
||||
.expect("request should succeed");
|
||||
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
let payload: serde_json::Value = response.json().await.expect("json body should parse");
|
||||
assert_eq!(payload["success"], 1);
|
||||
assert_eq!(payload["failed"], 0);
|
||||
assert_eq!(payload["total"], 1);
|
||||
assert_eq!(payload["results"][0]["status"], "success");
|
||||
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
|
||||
|
||||
let seen_execution_runtime_request = seen_execution_runtime
|
||||
.lock()
|
||||
.expect("mutex should lock")
|
||||
.clone()
|
||||
.expect("execution runtime request should be captured");
|
||||
assert_eq!(
|
||||
seen_execution_runtime_request.url,
|
||||
"https://daily-cloudcode-pa.googleapis.com/v1internal:fetchAvailableModels"
|
||||
);
|
||||
assert_eq!(
|
||||
seen_execution_runtime_request.authorization,
|
||||
"Bearer ya29.ant-token"
|
||||
);
|
||||
assert_eq!(
|
||||
seen_execution_runtime_request.provider_api_format,
|
||||
"antigravity:fetch_available_models"
|
||||
);
|
||||
assert_eq!(
|
||||
seen_execution_runtime_request.request_body,
|
||||
Some(json!({ "project": "project-ant-123" }))
|
||||
);
|
||||
|
||||
let reloaded = provider_catalog_repository
|
||||
.list_keys_by_ids(&["key-antigravity-a".to_string()])
|
||||
.await
|
||||
.expect("keys should read");
|
||||
assert_eq!(reloaded.len(), 1);
|
||||
assert_eq!(reloaded[0].oauth_invalid_reason, None);
|
||||
assert_eq!(
|
||||
reloaded[0]
|
||||
.upstream_metadata
|
||||
.as_ref()
|
||||
.and_then(|value| value.get("antigravity"))
|
||||
.and_then(|value| value.get("models"))
|
||||
.and_then(|value| value.get("claude-sonnet-4"))
|
||||
.and_then(|value| value.get("remaining_fraction")),
|
||||
Some(&json!(0.25))
|
||||
);
|
||||
assert_eq!(
|
||||
reloaded[0]
|
||||
.upstream_metadata
|
||||
.as_ref()
|
||||
.and_then(|value| value.get("antigravity"))
|
||||
.and_then(|value| value.get("models"))
|
||||
.and_then(|value| value.get("claude-sonnet-4"))
|
||||
.and_then(|value| value.get("used_percent")),
|
||||
Some(&json!(75.0))
|
||||
);
|
||||
|
||||
gateway_handle.abort();
|
||||
execution_runtime_handle.abort();
|
||||
upstream_handle.abort();
|
||||
}
|
||||
609
apps/aether-gateway/src/tests/control/admin/endpoints/routes.rs
Normal file
609
apps/aether-gateway/src/tests/control/admin/endpoints/routes.rs
Normal file
@@ -0,0 +1,609 @@
|
||||
use std::sync::{Arc, Mutex};
|
||||
|
||||
use aether_data::repository::provider_catalog::{
|
||||
InMemoryProviderCatalogReadRepository, ProviderCatalogReadRepository,
|
||||
};
|
||||
use axum::body::Body;
|
||||
use axum::routing::any;
|
||||
use axum::{extract::Request, Router};
|
||||
use http::StatusCode;
|
||||
use serde_json::json;
|
||||
|
||||
use super::super::super::{
|
||||
build_router_with_state, sample_endpoint, sample_key, sample_provider, start_server, AppState,
|
||||
};
|
||||
use crate::constants::{
|
||||
GATEWAY_HEADER, TRUSTED_ADMIN_SESSION_ID_HEADER, TRUSTED_ADMIN_USER_ID_HEADER,
|
||||
TRUSTED_ADMIN_USER_ROLE_HEADER,
|
||||
};
|
||||
use crate::data::GatewayDataState;
|
||||
|
||||
const ADMIN_ENDPOINTS_DATA_UNAVAILABLE_DETAIL: &str = "Admin endpoint data unavailable";
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_returns_service_unavailable_for_admin_provider_endpoints_when_catalog_reader_unavailable(
|
||||
) {
|
||||
let upstream_hits = Arc::new(Mutex::new(0usize));
|
||||
let upstream_hits_clone = Arc::clone(&upstream_hits);
|
||||
let upstream = Router::new().route(
|
||||
"/api/admin/endpoints/providers/provider-openai/endpoints",
|
||||
any(move |_request: Request| {
|
||||
let upstream_hits_inner = Arc::clone(&upstream_hits_clone);
|
||||
async move {
|
||||
*upstream_hits_inner.lock().expect("mutex should lock") += 1;
|
||||
(StatusCode::OK, Body::from("unexpected upstream hit"))
|
||||
}
|
||||
}),
|
||||
);
|
||||
|
||||
let (upstream_url, upstream_handle) = start_server(upstream).await;
|
||||
let gateway = build_router_with_state(AppState::new().expect("gateway should build"));
|
||||
let (gateway_url, gateway_handle) = start_server(gateway).await;
|
||||
|
||||
let response = reqwest::Client::new()
|
||||
.get(format!(
|
||||
"{gateway_url}/api/admin/endpoints/providers/provider-openai/endpoints?skip=0&limit=50"
|
||||
))
|
||||
.header(crate::constants::GATEWAY_HEADER, "rust-phase3b")
|
||||
.header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
|
||||
.header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
|
||||
.header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
|
||||
.send()
|
||||
.await
|
||||
.expect("request should succeed");
|
||||
|
||||
assert_eq!(response.status(), StatusCode::SERVICE_UNAVAILABLE);
|
||||
let payload: serde_json::Value = response.json().await.expect("json body should parse");
|
||||
assert_eq!(payload["detail"], ADMIN_ENDPOINTS_DATA_UNAVAILABLE_DETAIL);
|
||||
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
|
||||
|
||||
gateway_handle.abort();
|
||||
upstream_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_handles_admin_provider_endpoints_locally_with_trusted_admin_principal() {
|
||||
let upstream_hits = Arc::new(Mutex::new(0usize));
|
||||
let upstream_hits_clone = Arc::clone(&upstream_hits);
|
||||
let upstream = Router::new().route(
|
||||
"/api/admin/endpoints/providers/provider-openai/endpoints",
|
||||
any(move |_request: Request| {
|
||||
let upstream_hits_inner = Arc::clone(&upstream_hits_clone);
|
||||
async move {
|
||||
*upstream_hits_inner.lock().expect("mutex should lock") += 1;
|
||||
(StatusCode::OK, Body::from("unexpected upstream hit"))
|
||||
}
|
||||
}),
|
||||
);
|
||||
|
||||
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
|
||||
vec![sample_provider("provider-openai", "openai", 10)],
|
||||
vec![
|
||||
sample_endpoint(
|
||||
"endpoint-openai-chat",
|
||||
"provider-openai",
|
||||
"openai:chat",
|
||||
"https://api.openai.example",
|
||||
)
|
||||
.with_timestamps(Some(1_711_000_000), Some(1_711_000_100)),
|
||||
sample_endpoint(
|
||||
"endpoint-openai-embed",
|
||||
"provider-openai",
|
||||
"openai:embedding",
|
||||
"https://api.openai.example",
|
||||
)
|
||||
.with_timestamps(Some(1_710_000_000), Some(1_710_000_100)),
|
||||
],
|
||||
vec![
|
||||
sample_key(
|
||||
"key-openai-a",
|
||||
"provider-openai",
|
||||
"openai:chat",
|
||||
"sk-test-a",
|
||||
),
|
||||
sample_key(
|
||||
"key-openai-b",
|
||||
"provider-openai",
|
||||
"openai:chat",
|
||||
"sk-test-b",
|
||||
),
|
||||
],
|
||||
));
|
||||
|
||||
let (upstream_url, upstream_handle) = start_server(upstream).await;
|
||||
let gateway = build_router_with_state(
|
||||
AppState::new()
|
||||
.expect("gateway should build")
|
||||
.with_data_state_for_tests(GatewayDataState::with_provider_catalog_reader_for_tests(
|
||||
provider_catalog_repository,
|
||||
)),
|
||||
);
|
||||
let (gateway_url, gateway_handle) = start_server(gateway).await;
|
||||
|
||||
let response = reqwest::Client::new()
|
||||
.get(format!(
|
||||
"{gateway_url}/api/admin/endpoints/providers/provider-openai/endpoints?skip=0&limit=50"
|
||||
))
|
||||
.header(crate::constants::GATEWAY_HEADER, "rust-phase3b")
|
||||
.header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
|
||||
.header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
|
||||
.header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
|
||||
.send()
|
||||
.await
|
||||
.expect("request should succeed");
|
||||
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
let payload: serde_json::Value = response.json().await.expect("json body should parse");
|
||||
let items = payload.as_array().expect("payload should be an array");
|
||||
assert_eq!(items.len(), 2);
|
||||
assert_eq!(items[0]["id"], "endpoint-openai-chat");
|
||||
assert_eq!(items[0]["provider_name"], "openai");
|
||||
assert_eq!(items[0]["api_format"], "openai:chat");
|
||||
assert_eq!(items[0]["total_keys"], 2);
|
||||
assert_eq!(items[0]["active_keys"], 2);
|
||||
assert_eq!(items[1]["id"], "endpoint-openai-embed");
|
||||
assert_eq!(items[1]["total_keys"], 0);
|
||||
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
|
||||
|
||||
gateway_handle.abort();
|
||||
upstream_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_returns_service_unavailable_for_admin_provider_endpoint_create_when_catalog_writer_unavailable(
|
||||
) {
|
||||
let upstream_hits = Arc::new(Mutex::new(0usize));
|
||||
let upstream_hits_clone = Arc::clone(&upstream_hits);
|
||||
let upstream = Router::new().route(
|
||||
"/api/admin/endpoints/providers/provider-openai/endpoints",
|
||||
any(move |_request: Request| {
|
||||
let upstream_hits_inner = Arc::clone(&upstream_hits_clone);
|
||||
async move {
|
||||
*upstream_hits_inner.lock().expect("mutex should lock") += 1;
|
||||
(StatusCode::OK, Body::from("unexpected upstream hit"))
|
||||
}
|
||||
}),
|
||||
);
|
||||
|
||||
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
|
||||
vec![sample_provider("provider-openai", "openai", 10)],
|
||||
vec![],
|
||||
vec![],
|
||||
));
|
||||
|
||||
let (upstream_url, upstream_handle) = start_server(upstream).await;
|
||||
let gateway = build_router_with_state(
|
||||
AppState::new()
|
||||
.expect("gateway should build")
|
||||
.with_data_state_for_tests(GatewayDataState::with_provider_catalog_reader_for_tests(
|
||||
provider_catalog_repository,
|
||||
)),
|
||||
);
|
||||
let (gateway_url, gateway_handle) = start_server(gateway).await;
|
||||
|
||||
let response = reqwest::Client::new()
|
||||
.post(format!(
|
||||
"{gateway_url}/api/admin/endpoints/providers/provider-openai/endpoints"
|
||||
))
|
||||
.header(crate::constants::GATEWAY_HEADER, "rust-phase3b")
|
||||
.header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
|
||||
.header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
|
||||
.header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
|
||||
.json(&json!({
|
||||
"provider_id": "provider-openai",
|
||||
"api_format": "openai:chat",
|
||||
"base_url": "https://api.openai.example/"
|
||||
}))
|
||||
.send()
|
||||
.await
|
||||
.expect("request should succeed");
|
||||
|
||||
assert_eq!(response.status(), StatusCode::SERVICE_UNAVAILABLE);
|
||||
let payload: serde_json::Value = response.json().await.expect("json body should parse");
|
||||
assert_eq!(payload["detail"], ADMIN_ENDPOINTS_DATA_UNAVAILABLE_DETAIL);
|
||||
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
|
||||
|
||||
gateway_handle.abort();
|
||||
upstream_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_creates_admin_provider_endpoint_locally_with_trusted_admin_principal() {
|
||||
let upstream_hits = Arc::new(Mutex::new(0usize));
|
||||
let upstream_hits_clone = Arc::clone(&upstream_hits);
|
||||
let upstream = Router::new().route(
|
||||
"/api/admin/endpoints/providers/provider-openai/endpoints",
|
||||
any(move |_request: Request| {
|
||||
let upstream_hits_inner = Arc::clone(&upstream_hits_clone);
|
||||
async move {
|
||||
*upstream_hits_inner.lock().expect("mutex should lock") += 1;
|
||||
(StatusCode::OK, Body::from("unexpected upstream hit"))
|
||||
}
|
||||
}),
|
||||
);
|
||||
|
||||
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
|
||||
vec![sample_provider("provider-openai", "openai", 10)],
|
||||
vec![],
|
||||
vec![],
|
||||
));
|
||||
|
||||
let (upstream_url, upstream_handle) = start_server(upstream).await;
|
||||
let gateway = build_router_with_state(
|
||||
AppState::new()
|
||||
.expect("gateway should build")
|
||||
.with_data_state_for_tests(
|
||||
GatewayDataState::with_provider_catalog_repository_for_tests(
|
||||
provider_catalog_repository.clone(),
|
||||
),
|
||||
),
|
||||
);
|
||||
let (gateway_url, gateway_handle) = start_server(gateway).await;
|
||||
|
||||
let response = reqwest::Client::new()
|
||||
.post(format!(
|
||||
"{gateway_url}/api/admin/endpoints/providers/provider-openai/endpoints"
|
||||
))
|
||||
.header(crate::constants::GATEWAY_HEADER, "rust-phase3b")
|
||||
.header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
|
||||
.header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
|
||||
.header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
|
||||
.json(&json!({
|
||||
"provider_id": "provider-openai",
|
||||
"api_format": "openai:chat",
|
||||
"base_url": "https://api.openai.example/",
|
||||
"custom_path": "/v1/chat/completions",
|
||||
"max_retries": 5,
|
||||
"config": {"foo": "bar"},
|
||||
"proxy": {"url": "http://proxy.internal", "password": "secret"}
|
||||
}))
|
||||
.send()
|
||||
.await
|
||||
.expect("request should succeed");
|
||||
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
let payload: serde_json::Value = response.json().await.expect("json body should parse");
|
||||
assert_eq!(payload["provider_id"], "provider-openai");
|
||||
assert_eq!(payload["provider_name"], "openai");
|
||||
assert_eq!(payload["api_format"], "openai:chat");
|
||||
assert_eq!(payload["base_url"], "https://api.openai.example");
|
||||
assert_eq!(payload["custom_path"], "/v1/chat/completions");
|
||||
assert_eq!(payload["max_retries"], 5);
|
||||
assert_eq!(payload["total_keys"], 0);
|
||||
assert_eq!(payload["active_keys"], 0);
|
||||
assert_eq!(payload["proxy"]["url"], "http://proxy.internal");
|
||||
assert_eq!(payload["proxy"]["password"], "***");
|
||||
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
|
||||
|
||||
let endpoints = provider_catalog_repository
|
||||
.list_endpoints_by_provider_ids(&["provider-openai".to_string()])
|
||||
.await
|
||||
.expect("endpoints should read");
|
||||
assert_eq!(endpoints.len(), 1);
|
||||
assert_eq!(endpoints[0].provider_id, "provider-openai");
|
||||
assert_eq!(endpoints[0].api_format, "openai:chat");
|
||||
assert_eq!(endpoints[0].base_url, "https://api.openai.example");
|
||||
assert_eq!(endpoints[0].max_retries, Some(5));
|
||||
|
||||
gateway_handle.abort();
|
||||
upstream_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_updates_admin_provider_endpoint_locally_with_trusted_admin_principal() {
|
||||
let upstream_hits = Arc::new(Mutex::new(0usize));
|
||||
let upstream_hits_clone = Arc::clone(&upstream_hits);
|
||||
let upstream = Router::new().route(
|
||||
"/api/admin/endpoints/endpoint-openai-chat",
|
||||
any(move |_request: Request| {
|
||||
let upstream_hits_inner = Arc::clone(&upstream_hits_clone);
|
||||
async move {
|
||||
*upstream_hits_inner.lock().expect("mutex should lock") += 1;
|
||||
(StatusCode::OK, Body::from("unexpected upstream hit"))
|
||||
}
|
||||
}),
|
||||
);
|
||||
|
||||
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
|
||||
vec![sample_provider("provider-openai", "openai", 10)],
|
||||
vec![sample_endpoint(
|
||||
"endpoint-openai-chat",
|
||||
"provider-openai",
|
||||
"openai:chat",
|
||||
"https://api.openai.example",
|
||||
)
|
||||
.with_transport_fields(
|
||||
"https://api.openai.example".to_string(),
|
||||
None,
|
||||
None,
|
||||
Some(2),
|
||||
Some("/v1/chat/completions".to_string()),
|
||||
Some(json!({"foo":"old"})),
|
||||
None,
|
||||
Some(json!({"url":"http://proxy.internal","password":"secret"})),
|
||||
)
|
||||
.expect("endpoint transport should build")],
|
||||
vec![sample_key(
|
||||
"key-openai-a",
|
||||
"provider-openai",
|
||||
"openai:chat",
|
||||
"sk-test-a",
|
||||
)],
|
||||
));
|
||||
|
||||
let (upstream_url, upstream_handle) = start_server(upstream).await;
|
||||
let gateway = build_router_with_state(
|
||||
AppState::new()
|
||||
.expect("gateway should build")
|
||||
.with_data_state_for_tests(
|
||||
GatewayDataState::with_provider_catalog_repository_for_tests(
|
||||
provider_catalog_repository.clone(),
|
||||
),
|
||||
),
|
||||
);
|
||||
let (gateway_url, gateway_handle) = start_server(gateway).await;
|
||||
|
||||
let response = reqwest::Client::new()
|
||||
.put(format!(
|
||||
"{gateway_url}/api/admin/endpoints/endpoint-openai-chat"
|
||||
))
|
||||
.header(crate::constants::GATEWAY_HEADER, "rust-phase3b")
|
||||
.header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
|
||||
.header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
|
||||
.header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
|
||||
.json(&json!({
|
||||
"base_url": "https://updated.openai.example/",
|
||||
"custom_path": "/v1/responses",
|
||||
"max_retries": 5,
|
||||
"is_active": false,
|
||||
"config": {"foo": "new"},
|
||||
"proxy": {"url": "http://proxy-2.internal"}
|
||||
}))
|
||||
.send()
|
||||
.await
|
||||
.expect("request should succeed");
|
||||
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
let payload: serde_json::Value = response.json().await.expect("json body should parse");
|
||||
assert_eq!(payload["id"], "endpoint-openai-chat");
|
||||
assert_eq!(payload["provider_id"], "provider-openai");
|
||||
assert_eq!(payload["api_format"], "openai:chat");
|
||||
assert_eq!(payload["base_url"], "https://updated.openai.example");
|
||||
assert_eq!(payload["custom_path"], "/v1/responses");
|
||||
assert_eq!(payload["max_retries"], 5);
|
||||
assert_eq!(payload["is_active"], false);
|
||||
assert_eq!(payload["total_keys"], 1);
|
||||
assert_eq!(payload["active_keys"], 1);
|
||||
assert_eq!(payload["proxy"]["url"], "http://proxy-2.internal");
|
||||
assert_eq!(payload["proxy"]["password"], "***");
|
||||
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
|
||||
|
||||
let endpoints = provider_catalog_repository
|
||||
.list_endpoints_by_ids(&["endpoint-openai-chat".to_string()])
|
||||
.await
|
||||
.expect("endpoints should read");
|
||||
assert_eq!(endpoints.len(), 1);
|
||||
assert_eq!(endpoints[0].base_url, "https://updated.openai.example");
|
||||
assert_eq!(endpoints[0].custom_path.as_deref(), Some("/v1/responses"));
|
||||
assert_eq!(endpoints[0].max_retries, Some(5));
|
||||
assert!(!endpoints[0].is_active);
|
||||
assert_eq!(endpoints[0].config, Some(json!({"foo":"new"})));
|
||||
assert_eq!(
|
||||
endpoints[0].proxy,
|
||||
Some(json!({"url":"http://proxy-2.internal","password":"secret"}))
|
||||
);
|
||||
|
||||
gateway_handle.abort();
|
||||
upstream_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_deletes_admin_provider_endpoint_locally_with_trusted_admin_principal() {
|
||||
let upstream_hits = Arc::new(Mutex::new(0usize));
|
||||
let upstream_hits_clone = Arc::clone(&upstream_hits);
|
||||
let upstream = Router::new().route(
|
||||
"/api/admin/endpoints/endpoint-openai-chat",
|
||||
any(move |_request: Request| {
|
||||
let upstream_hits_inner = Arc::clone(&upstream_hits_clone);
|
||||
async move {
|
||||
*upstream_hits_inner.lock().expect("mutex should lock") += 1;
|
||||
(StatusCode::OK, Body::from("unexpected upstream hit"))
|
||||
}
|
||||
}),
|
||||
);
|
||||
|
||||
let mut key_a = sample_key(
|
||||
"key-openai-a",
|
||||
"provider-openai",
|
||||
"openai:chat",
|
||||
"sk-test-a",
|
||||
);
|
||||
key_a.api_formats = Some(json!(["openai:chat", "openai:embedding"]));
|
||||
|
||||
let mut key_b = sample_key(
|
||||
"key-openai-b",
|
||||
"provider-openai",
|
||||
"openai:chat",
|
||||
"sk-test-b",
|
||||
);
|
||||
key_b.api_formats = Some(json!(["openai:chat"]));
|
||||
|
||||
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
|
||||
vec![sample_provider("provider-openai", "openai", 10)],
|
||||
vec![sample_endpoint(
|
||||
"endpoint-openai-chat",
|
||||
"provider-openai",
|
||||
"openai:chat",
|
||||
"https://api.openai.example",
|
||||
)],
|
||||
vec![key_a, key_b],
|
||||
));
|
||||
|
||||
let (upstream_url, upstream_handle) = start_server(upstream).await;
|
||||
let gateway = build_router_with_state(
|
||||
AppState::new()
|
||||
.expect("gateway should build")
|
||||
.with_data_state_for_tests(
|
||||
GatewayDataState::with_provider_catalog_repository_for_tests(
|
||||
provider_catalog_repository.clone(),
|
||||
),
|
||||
),
|
||||
);
|
||||
let (gateway_url, gateway_handle) = start_server(gateway).await;
|
||||
|
||||
let response = reqwest::Client::new()
|
||||
.delete(format!(
|
||||
"{gateway_url}/api/admin/endpoints/endpoint-openai-chat"
|
||||
))
|
||||
.header(crate::constants::GATEWAY_HEADER, "rust-phase3b")
|
||||
.header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
|
||||
.header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
|
||||
.header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
|
||||
.send()
|
||||
.await
|
||||
.expect("request should succeed");
|
||||
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
let payload: serde_json::Value = response.json().await.expect("json body should parse");
|
||||
assert_eq!(payload["message"], "Endpoint endpoint-openai-chat 已删除");
|
||||
assert_eq!(payload["affected_keys_count"], 2);
|
||||
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
|
||||
|
||||
let endpoints = provider_catalog_repository
|
||||
.list_endpoints_by_ids(&["endpoint-openai-chat".to_string()])
|
||||
.await
|
||||
.expect("endpoints should read");
|
||||
assert!(endpoints.is_empty());
|
||||
|
||||
let keys = provider_catalog_repository
|
||||
.list_keys_by_provider_ids(&["provider-openai".to_string()])
|
||||
.await
|
||||
.expect("keys should read");
|
||||
let key_a = keys
|
||||
.iter()
|
||||
.find(|key| key.id == "key-openai-a")
|
||||
.expect("key a should exist");
|
||||
let key_b = keys
|
||||
.iter()
|
||||
.find(|key| key.id == "key-openai-b")
|
||||
.expect("key b should exist");
|
||||
assert_eq!(key_a.api_formats, Some(json!(["openai:embedding"])));
|
||||
assert_eq!(key_b.api_formats, Some(json!([])));
|
||||
|
||||
gateway_handle.abort();
|
||||
upstream_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_handles_admin_get_endpoint_locally_with_trusted_admin_principal() {
|
||||
let upstream_hits = Arc::new(Mutex::new(0usize));
|
||||
let upstream_hits_clone = Arc::clone(&upstream_hits);
|
||||
let upstream = Router::new().route(
|
||||
"/api/admin/endpoints/endpoint-openai-chat",
|
||||
any(move |_request: Request| {
|
||||
let upstream_hits_inner = Arc::clone(&upstream_hits_clone);
|
||||
async move {
|
||||
*upstream_hits_inner.lock().expect("mutex should lock") += 1;
|
||||
(StatusCode::OK, Body::from("unexpected upstream hit"))
|
||||
}
|
||||
}),
|
||||
);
|
||||
|
||||
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
|
||||
vec![sample_provider("provider-openai", "openai", 10)],
|
||||
vec![sample_endpoint(
|
||||
"endpoint-openai-chat",
|
||||
"provider-openai",
|
||||
"openai:chat",
|
||||
"https://api.openai.example",
|
||||
)
|
||||
.with_timestamps(Some(1_711_000_000), Some(1_711_000_100))],
|
||||
vec![sample_key(
|
||||
"key-openai-a",
|
||||
"provider-openai",
|
||||
"openai:chat",
|
||||
"sk-test-a",
|
||||
)],
|
||||
));
|
||||
|
||||
let (upstream_url, upstream_handle) = start_server(upstream).await;
|
||||
let gateway = build_router_with_state(
|
||||
AppState::new()
|
||||
.expect("gateway should build")
|
||||
.with_data_state_for_tests(GatewayDataState::with_provider_catalog_reader_for_tests(
|
||||
provider_catalog_repository,
|
||||
)),
|
||||
);
|
||||
let (gateway_url, gateway_handle) = start_server(gateway).await;
|
||||
|
||||
let response = reqwest::Client::new()
|
||||
.get(format!(
|
||||
"{gateway_url}/api/admin/endpoints/endpoint-openai-chat"
|
||||
))
|
||||
.header(crate::constants::GATEWAY_HEADER, "rust-phase3b")
|
||||
.header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
|
||||
.header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
|
||||
.header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
|
||||
.send()
|
||||
.await
|
||||
.expect("request should succeed");
|
||||
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
let payload: serde_json::Value = response.json().await.expect("json body should parse");
|
||||
assert_eq!(payload["id"], "endpoint-openai-chat");
|
||||
assert_eq!(payload["provider_id"], "provider-openai");
|
||||
assert_eq!(payload["provider_name"], "openai");
|
||||
assert_eq!(payload["api_format"], "openai:chat");
|
||||
assert_eq!(payload["total_keys"], 1);
|
||||
assert_eq!(payload["active_keys"], 1);
|
||||
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
|
||||
|
||||
gateway_handle.abort();
|
||||
upstream_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_handles_admin_default_body_rules_locally_with_trusted_admin_principal() {
|
||||
let upstream_hits = Arc::new(Mutex::new(0usize));
|
||||
let upstream_hits_clone = Arc::clone(&upstream_hits);
|
||||
let upstream = Router::new().route(
|
||||
"/api/admin/endpoints/defaults/openai:cli/body-rules",
|
||||
any(move |_request: Request| {
|
||||
let upstream_hits_inner = Arc::clone(&upstream_hits_clone);
|
||||
async move {
|
||||
*upstream_hits_inner.lock().expect("mutex should lock") += 1;
|
||||
(StatusCode::OK, Body::from("unexpected upstream hit"))
|
||||
}
|
||||
}),
|
||||
);
|
||||
|
||||
let (upstream_url, upstream_handle) = start_server(upstream).await;
|
||||
let gateway = build_router_with_state(AppState::new().expect("gateway should build"));
|
||||
let (gateway_url, gateway_handle) = start_server(gateway).await;
|
||||
|
||||
let response = reqwest::Client::new()
|
||||
.get(format!(
|
||||
"{gateway_url}/api/admin/endpoints/defaults/openai:cli/body-rules?provider_type=codex"
|
||||
))
|
||||
.header(crate::constants::GATEWAY_HEADER, "rust-phase3b")
|
||||
.header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
|
||||
.header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
|
||||
.header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
|
||||
.send()
|
||||
.await
|
||||
.expect("request should succeed");
|
||||
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
let payload: serde_json::Value = response.json().await.expect("json body should parse");
|
||||
assert_eq!(payload["api_format"], "openai:cli");
|
||||
let rules = payload["body_rules"]
|
||||
.as_array()
|
||||
.expect("body_rules should be an array");
|
||||
assert_eq!(rules.len(), 5);
|
||||
assert_eq!(rules[0]["action"], "drop");
|
||||
assert_eq!(rules[0]["path"], "max_output_tokens");
|
||||
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
|
||||
|
||||
gateway_handle.abort();
|
||||
upstream_handle.abort();
|
||||
}
|
||||
349
apps/aether-gateway/src/tests/control/admin/gemini_files.rs
Normal file
349
apps/aether-gateway/src/tests/control/admin/gemini_files.rs
Normal file
@@ -0,0 +1,349 @@
|
||||
use aether_crypto::{encrypt_python_fernet_plaintext, DEVELOPMENT_ENCRYPTION_KEY};
|
||||
use aether_data::repository::candidates::InMemoryRequestCandidateRepository;
|
||||
use aether_data::repository::gemini_file_mappings::{
|
||||
GeminiFileMappingReadRepository, InMemoryGeminiFileMappingRepository,
|
||||
};
|
||||
use aether_data::repository::provider_catalog::{
|
||||
InMemoryProviderCatalogReadRepository, StoredProviderCatalogEndpoint, StoredProviderCatalogKey,
|
||||
StoredProviderCatalogProvider,
|
||||
};
|
||||
use axum::body::{to_bytes, Body};
|
||||
use axum::routing::any;
|
||||
use axum::{extract::Request, Json, Router};
|
||||
use base64::{engine::general_purpose::STANDARD as BASE64_STANDARD, Engine as _};
|
||||
use http::StatusCode;
|
||||
use serde_json::json;
|
||||
use std::sync::{Arc, Mutex};
|
||||
|
||||
use super::super::{
|
||||
build_router_with_state, build_state_with_execution_runtime_override, sample_provider,
|
||||
start_server,
|
||||
};
|
||||
use crate::constants::{
|
||||
GATEWAY_HEADER, TRACE_ID_HEADER, TRUSTED_ADMIN_SESSION_ID_HEADER, TRUSTED_ADMIN_USER_ID_HEADER,
|
||||
TRUSTED_ADMIN_USER_ROLE_HEADER,
|
||||
};
|
||||
use crate::data::GatewayDataState;
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
struct SeenAdminGeminiUploadExecution {
|
||||
key_id: String,
|
||||
method: String,
|
||||
url: String,
|
||||
auth_header_value: String,
|
||||
content_type: String,
|
||||
body_bytes_b64: String,
|
||||
}
|
||||
|
||||
fn sample_admin_gemini_provider() -> StoredProviderCatalogProvider {
|
||||
sample_provider("provider-gemini-admin-1", "gemini", 10).with_transport_fields(
|
||||
true,
|
||||
false,
|
||||
false,
|
||||
None,
|
||||
Some(2),
|
||||
None,
|
||||
Some(20.0),
|
||||
None,
|
||||
None,
|
||||
)
|
||||
}
|
||||
|
||||
fn sample_admin_gemini_endpoint() -> StoredProviderCatalogEndpoint {
|
||||
StoredProviderCatalogEndpoint::new(
|
||||
"endpoint-gemini-admin-1".to_string(),
|
||||
"provider-gemini-admin-1".to_string(),
|
||||
"gemini:chat".to_string(),
|
||||
Some("gemini".to_string()),
|
||||
Some("chat".to_string()),
|
||||
true,
|
||||
)
|
||||
.expect("endpoint should build")
|
||||
.with_transport_fields(
|
||||
"https://generativelanguage.googleapis.com".to_string(),
|
||||
None,
|
||||
None,
|
||||
Some(2),
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
)
|
||||
.expect("endpoint transport should build")
|
||||
}
|
||||
|
||||
fn sample_admin_gemini_key(id: &str, name: &str, secret: &str) -> StoredProviderCatalogKey {
|
||||
StoredProviderCatalogKey::new(
|
||||
id.to_string(),
|
||||
"provider-gemini-admin-1".to_string(),
|
||||
name.to_string(),
|
||||
"api_key".to_string(),
|
||||
Some(json!({"gemini_files": true})),
|
||||
true,
|
||||
)
|
||||
.expect("key should build")
|
||||
.with_transport_fields(
|
||||
Some(json!(["gemini:chat"])),
|
||||
encrypt_python_fernet_plaintext(DEVELOPMENT_ENCRYPTION_KEY, secret)
|
||||
.expect("api key should encrypt"),
|
||||
None,
|
||||
None,
|
||||
Some(json!({"gemini:chat": 1})),
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
)
|
||||
.expect("key transport should build")
|
||||
}
|
||||
|
||||
fn build_admin_gemini_upload_multipart(
|
||||
boundary: &str,
|
||||
file_name: &str,
|
||||
mime_type: &str,
|
||||
body: &[u8],
|
||||
) -> Vec<u8> {
|
||||
let mut payload = Vec::new();
|
||||
payload.extend_from_slice(
|
||||
format!(
|
||||
"--{boundary}\r\nContent-Disposition: form-data; name=\"file\"; filename=\"{file_name}\"\r\nContent-Type: {mime_type}\r\n\r\n"
|
||||
)
|
||||
.as_bytes(),
|
||||
);
|
||||
payload.extend_from_slice(body);
|
||||
payload.extend_from_slice(format!("\r\n--{boundary}--\r\n").as_bytes());
|
||||
payload
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_uploads_admin_gemini_file_locally_with_trusted_admin_principal() {
|
||||
let upstream_hits = Arc::new(Mutex::new(0usize));
|
||||
let upstream_hits_clone = Arc::clone(&upstream_hits);
|
||||
let upstream = Router::new().route(
|
||||
"/api/admin/gemini-files/upload",
|
||||
any(move |_request: Request| {
|
||||
let upstream_hits_inner = Arc::clone(&upstream_hits_clone);
|
||||
async move {
|
||||
*upstream_hits_inner.lock().expect("mutex should lock") += 1;
|
||||
(StatusCode::OK, Body::from("unexpected upstream hit"))
|
||||
}
|
||||
}),
|
||||
);
|
||||
|
||||
let seen_execution_runtime = Arc::new(Mutex::new(Vec::<SeenAdminGeminiUploadExecution>::new()));
|
||||
let seen_execution_runtime_clone = Arc::clone(&seen_execution_runtime);
|
||||
let execution_runtime = Router::new().route(
|
||||
"/v1/execute/sync",
|
||||
any(move |request: Request| {
|
||||
let seen_execution_runtime_inner = Arc::clone(&seen_execution_runtime_clone);
|
||||
async move {
|
||||
let (_parts, body) = request.into_parts();
|
||||
let raw_body = to_bytes(body, usize::MAX).await.expect("body should read");
|
||||
let payload: serde_json::Value = serde_json::from_slice(&raw_body)
|
||||
.expect("execution runtime payload should parse");
|
||||
let key_id = payload
|
||||
.get("key_id")
|
||||
.and_then(|value| value.as_str())
|
||||
.unwrap_or_default()
|
||||
.to_string();
|
||||
seen_execution_runtime_inner
|
||||
.lock()
|
||||
.expect("mutex should lock")
|
||||
.push(SeenAdminGeminiUploadExecution {
|
||||
key_id: key_id.clone(),
|
||||
method: payload
|
||||
.get("method")
|
||||
.and_then(|value| value.as_str())
|
||||
.unwrap_or_default()
|
||||
.to_string(),
|
||||
url: payload
|
||||
.get("url")
|
||||
.and_then(|value| value.as_str())
|
||||
.unwrap_or_default()
|
||||
.to_string(),
|
||||
auth_header_value: payload
|
||||
.get("headers")
|
||||
.and_then(|value| value.get("x-goog-api-key"))
|
||||
.and_then(|value| value.as_str())
|
||||
.unwrap_or_default()
|
||||
.to_string(),
|
||||
content_type: payload
|
||||
.get("headers")
|
||||
.and_then(|value| value.get("content-type"))
|
||||
.and_then(|value| value.as_str())
|
||||
.unwrap_or_default()
|
||||
.to_string(),
|
||||
body_bytes_b64: payload
|
||||
.get("body")
|
||||
.and_then(|value| value.get("body_bytes_b64"))
|
||||
.and_then(|value| value.as_str())
|
||||
.unwrap_or_default()
|
||||
.to_string(),
|
||||
});
|
||||
match key_id.as_str() {
|
||||
"key-gemini-admin-ok" => Json(json!({
|
||||
"request_id": "trace-admin-gemini-upload-123:ok",
|
||||
"status_code": 200,
|
||||
"headers": {
|
||||
"content-type": "application/json"
|
||||
},
|
||||
"body": {
|
||||
"json_body": {
|
||||
"file": {
|
||||
"name": "files/admin-upload-ok"
|
||||
}
|
||||
}
|
||||
}
|
||||
})),
|
||||
"key-gemini-admin-fail" => Json(json!({
|
||||
"request_id": "trace-admin-gemini-upload-123:fail",
|
||||
"status_code": 429,
|
||||
"headers": {
|
||||
"content-type": "application/json"
|
||||
},
|
||||
"body": {
|
||||
"json_body": {
|
||||
"error": {
|
||||
"message": "quota exceeded"
|
||||
}
|
||||
}
|
||||
}
|
||||
})),
|
||||
other => panic!("unexpected key id: {other}"),
|
||||
}
|
||||
}
|
||||
}),
|
||||
);
|
||||
|
||||
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
|
||||
vec![sample_admin_gemini_provider()],
|
||||
vec![sample_admin_gemini_endpoint()],
|
||||
vec![
|
||||
sample_admin_gemini_key("key-gemini-admin-ok", "ok-key", "sk-admin-gemini-ok"),
|
||||
sample_admin_gemini_key("key-gemini-admin-fail", "fail-key", "sk-admin-gemini-fail"),
|
||||
],
|
||||
));
|
||||
let gemini_file_mapping_repository = Arc::new(InMemoryGeminiFileMappingRepository::default());
|
||||
let request_candidate_repository = Arc::new(InMemoryRequestCandidateRepository::default());
|
||||
let data_state =
|
||||
GatewayDataState::with_request_candidate_and_gemini_file_mapping_repository_for_tests(
|
||||
request_candidate_repository,
|
||||
Arc::clone(&gemini_file_mapping_repository),
|
||||
)
|
||||
.attach_provider_catalog_repository_for_tests(Arc::clone(&provider_catalog_repository))
|
||||
.with_encryption_key_for_tests(DEVELOPMENT_ENCRYPTION_KEY);
|
||||
|
||||
let (upstream_url, upstream_handle) = start_server(upstream).await;
|
||||
let (execution_runtime_url, execution_runtime_handle) = start_server(execution_runtime).await;
|
||||
let gateway = build_router_with_state(
|
||||
build_state_with_execution_runtime_override(execution_runtime_url)
|
||||
.with_data_state_for_tests(data_state),
|
||||
);
|
||||
let (gateway_url, gateway_handle) = start_server(gateway).await;
|
||||
|
||||
let boundary = "----AetherAdminGeminiUploadBoundary";
|
||||
let upload_body = build_admin_gemini_upload_multipart(
|
||||
boundary,
|
||||
"admin-upload.txt",
|
||||
"text/plain",
|
||||
b"admin-upload-body",
|
||||
);
|
||||
let response = reqwest::Client::new()
|
||||
.post(format!(
|
||||
"{gateway_url}/api/admin/gemini-files/upload?key_ids=key-gemini-admin-ok,key-gemini-admin-fail"
|
||||
))
|
||||
.header(crate::constants::GATEWAY_HEADER, "rust-phase3b")
|
||||
.header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
|
||||
.header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
|
||||
.header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
|
||||
.header(
|
||||
http::header::CONTENT_TYPE,
|
||||
format!("multipart/form-data; boundary={boundary}"),
|
||||
)
|
||||
.header(TRACE_ID_HEADER, "trace-admin-gemini-upload-123")
|
||||
.body(upload_body)
|
||||
.send()
|
||||
.await
|
||||
.expect("request should succeed");
|
||||
|
||||
let status = response.status();
|
||||
let payload: serde_json::Value = response.json().await.expect("json body should parse");
|
||||
assert_eq!(status, StatusCode::OK, "payload={payload}");
|
||||
assert_eq!(payload["display_name"], "admin-upload.txt");
|
||||
assert_eq!(payload["mime_type"], "text/plain");
|
||||
assert_eq!(payload["size_bytes"], json!(17));
|
||||
assert_eq!(payload["success_count"], json!(1));
|
||||
assert_eq!(payload["fail_count"], json!(1));
|
||||
assert_eq!(payload["results"].as_array().map(Vec::len), Some(2));
|
||||
assert_eq!(payload["results"][0]["key_id"], "key-gemini-admin-ok");
|
||||
assert_eq!(payload["results"][0]["key_name"], "ok-key");
|
||||
assert_eq!(payload["results"][0]["success"], json!(true));
|
||||
assert_eq!(payload["results"][0]["file_name"], "files/admin-upload-ok");
|
||||
assert_eq!(payload["results"][0]["error"], serde_json::Value::Null);
|
||||
assert_eq!(payload["results"][1]["key_id"], "key-gemini-admin-fail");
|
||||
assert_eq!(payload["results"][1]["key_name"], "fail-key");
|
||||
assert_eq!(payload["results"][1]["success"], json!(false));
|
||||
assert_eq!(payload["results"][1]["file_name"], serde_json::Value::Null);
|
||||
assert_eq!(payload["results"][1]["error"], "quota exceeded");
|
||||
|
||||
let seen_requests = seen_execution_runtime
|
||||
.lock()
|
||||
.expect("mutex should lock")
|
||||
.clone();
|
||||
assert_eq!(seen_requests.len(), 2);
|
||||
for request in &seen_requests {
|
||||
assert_eq!(request.method, "POST");
|
||||
assert_eq!(
|
||||
request.url,
|
||||
"https://generativelanguage.googleapis.com/upload/v1beta/files?uploadType=resumable"
|
||||
);
|
||||
assert_eq!(request.content_type, "text/plain");
|
||||
assert_eq!(
|
||||
BASE64_STANDARD
|
||||
.decode(&request.body_bytes_b64)
|
||||
.expect("execution runtime body should decode"),
|
||||
b"admin-upload-body"
|
||||
);
|
||||
}
|
||||
let seen_by_key_id = seen_requests
|
||||
.iter()
|
||||
.map(|request| (request.key_id.as_str(), request))
|
||||
.collect::<std::collections::BTreeMap<_, _>>();
|
||||
assert_eq!(
|
||||
seen_by_key_id
|
||||
.get("key-gemini-admin-ok")
|
||||
.expect("ok request should exist")
|
||||
.auth_header_value,
|
||||
"sk-admin-gemini-ok"
|
||||
);
|
||||
assert_eq!(
|
||||
seen_by_key_id
|
||||
.get("key-gemini-admin-fail")
|
||||
.expect("fail request should exist")
|
||||
.auth_header_value,
|
||||
"sk-admin-gemini-fail"
|
||||
);
|
||||
|
||||
let stored_mapping = gemini_file_mapping_repository
|
||||
.find_by_file_name("files/admin-upload-ok")
|
||||
.await
|
||||
.expect("mapping lookup should succeed")
|
||||
.expect("mapping should exist");
|
||||
assert_eq!(stored_mapping.key_id, "key-gemini-admin-ok");
|
||||
assert_eq!(stored_mapping.user_id, None);
|
||||
assert_eq!(
|
||||
stored_mapping.display_name.as_deref(),
|
||||
Some("admin-upload.txt")
|
||||
);
|
||||
assert_eq!(stored_mapping.mime_type.as_deref(), Some("text/plain"));
|
||||
assert!(gemini_file_mapping_repository
|
||||
.find_by_file_name("files/admin-upload-fail")
|
||||
.await
|
||||
.expect("mapping lookup should succeed")
|
||||
.is_none());
|
||||
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
|
||||
|
||||
gateway_handle.abort();
|
||||
execution_runtime_handle.abort();
|
||||
upstream_handle.abort();
|
||||
}
|
||||
1035
apps/aether-gateway/src/tests/control/admin/health_access.rs
Normal file
1035
apps/aether-gateway/src/tests/control/admin/health_access.rs
Normal file
File diff suppressed because it is too large
Load Diff
225
apps/aether-gateway/src/tests/control/admin/ldap.rs
Normal file
225
apps/aether-gateway/src/tests/control/admin/ldap.rs
Normal file
@@ -0,0 +1,225 @@
|
||||
use std::sync::{Arc, Mutex};
|
||||
|
||||
use aether_data::repository::auth_modules::{
|
||||
InMemoryAuthModuleReadRepository, StoredOAuthProviderModuleConfig,
|
||||
};
|
||||
use axum::body::Body;
|
||||
use axum::routing::any;
|
||||
use axum::{extract::Request, Router};
|
||||
use http::StatusCode;
|
||||
use serde_json::json;
|
||||
|
||||
use super::super::{build_router_with_state, sample_ldap_module_config, start_server, AppState};
|
||||
use crate::constants::{
|
||||
GATEWAY_HEADER, TRUSTED_ADMIN_SESSION_ID_HEADER, TRUSTED_ADMIN_USER_ID_HEADER,
|
||||
TRUSTED_ADMIN_USER_ROLE_HEADER,
|
||||
};
|
||||
use crate::data::GatewayDataState;
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_handles_admin_ldap_config_locally_with_trusted_admin_principal() {
|
||||
let upstream_hits = Arc::new(Mutex::new(0usize));
|
||||
let upstream_hits_clone = Arc::clone(&upstream_hits);
|
||||
let upstream = Router::new().route(
|
||||
"/api/admin/ldap/config",
|
||||
any(move |_request: Request| {
|
||||
let upstream_hits_inner = Arc::clone(&upstream_hits_clone);
|
||||
async move {
|
||||
*upstream_hits_inner.lock().expect("mutex should lock") += 1;
|
||||
(StatusCode::OK, Body::from("unexpected upstream hit"))
|
||||
}
|
||||
}),
|
||||
);
|
||||
|
||||
let auth_module_repository = Arc::new(InMemoryAuthModuleReadRepository::seed(
|
||||
Vec::<StoredOAuthProviderModuleConfig>::new(),
|
||||
Some(sample_ldap_module_config()),
|
||||
));
|
||||
|
||||
let (upstream_url, upstream_handle) = start_server(upstream).await;
|
||||
let gateway = build_router_with_state(
|
||||
AppState::new()
|
||||
.expect("gateway should build")
|
||||
.with_data_state_for_tests(GatewayDataState::with_auth_module_repository_for_tests(
|
||||
auth_module_repository,
|
||||
)),
|
||||
);
|
||||
let (gateway_url, gateway_handle) = start_server(gateway).await;
|
||||
|
||||
let response = reqwest::Client::new()
|
||||
.get(format!("{gateway_url}/api/admin/ldap/config"))
|
||||
.header(crate::constants::GATEWAY_HEADER, "rust-phase3b")
|
||||
.header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
|
||||
.header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
|
||||
.header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
|
||||
.send()
|
||||
.await
|
||||
.expect("request should succeed");
|
||||
|
||||
let status = response.status();
|
||||
let payload: serde_json::Value = response.json().await.expect("json body should parse");
|
||||
assert_eq!(status, StatusCode::OK, "payload={payload}");
|
||||
assert_eq!(payload["server_url"], "ldaps://ldap.example.com");
|
||||
assert_eq!(payload["bind_dn"], "cn=admin,dc=example,dc=com");
|
||||
assert_eq!(payload["base_dn"], "dc=example,dc=com");
|
||||
assert_eq!(payload["has_bind_password"], json!(true));
|
||||
assert_eq!(payload["user_search_filter"], "(uid={username})");
|
||||
assert_eq!(payload["username_attr"], "uid");
|
||||
assert_eq!(payload["email_attr"], "mail");
|
||||
assert_eq!(payload["display_name_attr"], "displayName");
|
||||
assert_eq!(payload["is_enabled"], json!(true));
|
||||
assert_eq!(payload["is_exclusive"], json!(false));
|
||||
assert_eq!(payload["use_starttls"], json!(true));
|
||||
assert_eq!(payload["connect_timeout"], json!(10));
|
||||
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
|
||||
|
||||
gateway_handle.abort();
|
||||
upstream_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_updates_admin_ldap_config_locally_with_trusted_admin_principal() {
|
||||
let upstream_hits = Arc::new(Mutex::new(0usize));
|
||||
let upstream_hits_clone = Arc::clone(&upstream_hits);
|
||||
let upstream = Router::new().route(
|
||||
"/api/admin/ldap/config",
|
||||
any(move |_request: Request| {
|
||||
let upstream_hits_inner = Arc::clone(&upstream_hits_clone);
|
||||
async move {
|
||||
*upstream_hits_inner.lock().expect("mutex should lock") += 1;
|
||||
(StatusCode::OK, Body::from("unexpected upstream hit"))
|
||||
}
|
||||
}),
|
||||
);
|
||||
|
||||
let auth_module_repository = Arc::new(InMemoryAuthModuleReadRepository::seed(
|
||||
Vec::<StoredOAuthProviderModuleConfig>::new(),
|
||||
Some(sample_ldap_module_config()),
|
||||
));
|
||||
|
||||
let (upstream_url, upstream_handle) = start_server(upstream).await;
|
||||
let gateway = build_router_with_state(
|
||||
AppState::new()
|
||||
.expect("gateway should build")
|
||||
.with_data_state_for_tests(GatewayDataState::with_auth_module_repository_for_tests(
|
||||
auth_module_repository,
|
||||
)),
|
||||
);
|
||||
let (gateway_url, gateway_handle) = start_server(gateway).await;
|
||||
|
||||
let client = reqwest::Client::new();
|
||||
let update_response = client
|
||||
.put(format!("{gateway_url}/api/admin/ldap/config"))
|
||||
.header(crate::constants::GATEWAY_HEADER, "rust-phase3b")
|
||||
.header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
|
||||
.header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
|
||||
.header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
|
||||
.json(&json!({
|
||||
"server_url": "mockldap://ldap.internal.example.com",
|
||||
"bind_dn": "cn=svc,dc=example,dc=com",
|
||||
"bind_password": "secret123",
|
||||
"base_dn": "ou=people,dc=example,dc=com",
|
||||
"user_search_filter": "(uid={username})",
|
||||
"username_attr": "uid",
|
||||
"email_attr": "mail",
|
||||
"display_name_attr": "cn",
|
||||
"is_enabled": true,
|
||||
"is_exclusive": false,
|
||||
"use_starttls": false,
|
||||
"connect_timeout": 20
|
||||
}))
|
||||
.send()
|
||||
.await
|
||||
.expect("update request should succeed");
|
||||
|
||||
let status = update_response.status();
|
||||
let update_payload: serde_json::Value = update_response
|
||||
.json()
|
||||
.await
|
||||
.expect("json body should parse");
|
||||
assert_eq!(status, StatusCode::OK, "payload={update_payload}");
|
||||
assert_eq!(update_payload["message"], "LDAP配置更新成功");
|
||||
|
||||
let get_response = client
|
||||
.get(format!("{gateway_url}/api/admin/ldap/config"))
|
||||
.header(crate::constants::GATEWAY_HEADER, "rust-phase3b")
|
||||
.header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
|
||||
.header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
|
||||
.header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
|
||||
.send()
|
||||
.await
|
||||
.expect("request should succeed");
|
||||
|
||||
let status = get_response.status();
|
||||
let payload: serde_json::Value = get_response.json().await.expect("json body should parse");
|
||||
assert_eq!(status, StatusCode::OK, "payload={payload}");
|
||||
assert_eq!(
|
||||
payload["server_url"],
|
||||
"mockldap://ldap.internal.example.com"
|
||||
);
|
||||
assert_eq!(payload["bind_dn"], "cn=svc,dc=example,dc=com");
|
||||
assert_eq!(payload["base_dn"], "ou=people,dc=example,dc=com");
|
||||
assert_eq!(payload["display_name_attr"], "cn");
|
||||
assert_eq!(payload["has_bind_password"], json!(true));
|
||||
assert_eq!(payload["connect_timeout"], json!(20));
|
||||
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
|
||||
|
||||
gateway_handle.abort();
|
||||
upstream_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_tests_admin_ldap_connection_locally_with_trusted_admin_principal() {
|
||||
let upstream_hits = Arc::new(Mutex::new(0usize));
|
||||
let upstream_hits_clone = Arc::clone(&upstream_hits);
|
||||
let upstream = Router::new().route(
|
||||
"/api/admin/ldap/test",
|
||||
any(move |_request: Request| {
|
||||
let upstream_hits_inner = Arc::clone(&upstream_hits_clone);
|
||||
async move {
|
||||
*upstream_hits_inner.lock().expect("mutex should lock") += 1;
|
||||
(StatusCode::OK, Body::from("unexpected upstream hit"))
|
||||
}
|
||||
}),
|
||||
);
|
||||
|
||||
let auth_module_repository = Arc::new(InMemoryAuthModuleReadRepository::default());
|
||||
|
||||
let (upstream_url, upstream_handle) = start_server(upstream).await;
|
||||
let gateway = build_router_with_state(
|
||||
AppState::new()
|
||||
.expect("gateway should build")
|
||||
.with_data_state_for_tests(GatewayDataState::with_auth_module_repository_for_tests(
|
||||
auth_module_repository,
|
||||
)),
|
||||
);
|
||||
let (gateway_url, gateway_handle) = start_server(gateway).await;
|
||||
|
||||
let response = reqwest::Client::new()
|
||||
.post(format!("{gateway_url}/api/admin/ldap/test"))
|
||||
.header(crate::constants::GATEWAY_HEADER, "rust-phase3b")
|
||||
.header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
|
||||
.header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
|
||||
.header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
|
||||
.json(&json!({
|
||||
"server_url": "mockldap://ldap.example.com",
|
||||
"bind_dn": "cn=admin,dc=example,dc=com",
|
||||
"bind_password": "secret123",
|
||||
"base_dn": "dc=example,dc=com",
|
||||
"use_starttls": false,
|
||||
"connect_timeout": 10
|
||||
}))
|
||||
.send()
|
||||
.await
|
||||
.expect("request should succeed");
|
||||
|
||||
let status = response.status();
|
||||
let payload: serde_json::Value = response.json().await.expect("json body should parse");
|
||||
assert_eq!(status, StatusCode::OK, "payload={payload}");
|
||||
assert_eq!(payload["success"], json!(true));
|
||||
assert_eq!(payload["message"], "连接成功");
|
||||
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
|
||||
|
||||
gateway_handle.abort();
|
||||
upstream_handle.abort();
|
||||
}
|
||||
@@ -0,0 +1,92 @@
|
||||
use std::sync::{Arc, Mutex};
|
||||
|
||||
use axum::body::Body;
|
||||
use axum::routing::any;
|
||||
use axum::{extract::Request, Router};
|
||||
use http::StatusCode;
|
||||
|
||||
use super::super::super::{build_router_with_state, start_server, AppState};
|
||||
use crate::constants::{
|
||||
GATEWAY_HEADER, TRUSTED_ADMIN_SESSION_ID_HEADER, TRUSTED_ADMIN_USER_ID_HEADER,
|
||||
TRUSTED_ADMIN_USER_ROLE_HEADER,
|
||||
};
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_handles_admin_external_models_locally_with_trusted_admin_principal() {
|
||||
let upstream_hits = Arc::new(Mutex::new(0usize));
|
||||
let upstream_hits_clone = Arc::clone(&upstream_hits);
|
||||
let upstream = Router::new().route(
|
||||
"/api/admin/models/external",
|
||||
any(move |_request: Request| {
|
||||
let upstream_hits_inner = Arc::clone(&upstream_hits_clone);
|
||||
async move {
|
||||
*upstream_hits_inner.lock().expect("mutex should lock") += 1;
|
||||
(StatusCode::OK, Body::from("unexpected upstream hit"))
|
||||
}
|
||||
}),
|
||||
);
|
||||
|
||||
let (upstream_url, upstream_handle) = start_server(upstream).await;
|
||||
let gateway = build_router_with_state(AppState::new().expect("gateway should build"));
|
||||
let (gateway_url, gateway_handle) = start_server(gateway).await;
|
||||
|
||||
let response = reqwest::Client::new()
|
||||
.get(format!("{gateway_url}/api/admin/models/external"))
|
||||
.header(crate::constants::GATEWAY_HEADER, "rust-phase3b")
|
||||
.header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
|
||||
.header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
|
||||
.header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
|
||||
.send()
|
||||
.await
|
||||
.expect("request should succeed");
|
||||
|
||||
assert_eq!(response.status(), StatusCode::SERVICE_UNAVAILABLE);
|
||||
let payload: serde_json::Value = response.json().await.expect("json body should parse");
|
||||
assert_eq!(
|
||||
payload["detail"],
|
||||
"External models catalog requires Rust admin backend"
|
||||
);
|
||||
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
|
||||
|
||||
gateway_handle.abort();
|
||||
upstream_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_clears_admin_external_models_cache_locally_with_trusted_admin_principal() {
|
||||
let upstream_hits = Arc::new(Mutex::new(0usize));
|
||||
let upstream_hits_clone = Arc::clone(&upstream_hits);
|
||||
let upstream = Router::new().route(
|
||||
"/api/admin/models/external/cache",
|
||||
any(move |_request: Request| {
|
||||
let upstream_hits_inner = Arc::clone(&upstream_hits_clone);
|
||||
async move {
|
||||
*upstream_hits_inner.lock().expect("mutex should lock") += 1;
|
||||
(StatusCode::OK, Body::from("unexpected upstream hit"))
|
||||
}
|
||||
}),
|
||||
);
|
||||
|
||||
let (upstream_url, upstream_handle) = start_server(upstream).await;
|
||||
let gateway = build_router_with_state(AppState::new().expect("gateway should build"));
|
||||
let (gateway_url, gateway_handle) = start_server(gateway).await;
|
||||
|
||||
let response = reqwest::Client::new()
|
||||
.delete(format!("{gateway_url}/api/admin/models/external/cache"))
|
||||
.header(crate::constants::GATEWAY_HEADER, "rust-phase3b")
|
||||
.header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
|
||||
.header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
|
||||
.header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
|
||||
.send()
|
||||
.await
|
||||
.expect("request should succeed");
|
||||
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
let payload: serde_json::Value = response.json().await.expect("json body should parse");
|
||||
assert_eq!(payload["cleared"], false);
|
||||
assert_eq!(payload["message"], "Redis 未启用");
|
||||
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
|
||||
|
||||
gateway_handle.abort();
|
||||
upstream_handle.abort();
|
||||
}
|
||||
1108
apps/aether-gateway/src/tests/control/admin/models/global.rs
Normal file
1108
apps/aether-gateway/src/tests/control/admin/models/global.rs
Normal file
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,3 @@
|
||||
mod external;
|
||||
mod global;
|
||||
mod provider;
|
||||
639
apps/aether-gateway/src/tests/control/admin/models/provider.rs
Normal file
639
apps/aether-gateway/src/tests/control/admin/models/provider.rs
Normal file
@@ -0,0 +1,639 @@
|
||||
use std::sync::{Arc, Mutex};
|
||||
|
||||
use aether_data::repository::global_models::{
|
||||
AdminProviderModelListQuery, GlobalModelReadRepository, InMemoryGlobalModelReadRepository,
|
||||
};
|
||||
use aether_data::repository::provider_catalog::InMemoryProviderCatalogReadRepository;
|
||||
use aether_data::repository::quota::InMemoryProviderQuotaRepository;
|
||||
use axum::body::Body;
|
||||
use axum::routing::any;
|
||||
use axum::{extract::Request, Router};
|
||||
use http::StatusCode;
|
||||
use serde_json::json;
|
||||
|
||||
use super::super::super::{
|
||||
build_router_with_state, sample_admin_global_model, sample_admin_provider_model,
|
||||
sample_provider, start_server, AppState,
|
||||
};
|
||||
use crate::constants::{
|
||||
GATEWAY_HEADER, TRUSTED_ADMIN_SESSION_ID_HEADER, TRUSTED_ADMIN_USER_ID_HEADER,
|
||||
TRUSTED_ADMIN_USER_ROLE_HEADER,
|
||||
};
|
||||
use crate::data::GatewayDataState;
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_handles_admin_provider_models_locally_with_trusted_admin_principal() {
|
||||
let upstream_hits = Arc::new(Mutex::new(0usize));
|
||||
let upstream_hits_clone = Arc::clone(&upstream_hits);
|
||||
let upstream = Router::new().route(
|
||||
"/api/admin/providers/provider-openai/models",
|
||||
any(move |_request: Request| {
|
||||
let upstream_hits_inner = Arc::clone(&upstream_hits_clone);
|
||||
async move {
|
||||
*upstream_hits_inner.lock().expect("mutex should lock") += 1;
|
||||
(StatusCode::OK, Body::from("unexpected upstream hit"))
|
||||
}
|
||||
}),
|
||||
);
|
||||
|
||||
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
|
||||
vec![sample_provider("provider-openai", "openai", 10)],
|
||||
Vec::new(),
|
||||
Vec::new(),
|
||||
));
|
||||
let global_model_repository = Arc::new(
|
||||
InMemoryGlobalModelReadRepository::seed(Vec::new()).with_admin_provider_models(vec![
|
||||
sample_admin_provider_model(
|
||||
"model-openai-gpt5",
|
||||
"provider-openai",
|
||||
"global-gpt-5",
|
||||
"gpt-5-upstream",
|
||||
),
|
||||
]),
|
||||
);
|
||||
let quota_repository = Arc::new(InMemoryProviderQuotaRepository::seed(vec![]));
|
||||
|
||||
let (upstream_url, upstream_handle) = start_server(upstream).await;
|
||||
let gateway = build_router_with_state(
|
||||
AppState::new()
|
||||
.expect("gateway should build")
|
||||
.with_data_state_for_tests(
|
||||
GatewayDataState::with_provider_catalog_global_model_and_quota_readers_for_tests(
|
||||
provider_catalog_repository,
|
||||
global_model_repository,
|
||||
quota_repository,
|
||||
),
|
||||
),
|
||||
);
|
||||
let (gateway_url, gateway_handle) = start_server(gateway).await;
|
||||
|
||||
let response = reqwest::Client::new()
|
||||
.get(format!(
|
||||
"{gateway_url}/api/admin/providers/provider-openai/models?skip=0&limit=20"
|
||||
))
|
||||
.header(crate::constants::GATEWAY_HEADER, "rust-phase3b")
|
||||
.header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
|
||||
.header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
|
||||
.header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
|
||||
.send()
|
||||
.await
|
||||
.expect("request should succeed");
|
||||
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
let payload: serde_json::Value = response.json().await.expect("json body should parse");
|
||||
let items = payload.as_array().expect("payload should be an array");
|
||||
assert_eq!(items.len(), 1);
|
||||
assert_eq!(items[0]["id"], "model-openai-gpt5");
|
||||
assert_eq!(items[0]["provider_id"], "provider-openai");
|
||||
assert_eq!(items[0]["global_model_id"], "global-gpt-5");
|
||||
assert_eq!(items[0]["provider_model_name"], "gpt-5-upstream");
|
||||
assert_eq!(items[0]["global_model_name"], "gpt-5");
|
||||
assert_eq!(items[0]["global_model_display_name"], "GPT 5");
|
||||
assert_eq!(items[0]["effective_input_price"], 3.0);
|
||||
assert_eq!(items[0]["effective_output_price"], 15.0);
|
||||
assert_eq!(items[0]["effective_supports_streaming"], true);
|
||||
assert_eq!(items[0]["created_at"], "2024-03-21T05:46:40Z");
|
||||
assert_eq!(items[0]["updated_at"], "2024-03-21T05:48:20Z");
|
||||
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
|
||||
|
||||
gateway_handle.abort();
|
||||
upstream_handle.abort();
|
||||
}
|
||||
#[tokio::test]
|
||||
async fn gateway_handles_admin_provider_model_detail_locally_with_trusted_admin_principal() {
|
||||
let upstream_hits = Arc::new(Mutex::new(0usize));
|
||||
let upstream_hits_clone = Arc::clone(&upstream_hits);
|
||||
let upstream = Router::new().route(
|
||||
"/api/admin/providers/provider-openai/models/model-openai-gpt5",
|
||||
any(move |_request: Request| {
|
||||
let upstream_hits_inner = Arc::clone(&upstream_hits_clone);
|
||||
async move {
|
||||
*upstream_hits_inner.lock().expect("mutex should lock") += 1;
|
||||
(StatusCode::OK, Body::from("unexpected upstream hit"))
|
||||
}
|
||||
}),
|
||||
);
|
||||
|
||||
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
|
||||
vec![sample_provider("provider-openai", "openai", 10)],
|
||||
Vec::new(),
|
||||
Vec::new(),
|
||||
));
|
||||
let global_model_repository = Arc::new(
|
||||
InMemoryGlobalModelReadRepository::seed(Vec::new()).with_admin_provider_models(vec![
|
||||
sample_admin_provider_model(
|
||||
"model-openai-gpt5",
|
||||
"provider-openai",
|
||||
"global-gpt-5",
|
||||
"gpt-5-upstream",
|
||||
),
|
||||
]),
|
||||
);
|
||||
let quota_repository = Arc::new(InMemoryProviderQuotaRepository::seed(vec![]));
|
||||
|
||||
let (upstream_url, upstream_handle) = start_server(upstream).await;
|
||||
let gateway = build_router_with_state(
|
||||
AppState::new()
|
||||
.expect("gateway should build")
|
||||
.with_data_state_for_tests(
|
||||
GatewayDataState::with_provider_catalog_global_model_and_quota_readers_for_tests(
|
||||
provider_catalog_repository,
|
||||
global_model_repository,
|
||||
quota_repository,
|
||||
),
|
||||
),
|
||||
);
|
||||
let (gateway_url, gateway_handle) = start_server(gateway).await;
|
||||
|
||||
let response = reqwest::Client::new()
|
||||
.get(format!(
|
||||
"{gateway_url}/api/admin/providers/provider-openai/models/model-openai-gpt5"
|
||||
))
|
||||
.header(crate::constants::GATEWAY_HEADER, "rust-phase3b")
|
||||
.header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
|
||||
.header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
|
||||
.header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
|
||||
.send()
|
||||
.await
|
||||
.expect("request should succeed");
|
||||
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
let payload: serde_json::Value = response.json().await.expect("json body should parse");
|
||||
assert_eq!(payload["id"], "model-openai-gpt5");
|
||||
assert_eq!(
|
||||
payload["provider_model_mappings"],
|
||||
json!([{"name": "gpt-5-upstream-alias", "priority": 1}])
|
||||
);
|
||||
assert_eq!(
|
||||
payload["effective_config"],
|
||||
json!({
|
||||
"billing": {"currency": "USD", "mode": "local"},
|
||||
"provider_hint": "gpt-5-upstream",
|
||||
"streaming": true,
|
||||
"vision": false
|
||||
})
|
||||
);
|
||||
assert_eq!(payload["effective_supports_vision"], true);
|
||||
assert_eq!(payload["effective_supports_image_generation"], false);
|
||||
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
|
||||
|
||||
gateway_handle.abort();
|
||||
upstream_handle.abort();
|
||||
}
|
||||
#[tokio::test]
|
||||
async fn gateway_creates_admin_provider_model_locally_with_trusted_admin_principal() {
|
||||
let upstream_hits = Arc::new(Mutex::new(0usize));
|
||||
let upstream_hits_clone = Arc::clone(&upstream_hits);
|
||||
let upstream = Router::new().route(
|
||||
"/api/admin/providers/provider-openai/models",
|
||||
any(move |_request: Request| {
|
||||
let upstream_hits_inner = Arc::clone(&upstream_hits_clone);
|
||||
async move {
|
||||
*upstream_hits_inner.lock().expect("mutex should lock") += 1;
|
||||
(StatusCode::OK, Body::from("unexpected upstream hit"))
|
||||
}
|
||||
}),
|
||||
);
|
||||
|
||||
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
|
||||
vec![sample_provider("provider-openai", "openai", 10)],
|
||||
Vec::new(),
|
||||
Vec::new(),
|
||||
));
|
||||
let global_model_repository = Arc::new(
|
||||
InMemoryGlobalModelReadRepository::seed(Vec::new()).with_admin_global_models(vec![
|
||||
sample_admin_global_model("global-gpt-5", "gpt-5", "GPT 5"),
|
||||
]),
|
||||
);
|
||||
|
||||
let (upstream_url, upstream_handle) = start_server(upstream).await;
|
||||
let gateway = build_router_with_state(
|
||||
AppState::new()
|
||||
.expect("gateway should build")
|
||||
.with_data_state_for_tests(
|
||||
GatewayDataState::with_provider_catalog_reader_for_tests(
|
||||
provider_catalog_repository,
|
||||
)
|
||||
.with_global_model_repository_for_tests(global_model_repository.clone()),
|
||||
),
|
||||
);
|
||||
let (gateway_url, gateway_handle) = start_server(gateway).await;
|
||||
|
||||
let response = reqwest::Client::new()
|
||||
.post(format!(
|
||||
"{gateway_url}/api/admin/providers/provider-openai/models"
|
||||
))
|
||||
.header(crate::constants::GATEWAY_HEADER, "rust-phase3b")
|
||||
.header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
|
||||
.header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
|
||||
.header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
|
||||
.json(&json!({
|
||||
"provider_model_name": "gpt-5-upstream",
|
||||
"global_model_id": "global-gpt-5",
|
||||
"supports_vision": true,
|
||||
"config": {"provider_hint": "gpt-5-upstream"}
|
||||
}))
|
||||
.send()
|
||||
.await
|
||||
.expect("request should succeed");
|
||||
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
let payload: serde_json::Value = response.json().await.expect("json body should parse");
|
||||
assert_eq!(payload["provider_id"], "provider-openai");
|
||||
assert_eq!(payload["global_model_id"], "global-gpt-5");
|
||||
assert_eq!(payload["provider_model_name"], "gpt-5-upstream");
|
||||
assert_eq!(payload["effective_supports_vision"], true);
|
||||
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
|
||||
|
||||
let created = global_model_repository
|
||||
.list_admin_provider_models(&AdminProviderModelListQuery {
|
||||
provider_id: "provider-openai".to_string(),
|
||||
is_active: None,
|
||||
offset: 0,
|
||||
limit: 20,
|
||||
})
|
||||
.await
|
||||
.expect("models should read");
|
||||
assert_eq!(created.len(), 1);
|
||||
assert_eq!(created[0].provider_model_name, "gpt-5-upstream");
|
||||
|
||||
gateway_handle.abort();
|
||||
upstream_handle.abort();
|
||||
}
|
||||
#[tokio::test]
|
||||
async fn gateway_updates_and_deletes_admin_provider_model_locally_with_trusted_admin_principal() {
|
||||
let upstream_hits = Arc::new(Mutex::new(0usize));
|
||||
let upstream_hits_clone = Arc::clone(&upstream_hits);
|
||||
let upstream = Router::new().route(
|
||||
"/api/admin/providers/provider-openai/models/model-openai-gpt5",
|
||||
any(move |_request: Request| {
|
||||
let upstream_hits_inner = Arc::clone(&upstream_hits_clone);
|
||||
async move {
|
||||
*upstream_hits_inner.lock().expect("mutex should lock") += 1;
|
||||
(StatusCode::OK, Body::from("unexpected upstream hit"))
|
||||
}
|
||||
}),
|
||||
);
|
||||
|
||||
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
|
||||
vec![sample_provider("provider-openai", "openai", 10)],
|
||||
Vec::new(),
|
||||
Vec::new(),
|
||||
));
|
||||
let global_model_repository = Arc::new(
|
||||
InMemoryGlobalModelReadRepository::seed(Vec::new())
|
||||
.with_admin_global_models(vec![
|
||||
sample_admin_global_model("global-gpt-5", "gpt-5", "GPT 5"),
|
||||
sample_admin_global_model("global-gpt-5-mini", "gpt-5-mini", "GPT 5 mini"),
|
||||
])
|
||||
.with_admin_provider_models(vec![sample_admin_provider_model(
|
||||
"model-openai-gpt5",
|
||||
"provider-openai",
|
||||
"global-gpt-5",
|
||||
"gpt-5-upstream",
|
||||
)]),
|
||||
);
|
||||
|
||||
let (upstream_url, upstream_handle) = start_server(upstream).await;
|
||||
let gateway = build_router_with_state(
|
||||
AppState::new()
|
||||
.expect("gateway should build")
|
||||
.with_data_state_for_tests(
|
||||
GatewayDataState::with_provider_catalog_reader_for_tests(
|
||||
provider_catalog_repository,
|
||||
)
|
||||
.with_global_model_repository_for_tests(global_model_repository.clone()),
|
||||
),
|
||||
);
|
||||
let (gateway_url, gateway_handle) = start_server(gateway).await;
|
||||
|
||||
let update_response = reqwest::Client::new()
|
||||
.patch(format!(
|
||||
"{gateway_url}/api/admin/providers/provider-openai/models/model-openai-gpt5"
|
||||
))
|
||||
.header(crate::constants::GATEWAY_HEADER, "rust-phase3b")
|
||||
.header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
|
||||
.header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
|
||||
.header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
|
||||
.json(&json!({
|
||||
"provider_model_name": "gpt-5-mini-upstream",
|
||||
"global_model_id": "global-gpt-5-mini",
|
||||
"supports_streaming": false,
|
||||
"is_available": false
|
||||
}))
|
||||
.send()
|
||||
.await
|
||||
.expect("request should succeed");
|
||||
assert_eq!(update_response.status(), StatusCode::OK);
|
||||
let update_payload: serde_json::Value = update_response
|
||||
.json()
|
||||
.await
|
||||
.expect("json body should parse");
|
||||
assert_eq!(update_payload["provider_model_name"], "gpt-5-mini-upstream");
|
||||
assert_eq!(update_payload["global_model_id"], "global-gpt-5-mini");
|
||||
assert_eq!(update_payload["is_available"], false);
|
||||
|
||||
let delete_response = reqwest::Client::new()
|
||||
.delete(format!(
|
||||
"{gateway_url}/api/admin/providers/provider-openai/models/model-openai-gpt5"
|
||||
))
|
||||
.header(crate::constants::GATEWAY_HEADER, "rust-phase3b")
|
||||
.header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
|
||||
.header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
|
||||
.header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
|
||||
.send()
|
||||
.await
|
||||
.expect("request should succeed");
|
||||
assert_eq!(delete_response.status(), StatusCode::OK);
|
||||
let delete_payload: serde_json::Value = delete_response
|
||||
.json()
|
||||
.await
|
||||
.expect("json body should parse");
|
||||
assert_eq!(
|
||||
delete_payload["message"],
|
||||
"Model 'gpt-5-mini-upstream' deleted successfully"
|
||||
);
|
||||
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
|
||||
|
||||
let models = global_model_repository
|
||||
.list_admin_provider_models(&AdminProviderModelListQuery {
|
||||
provider_id: "provider-openai".to_string(),
|
||||
is_active: None,
|
||||
offset: 0,
|
||||
limit: 20,
|
||||
})
|
||||
.await
|
||||
.expect("models should read");
|
||||
assert!(models.is_empty());
|
||||
|
||||
gateway_handle.abort();
|
||||
upstream_handle.abort();
|
||||
}
|
||||
#[tokio::test]
|
||||
async fn gateway_batch_creates_admin_provider_models_locally_with_trusted_admin_principal() {
|
||||
let upstream_hits = Arc::new(Mutex::new(0usize));
|
||||
let upstream_hits_clone = Arc::clone(&upstream_hits);
|
||||
let upstream = Router::new().route(
|
||||
"/api/admin/providers/provider-openai/models/batch",
|
||||
any(move |_request: Request| {
|
||||
let upstream_hits_inner = Arc::clone(&upstream_hits_clone);
|
||||
async move {
|
||||
*upstream_hits_inner.lock().expect("mutex should lock") += 1;
|
||||
(StatusCode::OK, Body::from("unexpected upstream hit"))
|
||||
}
|
||||
}),
|
||||
);
|
||||
|
||||
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
|
||||
vec![sample_provider("provider-openai", "openai", 10)],
|
||||
Vec::new(),
|
||||
Vec::new(),
|
||||
));
|
||||
let global_model_repository = Arc::new(
|
||||
InMemoryGlobalModelReadRepository::seed(Vec::new()).with_admin_global_models(vec![
|
||||
sample_admin_global_model("global-gpt-5", "gpt-5", "GPT 5"),
|
||||
sample_admin_global_model("global-gpt-4.1", "gpt-4.1", "GPT 4.1"),
|
||||
]),
|
||||
);
|
||||
|
||||
let (upstream_url, upstream_handle) = start_server(upstream).await;
|
||||
let gateway = build_router_with_state(
|
||||
AppState::new()
|
||||
.expect("gateway should build")
|
||||
.with_data_state_for_tests(
|
||||
GatewayDataState::with_provider_catalog_reader_for_tests(
|
||||
provider_catalog_repository,
|
||||
)
|
||||
.with_global_model_repository_for_tests(global_model_repository.clone()),
|
||||
),
|
||||
);
|
||||
let (gateway_url, gateway_handle) = start_server(gateway).await;
|
||||
|
||||
let response = reqwest::Client::new()
|
||||
.post(format!(
|
||||
"{gateway_url}/api/admin/providers/provider-openai/models/batch"
|
||||
))
|
||||
.header(crate::constants::GATEWAY_HEADER, "rust-phase3b")
|
||||
.header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
|
||||
.header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
|
||||
.header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
|
||||
.json(&json!([
|
||||
{"provider_model_name": "gpt-5-upstream", "global_model_id": "global-gpt-5"},
|
||||
{"provider_model_name": "gpt-4.1-upstream", "global_model_id": "global-gpt-4.1"}
|
||||
]))
|
||||
.send()
|
||||
.await
|
||||
.expect("request should succeed");
|
||||
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
let payload: serde_json::Value = response.json().await.expect("json body should parse");
|
||||
assert_eq!(payload.as_array().expect("payload array").len(), 2);
|
||||
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
|
||||
|
||||
gateway_handle.abort();
|
||||
upstream_handle.abort();
|
||||
}
|
||||
#[tokio::test]
|
||||
async fn gateway_handles_admin_provider_available_source_models_locally_with_trusted_admin_principal(
|
||||
) {
|
||||
let upstream_hits = Arc::new(Mutex::new(0usize));
|
||||
let upstream_hits_clone = Arc::clone(&upstream_hits);
|
||||
let upstream = Router::new().route(
|
||||
"/api/admin/providers/provider-openai/available-source-models",
|
||||
any(move |_request: Request| {
|
||||
let upstream_hits_inner = Arc::clone(&upstream_hits_clone);
|
||||
async move {
|
||||
*upstream_hits_inner.lock().expect("mutex should lock") += 1;
|
||||
(StatusCode::OK, Body::from("unexpected upstream hit"))
|
||||
}
|
||||
}),
|
||||
);
|
||||
|
||||
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
|
||||
vec![sample_provider("provider-openai", "openai", 10)],
|
||||
Vec::new(),
|
||||
Vec::new(),
|
||||
));
|
||||
let global_model_repository = Arc::new(
|
||||
InMemoryGlobalModelReadRepository::seed(Vec::new())
|
||||
.with_admin_global_models(vec![sample_admin_global_model(
|
||||
"global-gpt-5",
|
||||
"gpt-5",
|
||||
"GPT 5",
|
||||
)])
|
||||
.with_admin_provider_models(vec![
|
||||
sample_admin_provider_model(
|
||||
"model-openai-gpt5",
|
||||
"provider-openai",
|
||||
"global-gpt-5",
|
||||
"gpt-5-upstream",
|
||||
),
|
||||
sample_admin_provider_model(
|
||||
"model-openai-gpt5-b",
|
||||
"provider-openai",
|
||||
"global-gpt-5",
|
||||
"gpt-5-alt",
|
||||
),
|
||||
]),
|
||||
);
|
||||
|
||||
let (upstream_url, upstream_handle) = start_server(upstream).await;
|
||||
let gateway = build_router_with_state(
|
||||
AppState::new()
|
||||
.expect("gateway should build")
|
||||
.with_data_state_for_tests(
|
||||
GatewayDataState::with_provider_catalog_reader_for_tests(
|
||||
provider_catalog_repository,
|
||||
)
|
||||
.with_global_model_repository_for_tests(global_model_repository),
|
||||
),
|
||||
);
|
||||
let (gateway_url, gateway_handle) = start_server(gateway).await;
|
||||
|
||||
let response = reqwest::Client::new()
|
||||
.get(format!(
|
||||
"{gateway_url}/api/admin/providers/provider-openai/available-source-models"
|
||||
))
|
||||
.header(crate::constants::GATEWAY_HEADER, "rust-phase3b")
|
||||
.header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
|
||||
.header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
|
||||
.header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
|
||||
.send()
|
||||
.await
|
||||
.expect("request should succeed");
|
||||
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
let payload: serde_json::Value = response.json().await.expect("json body should parse");
|
||||
assert_eq!(payload["total"], 1);
|
||||
assert_eq!(payload["models"][0]["global_model_name"], "gpt-5");
|
||||
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
|
||||
|
||||
gateway_handle.abort();
|
||||
upstream_handle.abort();
|
||||
}
|
||||
#[tokio::test]
|
||||
async fn gateway_assigns_and_imports_admin_provider_models_locally_with_trusted_admin_principal() {
|
||||
let upstream_hits = Arc::new(Mutex::new(0usize));
|
||||
let upstream_hits_clone = Arc::clone(&upstream_hits);
|
||||
let upstream = Router::new()
|
||||
.route(
|
||||
"/api/admin/providers/provider-openai/assign-global-models",
|
||||
any({
|
||||
let upstream_hits_inner = Arc::clone(&upstream_hits_clone);
|
||||
move |_request: Request| {
|
||||
let upstream_hits_inner = Arc::clone(&upstream_hits_inner);
|
||||
async move {
|
||||
*upstream_hits_inner.lock().expect("mutex should lock") += 1;
|
||||
(StatusCode::OK, Body::from("unexpected upstream hit"))
|
||||
}
|
||||
}
|
||||
}),
|
||||
)
|
||||
.route(
|
||||
"/api/admin/providers/provider-openai/import-from-upstream",
|
||||
any(move |_request: Request| {
|
||||
let upstream_hits_inner = Arc::clone(&upstream_hits);
|
||||
async move {
|
||||
*upstream_hits_inner.lock().expect("mutex should lock") += 1;
|
||||
(StatusCode::OK, Body::from("unexpected upstream hit"))
|
||||
}
|
||||
}),
|
||||
);
|
||||
|
||||
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
|
||||
vec![sample_provider("provider-openai", "openai", 10)],
|
||||
Vec::new(),
|
||||
Vec::new(),
|
||||
));
|
||||
let global_model_repository = Arc::new(
|
||||
InMemoryGlobalModelReadRepository::seed(Vec::new())
|
||||
.with_admin_global_models(vec![
|
||||
sample_admin_global_model("global-gpt-5", "gpt-5", "GPT 5"),
|
||||
sample_admin_global_model("global-gpt-4.1", "gpt-4.1", "GPT 4.1"),
|
||||
])
|
||||
.with_admin_provider_models(vec![sample_admin_provider_model(
|
||||
"model-openai-gpt5",
|
||||
"provider-openai",
|
||||
"global-gpt-5",
|
||||
"gpt-5-upstream",
|
||||
)]),
|
||||
);
|
||||
|
||||
let (upstream_url, upstream_handle) = start_server(upstream).await;
|
||||
let gateway = build_router_with_state(
|
||||
AppState::new()
|
||||
.expect("gateway should build")
|
||||
.with_data_state_for_tests(
|
||||
GatewayDataState::with_provider_catalog_reader_for_tests(
|
||||
provider_catalog_repository,
|
||||
)
|
||||
.with_global_model_repository_for_tests(global_model_repository.clone()),
|
||||
),
|
||||
);
|
||||
let (gateway_url, gateway_handle) = start_server(gateway).await;
|
||||
|
||||
let assign_response = reqwest::Client::new()
|
||||
.post(format!(
|
||||
"{gateway_url}/api/admin/providers/provider-openai/assign-global-models"
|
||||
))
|
||||
.header(crate::constants::GATEWAY_HEADER, "rust-phase3b")
|
||||
.header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
|
||||
.header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
|
||||
.header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
|
||||
.json(&json!({"global_model_ids": ["global-gpt-5", "global-gpt-4.1"]}))
|
||||
.send()
|
||||
.await
|
||||
.expect("request should succeed");
|
||||
assert_eq!(assign_response.status(), StatusCode::OK);
|
||||
let assign_payload: serde_json::Value = assign_response
|
||||
.json()
|
||||
.await
|
||||
.expect("json body should parse");
|
||||
assert_eq!(
|
||||
assign_payload["success"]
|
||||
.as_array()
|
||||
.expect("success array")
|
||||
.len(),
|
||||
1
|
||||
);
|
||||
assert_eq!(
|
||||
assign_payload["errors"]
|
||||
.as_array()
|
||||
.expect("errors array")
|
||||
.len(),
|
||||
1
|
||||
);
|
||||
|
||||
let import_response = reqwest::Client::new()
|
||||
.post(format!(
|
||||
"{gateway_url}/api/admin/providers/provider-openai/import-from-upstream"
|
||||
))
|
||||
.header(crate::constants::GATEWAY_HEADER, "rust-phase3b")
|
||||
.header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
|
||||
.header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
|
||||
.header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
|
||||
.json(&json!({
|
||||
"model_ids": ["gpt-5-upstream", "brand-new-model"],
|
||||
"price_per_request": 0.1
|
||||
}))
|
||||
.send()
|
||||
.await
|
||||
.expect("request should succeed");
|
||||
assert_eq!(import_response.status(), StatusCode::OK);
|
||||
let import_payload: serde_json::Value = import_response
|
||||
.json()
|
||||
.await
|
||||
.expect("json body should parse");
|
||||
assert_eq!(
|
||||
import_payload["success"]
|
||||
.as_array()
|
||||
.expect("success array")
|
||||
.len(),
|
||||
2
|
||||
);
|
||||
assert_eq!(import_payload["success"][1]["created_global_model"], true);
|
||||
assert_eq!(*upstream_hits_clone.lock().expect("mutex should lock"), 0);
|
||||
|
||||
gateway_handle.abort();
|
||||
upstream_handle.abort();
|
||||
}
|
||||
2142
apps/aether-gateway/src/tests/control/admin/monitoring.rs
Normal file
2142
apps/aether-gateway/src/tests/control/admin/monitoring.rs
Normal file
File diff suppressed because it is too large
Load Diff
2634
apps/aether-gateway/src/tests/control/admin/oauth.rs
Normal file
2634
apps/aether-gateway/src/tests/control/admin/oauth.rs
Normal file
File diff suppressed because it is too large
Load Diff
752
apps/aether-gateway/src/tests/control/admin/payments.rs
Normal file
752
apps/aether-gateway/src/tests/control/admin/payments.rs
Normal file
@@ -0,0 +1,752 @@
|
||||
use std::sync::{Arc, Mutex};
|
||||
|
||||
use aether_data::repository::wallet::StoredWalletSnapshot;
|
||||
use axum::body::Body;
|
||||
use axum::routing::any;
|
||||
use axum::{extract::Request, Router};
|
||||
use http::StatusCode;
|
||||
use serde_json::json;
|
||||
|
||||
use super::super::{build_router_with_state, start_server, AppState};
|
||||
use crate::constants::{
|
||||
GATEWAY_HEADER, TRUSTED_ADMIN_SESSION_ID_HEADER, TRUSTED_ADMIN_USER_ID_HEADER,
|
||||
TRUSTED_ADMIN_USER_ROLE_HEADER,
|
||||
};
|
||||
|
||||
fn admin_request(builder: reqwest::RequestBuilder) -> reqwest::RequestBuilder {
|
||||
builder
|
||||
.header(crate::constants::GATEWAY_HEADER, "rust-phase3b")
|
||||
.header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
|
||||
.header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
|
||||
.header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
|
||||
}
|
||||
|
||||
async fn start_payments_upstream(
|
||||
path: &'static str,
|
||||
) -> (String, Arc<Mutex<usize>>, tokio::task::JoinHandle<()>) {
|
||||
let upstream_hits = Arc::new(Mutex::new(0usize));
|
||||
let upstream_hits_clone = Arc::clone(&upstream_hits);
|
||||
let upstream = Router::new().route(
|
||||
path,
|
||||
any(move |_request: Request| {
|
||||
let upstream_hits_inner = Arc::clone(&upstream_hits_clone);
|
||||
async move {
|
||||
*upstream_hits_inner.lock().expect("mutex should lock") += 1;
|
||||
(StatusCode::OK, Body::from("unexpected upstream hit"))
|
||||
}
|
||||
}),
|
||||
);
|
||||
|
||||
let (upstream_url, upstream_handle) = start_server(upstream).await;
|
||||
(upstream_url, upstream_hits, upstream_handle)
|
||||
}
|
||||
|
||||
fn sample_wallet(wallet_id: &str, user_id: &str) -> StoredWalletSnapshot {
|
||||
StoredWalletSnapshot::new(
|
||||
wallet_id.to_string(),
|
||||
Some(user_id.to_string()),
|
||||
None,
|
||||
12.5,
|
||||
2.5,
|
||||
"finite".to_string(),
|
||||
"USD".to_string(),
|
||||
"active".to_string(),
|
||||
30.0,
|
||||
10.0,
|
||||
3.0,
|
||||
1.5,
|
||||
1_710_000_000,
|
||||
)
|
||||
.expect("wallet should build")
|
||||
}
|
||||
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
fn sample_payment_order(
|
||||
order_id: &str,
|
||||
wallet_id: &str,
|
||||
user_id: &str,
|
||||
amount_usd: f64,
|
||||
payment_method: &str,
|
||||
status: &str,
|
||||
created_at_unix_secs: u64,
|
||||
expires_at_unix_secs: Option<u64>,
|
||||
) -> crate::AdminWalletPaymentOrderRecord {
|
||||
crate::AdminWalletPaymentOrderRecord {
|
||||
id: order_id.to_string(),
|
||||
order_no: format!("po-{order_id}"),
|
||||
wallet_id: wallet_id.to_string(),
|
||||
user_id: Some(user_id.to_string()),
|
||||
amount_usd,
|
||||
pay_amount: None,
|
||||
pay_currency: None,
|
||||
exchange_rate: None,
|
||||
refunded_amount_usd: 0.0,
|
||||
refundable_amount_usd: amount_usd,
|
||||
payment_method: payment_method.to_string(),
|
||||
gateway_order_id: None,
|
||||
status: status.to_string(),
|
||||
gateway_response: None,
|
||||
created_at_unix_secs,
|
||||
paid_at_unix_secs: None,
|
||||
credited_at_unix_secs: None,
|
||||
expires_at_unix_secs,
|
||||
}
|
||||
}
|
||||
|
||||
fn sample_payment_callback(
|
||||
callback_id: &str,
|
||||
payment_order_id: Option<&str>,
|
||||
payment_method: &str,
|
||||
callback_key: &str,
|
||||
status: &str,
|
||||
created_at_unix_secs: u64,
|
||||
) -> crate::state::AdminPaymentCallbackRecord {
|
||||
crate::state::AdminPaymentCallbackRecord {
|
||||
id: callback_id.to_string(),
|
||||
payment_order_id: payment_order_id.map(str::to_string),
|
||||
payment_method: payment_method.to_string(),
|
||||
callback_key: callback_key.to_string(),
|
||||
order_no: payment_order_id.map(|value| format!("po-{value}")),
|
||||
gateway_order_id: Some(format!("gw-{callback_id}")),
|
||||
payload_hash: Some(format!("hash-{callback_id}")),
|
||||
signature_valid: status != "failed",
|
||||
status: status.to_string(),
|
||||
payload: Some(json!({ "source": callback_id })),
|
||||
error_message: (status == "failed").then(|| "signature mismatch".to_string()),
|
||||
created_at_unix_secs,
|
||||
processed_at_unix_secs: Some(created_at_unix_secs.saturating_add(60)),
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_handles_admin_payments_list_orders_locally_with_trusted_admin_principal() {
|
||||
let (upstream_url, upstream_hits, upstream_handle) =
|
||||
start_payments_upstream("/api/admin/payments/orders").await;
|
||||
|
||||
let gateway = build_router_with_state(
|
||||
AppState::new()
|
||||
.expect("gateway should build")
|
||||
.with_admin_wallet_payment_orders_for_tests([
|
||||
sample_payment_order(
|
||||
"order-1",
|
||||
"wallet-1",
|
||||
"user-1",
|
||||
12.5,
|
||||
"alipay",
|
||||
"pending",
|
||||
1_710_000_000,
|
||||
Some(4_102_444_800),
|
||||
),
|
||||
sample_payment_order(
|
||||
"order-2",
|
||||
"wallet-2",
|
||||
"user-2",
|
||||
8.0,
|
||||
"manual",
|
||||
"failed",
|
||||
1_709_000_000,
|
||||
None,
|
||||
),
|
||||
]),
|
||||
);
|
||||
let (gateway_url, gateway_handle) = start_server(gateway).await;
|
||||
|
||||
let response = admin_request(reqwest::Client::new().get(format!(
|
||||
"{gateway_url}/api/admin/payments/orders?status=pending&payment_method=alipay&limit=20&offset=0"
|
||||
)))
|
||||
.send()
|
||||
.await
|
||||
.expect("request should succeed");
|
||||
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
let payload: serde_json::Value = response.json().await.expect("json body should parse");
|
||||
let items = payload["items"].as_array().expect("items should be array");
|
||||
assert_eq!(items.len(), 1);
|
||||
assert_eq!(items[0]["id"], "order-1");
|
||||
assert_eq!(items[0]["wallet_id"], "wallet-1");
|
||||
assert_eq!(items[0]["payment_method"], "alipay");
|
||||
assert_eq!(items[0]["status"], "pending");
|
||||
assert_eq!(payload["total"], 1);
|
||||
assert_eq!(payload["limit"], 20);
|
||||
assert_eq!(payload["offset"], 0);
|
||||
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
|
||||
|
||||
gateway_handle.abort();
|
||||
upstream_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_handles_admin_payments_get_order_locally_with_trusted_admin_principal() {
|
||||
let (upstream_url, upstream_hits, upstream_handle) =
|
||||
start_payments_upstream("/api/admin/payments/orders/order-1").await;
|
||||
|
||||
let gateway = build_router_with_state(
|
||||
AppState::new()
|
||||
.expect("gateway should build")
|
||||
.with_admin_wallet_payment_orders_for_tests([sample_payment_order(
|
||||
"order-1",
|
||||
"wallet-1",
|
||||
"user-1",
|
||||
12.5,
|
||||
"wechat",
|
||||
"pending",
|
||||
1_710_000_000,
|
||||
Some(4_102_444_800),
|
||||
)]),
|
||||
);
|
||||
let (gateway_url, gateway_handle) = start_server(gateway).await;
|
||||
|
||||
let response = admin_request(
|
||||
reqwest::Client::new().get(format!("{gateway_url}/api/admin/payments/orders/order-1")),
|
||||
)
|
||||
.send()
|
||||
.await
|
||||
.expect("request should succeed");
|
||||
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
let payload: serde_json::Value = response.json().await.expect("json body should parse");
|
||||
assert_eq!(payload["order"]["id"], "order-1");
|
||||
assert_eq!(payload["order"]["payment_method"], "wechat");
|
||||
assert_eq!(payload["order"]["amount_usd"], 12.5);
|
||||
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
|
||||
|
||||
gateway_handle.abort();
|
||||
upstream_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_handles_admin_payments_expire_order_locally_with_trusted_admin_principal() {
|
||||
let (upstream_url, upstream_hits, upstream_handle) =
|
||||
start_payments_upstream("/api/admin/payments/orders/order-1/expire").await;
|
||||
|
||||
let gateway = build_router_with_state(
|
||||
AppState::new()
|
||||
.expect("gateway should build")
|
||||
.with_admin_wallet_payment_orders_for_tests([sample_payment_order(
|
||||
"order-1",
|
||||
"wallet-1",
|
||||
"user-1",
|
||||
12.5,
|
||||
"wechat",
|
||||
"pending",
|
||||
1_710_000_000,
|
||||
Some(4_102_444_800),
|
||||
)]),
|
||||
);
|
||||
let (gateway_url, gateway_handle) = start_server(gateway).await;
|
||||
|
||||
let response = admin_request(reqwest::Client::new().post(format!(
|
||||
"{gateway_url}/api/admin/payments/orders/order-1/expire"
|
||||
)))
|
||||
.json(&json!({}))
|
||||
.send()
|
||||
.await
|
||||
.expect("request should succeed");
|
||||
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
let payload: serde_json::Value = response.json().await.expect("json body should parse");
|
||||
assert_eq!(payload["expired"], true);
|
||||
assert_eq!(payload["order"]["status"], "expired");
|
||||
assert_eq!(
|
||||
payload["order"]["gateway_response"]["expire_reason"],
|
||||
"admin_mark_expired"
|
||||
);
|
||||
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
|
||||
|
||||
gateway_handle.abort();
|
||||
upstream_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_handles_admin_payments_credit_order_locally_with_trusted_admin_principal() {
|
||||
let (upstream_url, upstream_hits, upstream_handle) =
|
||||
start_payments_upstream("/api/admin/payments/orders/order-1/credit").await;
|
||||
|
||||
let gateway = build_router_with_state(
|
||||
AppState::new()
|
||||
.expect("gateway should build")
|
||||
.with_auth_wallets_for_tests([sample_wallet("wallet-1", "user-1")])
|
||||
.with_admin_wallet_payment_orders_for_tests([sample_payment_order(
|
||||
"order-1",
|
||||
"wallet-1",
|
||||
"user-1",
|
||||
12.5,
|
||||
"wechat",
|
||||
"pending",
|
||||
1_710_000_000,
|
||||
Some(4_102_444_800),
|
||||
)]),
|
||||
);
|
||||
let (gateway_url, gateway_handle) = start_server(gateway).await;
|
||||
|
||||
let response = admin_request(reqwest::Client::new().post(format!(
|
||||
"{gateway_url}/api/admin/payments/orders/order-1/credit"
|
||||
)))
|
||||
.json(&json!({
|
||||
"gateway_order_id": "gateway-order-1",
|
||||
"pay_amount": 91.25,
|
||||
"pay_currency": "cny",
|
||||
"exchange_rate": 7.3,
|
||||
"gateway_response": {
|
||||
"channel": "manual-review"
|
||||
}
|
||||
}))
|
||||
.send()
|
||||
.await
|
||||
.expect("request should succeed");
|
||||
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
let payload: serde_json::Value = response.json().await.expect("json body should parse");
|
||||
assert_eq!(payload["credited"], true);
|
||||
assert_eq!(payload["order"]["status"], "credited");
|
||||
assert_eq!(payload["order"]["gateway_order_id"], "gateway-order-1");
|
||||
assert_eq!(payload["order"]["pay_amount"], 91.25);
|
||||
assert_eq!(payload["order"]["pay_currency"], "CNY");
|
||||
assert_eq!(payload["order"]["exchange_rate"], 7.3);
|
||||
assert_eq!(
|
||||
payload["order"]["gateway_response"]["channel"],
|
||||
"manual-review"
|
||||
);
|
||||
assert_eq!(payload["order"]["gateway_response"]["manual_credit"], true);
|
||||
assert_eq!(
|
||||
payload["order"]["gateway_response"]["credited_by"],
|
||||
"admin-user-123"
|
||||
);
|
||||
assert!(payload["order"]["paid_at"].as_str().is_some());
|
||||
assert!(payload["order"]["credited_at"].as_str().is_some());
|
||||
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
|
||||
|
||||
gateway_handle.abort();
|
||||
upstream_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_handles_admin_payments_fail_order_locally_with_trusted_admin_principal() {
|
||||
let (upstream_url, upstream_hits, upstream_handle) =
|
||||
start_payments_upstream("/api/admin/payments/orders/order-1/fail").await;
|
||||
|
||||
let gateway = build_router_with_state(
|
||||
AppState::new()
|
||||
.expect("gateway should build")
|
||||
.with_admin_wallet_payment_orders_for_tests([sample_payment_order(
|
||||
"order-1",
|
||||
"wallet-1",
|
||||
"user-1",
|
||||
12.5,
|
||||
"wechat",
|
||||
"pending",
|
||||
1_710_000_000,
|
||||
Some(4_102_444_800),
|
||||
)]),
|
||||
);
|
||||
let (gateway_url, gateway_handle) = start_server(gateway).await;
|
||||
|
||||
let response = admin_request(reqwest::Client::new().post(format!(
|
||||
"{gateway_url}/api/admin/payments/orders/order-1/fail"
|
||||
)))
|
||||
.json(&json!({}))
|
||||
.send()
|
||||
.await
|
||||
.expect("request should succeed");
|
||||
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
let payload: serde_json::Value = response.json().await.expect("json body should parse");
|
||||
assert_eq!(payload["order"]["status"], "failed");
|
||||
assert_eq!(
|
||||
payload["order"]["gateway_response"]["failure_reason"],
|
||||
"admin_mark_failed"
|
||||
);
|
||||
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
|
||||
|
||||
gateway_handle.abort();
|
||||
upstream_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_handles_admin_payments_callbacks_locally_with_trusted_admin_principal() {
|
||||
let (upstream_url, upstream_hits, upstream_handle) =
|
||||
start_payments_upstream("/api/admin/payments/callbacks").await;
|
||||
|
||||
let gateway = build_router_with_state(
|
||||
AppState::new()
|
||||
.expect("gateway should build")
|
||||
.with_admin_payment_callbacks_for_tests([
|
||||
sample_payment_callback(
|
||||
"callback-1",
|
||||
Some("order-1"),
|
||||
"alipay",
|
||||
"callback-key-1",
|
||||
"processed",
|
||||
1_710_100_000,
|
||||
),
|
||||
sample_payment_callback(
|
||||
"callback-2",
|
||||
Some("order-2"),
|
||||
"wechat",
|
||||
"callback-key-2",
|
||||
"failed",
|
||||
1_710_000_000,
|
||||
),
|
||||
]),
|
||||
);
|
||||
let (gateway_url, gateway_handle) = start_server(gateway).await;
|
||||
|
||||
let response = admin_request(reqwest::Client::new().get(format!(
|
||||
"{gateway_url}/api/admin/payments/callbacks?payment_method=alipay&limit=10&offset=0"
|
||||
)))
|
||||
.send()
|
||||
.await
|
||||
.expect("request should succeed");
|
||||
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
let payload: serde_json::Value = response.json().await.expect("json body should parse");
|
||||
let items = payload["items"].as_array().expect("items should be array");
|
||||
assert_eq!(items.len(), 1);
|
||||
assert_eq!(items[0]["id"], "callback-1");
|
||||
assert_eq!(items[0]["payment_order_id"], "order-1");
|
||||
assert_eq!(items[0]["payment_method"], "alipay");
|
||||
assert_eq!(items[0]["callback_key"], "callback-key-1");
|
||||
assert_eq!(items[0]["signature_valid"], true);
|
||||
assert_eq!(items[0]["status"], "processed");
|
||||
assert_eq!(items[0]["payload"]["source"], "callback-1");
|
||||
assert!(items[0]["processed_at"].as_str().is_some());
|
||||
assert_eq!(payload["total"], 1);
|
||||
assert_eq!(payload["limit"], 10);
|
||||
assert_eq!(payload["offset"], 0);
|
||||
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
|
||||
|
||||
gateway_handle.abort();
|
||||
upstream_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_handles_admin_payments_trailing_slash_routes_locally_with_trusted_admin_principal()
|
||||
{
|
||||
let upstream_hits = Arc::new(Mutex::new(0usize));
|
||||
let upstream_hits_clone = Arc::clone(&upstream_hits);
|
||||
let upstream = Router::new()
|
||||
.route(
|
||||
"/api/admin/payments/orders/order-1",
|
||||
any({
|
||||
let upstream_hits_inner = Arc::clone(&upstream_hits_clone);
|
||||
move |_request: Request| {
|
||||
let upstream_hits_inner = Arc::clone(&upstream_hits_inner);
|
||||
async move {
|
||||
*upstream_hits_inner.lock().expect("mutex should lock") += 1;
|
||||
(StatusCode::OK, Body::from("unexpected upstream hit"))
|
||||
}
|
||||
}
|
||||
}),
|
||||
)
|
||||
.route(
|
||||
"/api/admin/payments/orders/order-1/expire",
|
||||
any({
|
||||
let upstream_hits_inner = Arc::clone(&upstream_hits_clone);
|
||||
move |_request: Request| {
|
||||
let upstream_hits_inner = Arc::clone(&upstream_hits_inner);
|
||||
async move {
|
||||
*upstream_hits_inner.lock().expect("mutex should lock") += 1;
|
||||
(StatusCode::OK, Body::from("unexpected upstream hit"))
|
||||
}
|
||||
}
|
||||
}),
|
||||
)
|
||||
.route(
|
||||
"/api/admin/payments/orders/order-1/credit",
|
||||
any({
|
||||
let upstream_hits_inner = Arc::clone(&upstream_hits_clone);
|
||||
move |_request: Request| {
|
||||
let upstream_hits_inner = Arc::clone(&upstream_hits_inner);
|
||||
async move {
|
||||
*upstream_hits_inner.lock().expect("mutex should lock") += 1;
|
||||
(StatusCode::OK, Body::from("unexpected upstream hit"))
|
||||
}
|
||||
}
|
||||
}),
|
||||
)
|
||||
.route(
|
||||
"/api/admin/payments/orders/order-1/fail",
|
||||
any({
|
||||
let upstream_hits_inner = Arc::clone(&upstream_hits_clone);
|
||||
move |_request: Request| {
|
||||
let upstream_hits_inner = Arc::clone(&upstream_hits_inner);
|
||||
async move {
|
||||
*upstream_hits_inner.lock().expect("mutex should lock") += 1;
|
||||
(StatusCode::OK, Body::from("unexpected upstream hit"))
|
||||
}
|
||||
}
|
||||
}),
|
||||
);
|
||||
|
||||
let (upstream_url, upstream_handle) = start_server(upstream).await;
|
||||
let gateway = build_router_with_state(
|
||||
AppState::new()
|
||||
.expect("gateway should build")
|
||||
.with_auth_wallets_for_tests([sample_wallet("wallet-1", "user-1")])
|
||||
.with_admin_wallet_payment_orders_for_tests([sample_payment_order(
|
||||
"order-1",
|
||||
"wallet-1",
|
||||
"user-1",
|
||||
12.5,
|
||||
"wechat",
|
||||
"pending",
|
||||
1_710_000_000,
|
||||
Some(4_102_444_800),
|
||||
)]),
|
||||
);
|
||||
let (gateway_url, gateway_handle) = start_server(gateway).await;
|
||||
|
||||
let detail = admin_request(
|
||||
reqwest::Client::new().get(format!("{gateway_url}/api/admin/payments/orders/order-1/")),
|
||||
)
|
||||
.send()
|
||||
.await
|
||||
.expect("detail request should succeed");
|
||||
assert_eq!(detail.status(), StatusCode::OK);
|
||||
let detail_payload: serde_json::Value = detail.json().await.expect("json body should parse");
|
||||
assert_eq!(detail_payload["order"]["id"], "order-1");
|
||||
|
||||
let expire = admin_request(reqwest::Client::new().post(format!(
|
||||
"{gateway_url}/api/admin/payments/orders/order-1/expire/"
|
||||
)))
|
||||
.json(&json!({}))
|
||||
.send()
|
||||
.await
|
||||
.expect("expire request should succeed");
|
||||
assert_eq!(expire.status(), StatusCode::OK);
|
||||
let expire_payload: serde_json::Value = expire.json().await.expect("json body should parse");
|
||||
assert_eq!(expire_payload["order"]["status"], "expired");
|
||||
|
||||
let mut state = AppState::new().expect("gateway should build");
|
||||
state = state
|
||||
.with_auth_wallets_for_tests([sample_wallet("wallet-1", "user-1")])
|
||||
.with_admin_wallet_payment_orders_for_tests([sample_payment_order(
|
||||
"order-1",
|
||||
"wallet-1",
|
||||
"user-1",
|
||||
12.5,
|
||||
"wechat",
|
||||
"pending",
|
||||
1_710_000_000,
|
||||
Some(4_102_444_800),
|
||||
)]);
|
||||
let gateway = build_router_with_state(state);
|
||||
let (gateway_url, gateway_handle_credit) = start_server(gateway).await;
|
||||
|
||||
let credit = admin_request(reqwest::Client::new().post(format!(
|
||||
"{gateway_url}/api/admin/payments/orders/order-1/credit/"
|
||||
)))
|
||||
.json(&json!({
|
||||
"gateway_order_id": "gateway-order-1",
|
||||
"pay_amount": 91.25,
|
||||
"pay_currency": "cny",
|
||||
"exchange_rate": 7.3
|
||||
}))
|
||||
.send()
|
||||
.await
|
||||
.expect("credit request should succeed");
|
||||
assert_eq!(credit.status(), StatusCode::OK);
|
||||
let credit_payload: serde_json::Value = credit.json().await.expect("json body should parse");
|
||||
assert_eq!(credit_payload["order"]["status"], "credited");
|
||||
|
||||
let gateway = build_router_with_state(
|
||||
AppState::new()
|
||||
.expect("gateway should build")
|
||||
.with_admin_wallet_payment_orders_for_tests([sample_payment_order(
|
||||
"order-1",
|
||||
"wallet-1",
|
||||
"user-1",
|
||||
12.5,
|
||||
"wechat",
|
||||
"pending",
|
||||
1_710_000_000,
|
||||
Some(4_102_444_800),
|
||||
)]),
|
||||
);
|
||||
let (gateway_url, gateway_handle_fail) = start_server(gateway).await;
|
||||
|
||||
let fail = admin_request(reqwest::Client::new().post(format!(
|
||||
"{gateway_url}/api/admin/payments/orders/order-1/fail/"
|
||||
)))
|
||||
.json(&json!({}))
|
||||
.send()
|
||||
.await
|
||||
.expect("fail request should succeed");
|
||||
assert_eq!(fail.status(), StatusCode::OK);
|
||||
let fail_payload: serde_json::Value = fail.json().await.expect("json body should parse");
|
||||
assert_eq!(fail_payload["order"]["status"], "failed");
|
||||
|
||||
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
|
||||
|
||||
gateway_handle.abort();
|
||||
gateway_handle_credit.abort();
|
||||
gateway_handle_fail.abort();
|
||||
upstream_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_handles_admin_payments_list_orders_locally_without_payment_backend() {
|
||||
let (upstream_url, upstream_hits, upstream_handle) =
|
||||
start_payments_upstream("/api/admin/payments/orders").await;
|
||||
|
||||
let mut state = AppState::new().expect("gateway should build");
|
||||
state.admin_wallet_payment_order_store = None;
|
||||
let gateway = build_router_with_state(state);
|
||||
let (gateway_url, gateway_handle) = start_server(gateway).await;
|
||||
|
||||
let response = admin_request(reqwest::Client::new().get(format!(
|
||||
"{gateway_url}/api/admin/payments/orders?limit=10&offset=0"
|
||||
)))
|
||||
.send()
|
||||
.await
|
||||
.expect("request should succeed");
|
||||
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
let payload: serde_json::Value = response.json().await.expect("json body should parse");
|
||||
assert_eq!(payload["items"], json!([]));
|
||||
assert_eq!(payload["total"], 0);
|
||||
assert_eq!(payload["limit"], 10);
|
||||
assert_eq!(payload["offset"], 0);
|
||||
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
|
||||
|
||||
gateway_handle.abort();
|
||||
upstream_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_handles_admin_payments_get_order_locally_without_payment_backend() {
|
||||
let (upstream_url, upstream_hits, upstream_handle) =
|
||||
start_payments_upstream("/api/admin/payments/orders/order-missing").await;
|
||||
|
||||
let mut state = AppState::new().expect("gateway should build");
|
||||
state.admin_wallet_payment_order_store = None;
|
||||
let gateway = build_router_with_state(state);
|
||||
let (gateway_url, gateway_handle) = start_server(gateway).await;
|
||||
|
||||
let response = admin_request(reqwest::Client::new().get(format!(
|
||||
"{gateway_url}/api/admin/payments/orders/order-missing"
|
||||
)))
|
||||
.send()
|
||||
.await
|
||||
.expect("request should succeed");
|
||||
|
||||
assert_eq!(response.status(), StatusCode::SERVICE_UNAVAILABLE);
|
||||
let payload: serde_json::Value = response.json().await.expect("json body should parse");
|
||||
assert_eq!(payload["detail"], "Payment order read backend unavailable");
|
||||
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
|
||||
|
||||
gateway_handle.abort();
|
||||
upstream_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_handles_admin_payments_expire_order_with_backend_unavailable_detail() {
|
||||
let (upstream_url, upstream_hits, upstream_handle) =
|
||||
start_payments_upstream("/api/admin/payments/orders/order-missing/expire").await;
|
||||
|
||||
let mut state = AppState::new().expect("gateway should build");
|
||||
state.admin_wallet_payment_order_store = None;
|
||||
let gateway = build_router_with_state(state);
|
||||
let (gateway_url, gateway_handle) = start_server(gateway).await;
|
||||
|
||||
let response = admin_request(reqwest::Client::new().post(format!(
|
||||
"{gateway_url}/api/admin/payments/orders/order-missing/expire"
|
||||
)))
|
||||
.json(&json!({}))
|
||||
.send()
|
||||
.await
|
||||
.expect("request should succeed");
|
||||
|
||||
assert_eq!(response.status(), StatusCode::SERVICE_UNAVAILABLE);
|
||||
let payload: serde_json::Value = response.json().await.expect("json body should parse");
|
||||
assert_eq!(payload["detail"], "Payment order write backend unavailable");
|
||||
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
|
||||
|
||||
gateway_handle.abort();
|
||||
upstream_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_handles_admin_payments_callbacks_locally_without_payment_backend() {
|
||||
let (upstream_url, upstream_hits, upstream_handle) =
|
||||
start_payments_upstream("/api/admin/payments/callbacks").await;
|
||||
|
||||
let mut state = AppState::new().expect("gateway should build");
|
||||
state.admin_payment_callback_store = None;
|
||||
let gateway = build_router_with_state(state);
|
||||
let (gateway_url, gateway_handle) = start_server(gateway).await;
|
||||
|
||||
let response = admin_request(reqwest::Client::new().get(format!(
|
||||
"{gateway_url}/api/admin/payments/callbacks?payment_method=alipay&limit=10&offset=5"
|
||||
)))
|
||||
.send()
|
||||
.await
|
||||
.expect("request should succeed");
|
||||
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
let payload: serde_json::Value = response.json().await.expect("json body should parse");
|
||||
assert_eq!(payload["items"], json!([]));
|
||||
assert_eq!(payload["total"], 0);
|
||||
assert_eq!(payload["limit"], 10);
|
||||
assert_eq!(payload["offset"], 5);
|
||||
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
|
||||
|
||||
gateway_handle.abort();
|
||||
upstream_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_rejects_admin_payments_empty_order_identifier_locally_with_trusted_admin_principal(
|
||||
) {
|
||||
let upstream_hits = Arc::new(Mutex::new(0usize));
|
||||
let upstream_hits_clone = Arc::clone(&upstream_hits);
|
||||
let upstream = Router::new().fallback(any(move |_request: Request| {
|
||||
let upstream_hits_inner = Arc::clone(&upstream_hits_clone);
|
||||
async move {
|
||||
*upstream_hits_inner.lock().expect("mutex should lock") += 1;
|
||||
(StatusCode::OK, Body::from("unexpected upstream hit"))
|
||||
}
|
||||
}));
|
||||
|
||||
let (upstream_url, upstream_handle) = start_server(upstream).await;
|
||||
let gateway = build_router_with_state(AppState::new().expect("gateway should build"));
|
||||
let (gateway_url, gateway_handle) = start_server(gateway).await;
|
||||
|
||||
let expire = admin_request(
|
||||
reqwest::Client::new().post(format!("{gateway_url}/api/admin/payments/orders//expire")),
|
||||
)
|
||||
.json(&json!({}))
|
||||
.send()
|
||||
.await
|
||||
.expect("expire request should succeed");
|
||||
assert_eq!(expire.status(), StatusCode::NOT_FOUND);
|
||||
let expire_payload: serde_json::Value = expire.json().await.expect("json body should parse");
|
||||
assert_eq!(expire_payload["detail"], "Payment order not found");
|
||||
|
||||
let credit = admin_request(
|
||||
reqwest::Client::new().post(format!("{gateway_url}/api/admin/payments/orders//credit")),
|
||||
)
|
||||
.json(&json!({}))
|
||||
.send()
|
||||
.await
|
||||
.expect("credit request should succeed");
|
||||
assert_eq!(credit.status(), StatusCode::NOT_FOUND);
|
||||
let credit_payload: serde_json::Value = credit.json().await.expect("json body should parse");
|
||||
assert_eq!(credit_payload["detail"], "Payment order not found");
|
||||
|
||||
let fail = admin_request(
|
||||
reqwest::Client::new().post(format!("{gateway_url}/api/admin/payments/orders//fail")),
|
||||
)
|
||||
.json(&json!({}))
|
||||
.send()
|
||||
.await
|
||||
.expect("fail request should succeed");
|
||||
assert_eq!(fail.status(), StatusCode::NOT_FOUND);
|
||||
let fail_payload: serde_json::Value = fail.json().await.expect("json body should parse");
|
||||
assert_eq!(fail_payload["detail"], "Payment order not found");
|
||||
|
||||
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
|
||||
|
||||
gateway_handle.abort();
|
||||
upstream_handle.abort();
|
||||
}
|
||||
1198
apps/aether-gateway/src/tests/control/admin/pool.rs
Normal file
1198
apps/aether-gateway/src/tests/control/admin/pool.rs
Normal file
File diff suppressed because it is too large
Load Diff
1568
apps/aether-gateway/src/tests/control/admin/provider_ops.rs
Normal file
1568
apps/aether-gateway/src/tests/control/admin/provider_ops.rs
Normal file
File diff suppressed because it is too large
Load Diff
344
apps/aether-gateway/src/tests/control/admin/provider_query.rs
Normal file
344
apps/aether-gateway/src/tests/control/admin/provider_query.rs
Normal file
@@ -0,0 +1,344 @@
|
||||
use std::sync::{Arc, Mutex};
|
||||
|
||||
use aether_data::repository::global_models::InMemoryGlobalModelReadRepository;
|
||||
use aether_data::repository::provider_catalog::{
|
||||
InMemoryProviderCatalogReadRepository, StoredProviderCatalogEndpoint,
|
||||
};
|
||||
use axum::body::Body;
|
||||
use axum::routing::any;
|
||||
use axum::{extract::Request, Router};
|
||||
use http::StatusCode;
|
||||
use serde_json::json;
|
||||
|
||||
use super::super::{
|
||||
build_router_with_state, sample_admin_global_model, sample_admin_provider_model, sample_key,
|
||||
sample_provider, start_server, AppState,
|
||||
};
|
||||
use crate::constants::{
|
||||
GATEWAY_HEADER, TRUSTED_ADMIN_SESSION_ID_HEADER, TRUSTED_ADMIN_USER_ID_HEADER,
|
||||
TRUSTED_ADMIN_USER_ROLE_HEADER,
|
||||
};
|
||||
use crate::data::GatewayDataState;
|
||||
|
||||
async fn assert_admin_provider_query_route(
|
||||
path: &str,
|
||||
request_payload: serde_json::Value,
|
||||
expected_status: StatusCode,
|
||||
expected_payload_assertions: impl FnOnce(&serde_json::Value),
|
||||
) {
|
||||
let upstream_hits = Arc::new(Mutex::new(0usize));
|
||||
let upstream_hits_clone = Arc::clone(&upstream_hits);
|
||||
let upstream = Router::new().route(
|
||||
path,
|
||||
any(move |_request: Request| {
|
||||
let upstream_hits_inner = Arc::clone(&upstream_hits_clone);
|
||||
async move {
|
||||
*upstream_hits_inner.lock().expect("mutex should lock") += 1;
|
||||
(StatusCode::OK, Body::from("unexpected upstream hit"))
|
||||
}
|
||||
}),
|
||||
);
|
||||
|
||||
let (upstream_url, upstream_handle) = start_server(upstream).await;
|
||||
let gateway = build_router_with_state(AppState::new().expect("gateway should build"));
|
||||
let (gateway_url, gateway_handle) = start_server(gateway).await;
|
||||
|
||||
let response = reqwest::Client::new()
|
||||
.post(format!("{gateway_url}{path}"))
|
||||
.header(crate::constants::GATEWAY_HEADER, "rust-phase3b")
|
||||
.header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
|
||||
.header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
|
||||
.header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
|
||||
.json(&request_payload)
|
||||
.send()
|
||||
.await
|
||||
.expect("request should succeed");
|
||||
|
||||
assert_eq!(response.status(), expected_status);
|
||||
let payload: serde_json::Value = response.json().await.expect("json body should parse");
|
||||
expected_payload_assertions(&payload);
|
||||
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
|
||||
|
||||
gateway_handle.abort();
|
||||
upstream_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_handles_admin_provider_query_models_locally_with_trusted_admin_principal() {
|
||||
let upstream_hits = Arc::new(Mutex::new(0usize));
|
||||
let upstream_hits_clone = Arc::clone(&upstream_hits);
|
||||
let upstream = Router::new().route(
|
||||
"/api/admin/provider-query/models",
|
||||
any(move |_request: Request| {
|
||||
let upstream_hits_inner = Arc::clone(&upstream_hits_clone);
|
||||
async move {
|
||||
*upstream_hits_inner.lock().expect("mutex should lock") += 1;
|
||||
(StatusCode::OK, Body::from("unexpected upstream hit"))
|
||||
}
|
||||
}),
|
||||
);
|
||||
|
||||
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
|
||||
vec![sample_provider("provider-openai", "OpenAI", 10)],
|
||||
vec![StoredProviderCatalogEndpoint::new(
|
||||
"endpoint-openai-chat".to_string(),
|
||||
"provider-openai".to_string(),
|
||||
"openai:chat".to_string(),
|
||||
Some("chat".to_string()),
|
||||
Some("primary".to_string()),
|
||||
true,
|
||||
)
|
||||
.expect("endpoint should build")
|
||||
.with_transport_fields(
|
||||
"https://api.openai.com/v1".to_string(),
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
)
|
||||
.expect("endpoint transport should build")],
|
||||
vec![
|
||||
{
|
||||
let mut key = sample_key(
|
||||
"key-openai-allowed",
|
||||
"provider-openai",
|
||||
"openai:chat",
|
||||
"sk-test",
|
||||
);
|
||||
key.allowed_models = Some(json!(["gpt-5"]));
|
||||
key
|
||||
},
|
||||
sample_key(
|
||||
"key-openai-all",
|
||||
"provider-openai",
|
||||
"openai:chat",
|
||||
"sk-test-2",
|
||||
),
|
||||
],
|
||||
));
|
||||
let global_model_repository = Arc::new(
|
||||
InMemoryGlobalModelReadRepository::seed(Vec::new())
|
||||
.with_admin_global_models(vec![
|
||||
sample_admin_global_model("global-gpt-5", "gpt-5", "GPT 5"),
|
||||
sample_admin_global_model("global-gpt-4.1", "gpt-4.1", "GPT 4.1"),
|
||||
])
|
||||
.with_admin_provider_models(vec![
|
||||
{
|
||||
let mut model = sample_admin_provider_model(
|
||||
"provider-model-gpt-5",
|
||||
"provider-openai",
|
||||
"global-gpt-5",
|
||||
"gpt-5",
|
||||
);
|
||||
model.global_model_name = Some("gpt-5".to_string());
|
||||
model.global_model_display_name = Some("GPT 5".to_string());
|
||||
model
|
||||
},
|
||||
{
|
||||
let mut model = sample_admin_provider_model(
|
||||
"provider-model-gpt-4.1",
|
||||
"provider-openai",
|
||||
"global-gpt-4.1",
|
||||
"gpt-4.1",
|
||||
);
|
||||
model.global_model_name = Some("gpt-4.1".to_string());
|
||||
model.global_model_display_name = Some("GPT 4.1".to_string());
|
||||
model
|
||||
},
|
||||
]),
|
||||
);
|
||||
|
||||
let (upstream_url, upstream_handle) = start_server(upstream).await;
|
||||
let gateway = build_router_with_state(
|
||||
AppState::new()
|
||||
.expect("gateway should build")
|
||||
.with_data_state_for_tests(
|
||||
GatewayDataState::with_provider_catalog_reader_for_tests(
|
||||
provider_catalog_repository,
|
||||
)
|
||||
.with_global_model_repository_for_tests(global_model_repository),
|
||||
),
|
||||
);
|
||||
let (gateway_url, gateway_handle) = start_server(gateway).await;
|
||||
|
||||
let response = reqwest::Client::new()
|
||||
.post(format!("{gateway_url}/api/admin/provider-query/models"))
|
||||
.header(crate::constants::GATEWAY_HEADER, "rust-phase3b")
|
||||
.header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
|
||||
.header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
|
||||
.header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
|
||||
.json(&json!({
|
||||
"provider_id": "provider-openai",
|
||||
"api_key_id": "key-openai-allowed"
|
||||
}))
|
||||
.send()
|
||||
.await
|
||||
.expect("request should succeed");
|
||||
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
let payload: serde_json::Value = response.json().await.expect("json body should parse");
|
||||
assert_eq!(payload["success"], json!(true));
|
||||
assert_eq!(payload["provider"]["id"], "provider-openai");
|
||||
assert_eq!(payload["provider"]["name"], "OpenAI");
|
||||
assert_eq!(payload["provider"]["display_name"], "OpenAI");
|
||||
assert_eq!(payload["data"]["error"], serde_json::Value::Null);
|
||||
assert_eq!(payload["data"]["from_cache"], json!(true));
|
||||
assert_eq!(payload["data"]["keys_total"], json!(2));
|
||||
let models = payload["data"]["models"]
|
||||
.as_array()
|
||||
.expect("models should be an array");
|
||||
assert_eq!(models.len(), 2);
|
||||
let model_ids: Vec<_> = models
|
||||
.iter()
|
||||
.map(|model| {
|
||||
(
|
||||
model["id"].as_str().expect("id should be present"),
|
||||
model["display_name"]
|
||||
.as_str()
|
||||
.expect("display_name should be present"),
|
||||
)
|
||||
})
|
||||
.collect();
|
||||
assert_eq!(model_ids, vec![("gpt-4.1", "GPT 4.1"), ("gpt-5", "GPT 5")]);
|
||||
for model in models {
|
||||
assert_eq!(model["owned_by"], "OpenAI");
|
||||
assert_eq!(model["api_format"], "openai:chat");
|
||||
assert_eq!(model["api_formats"], json!(["openai:chat"]));
|
||||
}
|
||||
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
|
||||
|
||||
gateway_handle.abort();
|
||||
upstream_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_handles_admin_provider_query_test_model_locally_with_trusted_admin_principal() {
|
||||
assert_admin_provider_query_route(
|
||||
"/api/admin/provider-query/test-model",
|
||||
json!({ "provider_id": "provider-openai", "model": "gpt-4.1" }),
|
||||
StatusCode::OK,
|
||||
|payload| {
|
||||
assert_eq!(payload["success"], json!(false));
|
||||
assert_eq!(payload["tested"], json!(false));
|
||||
assert!(payload["provider_id"].as_str().is_some());
|
||||
},
|
||||
)
|
||||
.await;
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_handles_admin_provider_query_test_model_failover_locally_with_trusted_admin_principal(
|
||||
) {
|
||||
assert_admin_provider_query_route(
|
||||
"/api/admin/provider-query/test-model-failover",
|
||||
json!({
|
||||
"provider_id": "provider-openai",
|
||||
"failover_models": ["gpt-4.1", "gpt-4o-mini"]
|
||||
}),
|
||||
StatusCode::OK,
|
||||
|payload| {
|
||||
assert_eq!(payload["success"], json!(false));
|
||||
assert_eq!(payload["tested"], json!(false));
|
||||
assert!(payload["provider_id"].as_str().is_some());
|
||||
},
|
||||
)
|
||||
.await;
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_rejects_admin_provider_query_invalid_json_body() {
|
||||
let upstream_hits = Arc::new(Mutex::new(0usize));
|
||||
let upstream_hits_clone = Arc::clone(&upstream_hits);
|
||||
let upstream = Router::new().route(
|
||||
"/api/admin/provider-query/models",
|
||||
any(move |_request: Request| {
|
||||
let upstream_hits_inner = Arc::clone(&upstream_hits_clone);
|
||||
async move {
|
||||
*upstream_hits_inner.lock().expect("mutex should lock") += 1;
|
||||
(StatusCode::OK, Body::from("unexpected upstream hit"))
|
||||
}
|
||||
}),
|
||||
);
|
||||
|
||||
let (upstream_url, upstream_handle) = start_server(upstream).await;
|
||||
let gateway = build_router_with_state(AppState::new().expect("gateway should build"));
|
||||
let (gateway_url, gateway_handle) = start_server(gateway).await;
|
||||
|
||||
let response = reqwest::Client::new()
|
||||
.post(format!("{gateway_url}/api/admin/provider-query/models"))
|
||||
.header(crate::constants::GATEWAY_HEADER, "rust-phase3b")
|
||||
.header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
|
||||
.header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
|
||||
.header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
|
||||
.header(reqwest::header::CONTENT_TYPE, "application/json")
|
||||
.body("{")
|
||||
.send()
|
||||
.await
|
||||
.expect("request should succeed");
|
||||
|
||||
assert_eq!(response.status(), StatusCode::BAD_REQUEST);
|
||||
let payload: serde_json::Value = response.json().await.expect("json body should parse");
|
||||
assert_eq!(payload["detail"], json!("Invalid JSON request body"));
|
||||
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
|
||||
|
||||
gateway_handle.abort();
|
||||
upstream_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_rejects_admin_provider_query_test_model_without_provider_id() {
|
||||
assert_admin_provider_query_route(
|
||||
"/api/admin/provider-query/test-model",
|
||||
json!({ "model": "gpt-4.1" }),
|
||||
StatusCode::BAD_REQUEST,
|
||||
|payload| {
|
||||
assert_eq!(payload["detail"], json!("provider_id is required"));
|
||||
},
|
||||
)
|
||||
.await;
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_rejects_admin_provider_query_test_model_without_model() {
|
||||
assert_admin_provider_query_route(
|
||||
"/api/admin/provider-query/test-model",
|
||||
json!({ "provider_id": "provider-openai" }),
|
||||
StatusCode::BAD_REQUEST,
|
||||
|payload| {
|
||||
assert_eq!(payload["detail"], json!("model is required"));
|
||||
},
|
||||
)
|
||||
.await;
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_rejects_admin_provider_query_test_model_failover_without_provider_id() {
|
||||
assert_admin_provider_query_route(
|
||||
"/api/admin/provider-query/test-model-failover",
|
||||
json!({ "failover_models": ["gpt-4.1"] }),
|
||||
StatusCode::BAD_REQUEST,
|
||||
|payload| {
|
||||
assert_eq!(payload["detail"], json!("provider_id is required"));
|
||||
},
|
||||
)
|
||||
.await;
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_rejects_admin_provider_query_test_model_failover_without_models() {
|
||||
assert_admin_provider_query_route(
|
||||
"/api/admin/provider-query/test-model-failover",
|
||||
json!({ "provider_id": "provider-openai", "failover_models": [] }),
|
||||
StatusCode::BAD_REQUEST,
|
||||
|payload| {
|
||||
assert_eq!(
|
||||
payload["detail"],
|
||||
json!("failover_models should not be empty")
|
||||
);
|
||||
},
|
||||
)
|
||||
.await;
|
||||
}
|
||||
617
apps/aether-gateway/src/tests/control/admin/provider_strategy.rs
Normal file
617
apps/aether-gateway/src/tests/control/admin/provider_strategy.rs
Normal file
@@ -0,0 +1,617 @@
|
||||
use std::sync::{Arc, Mutex};
|
||||
|
||||
use aether_data::repository::provider_catalog::{
|
||||
InMemoryProviderCatalogReadRepository, ProviderCatalogReadRepository,
|
||||
};
|
||||
use aether_data::repository::usage::{InMemoryUsageReadRepository, StoredProviderUsageWindow};
|
||||
use axum::body::Body;
|
||||
use axum::routing::any;
|
||||
use axum::{extract::Request, Router};
|
||||
use http::{Method, StatusCode};
|
||||
use serde_json::json;
|
||||
|
||||
use super::super::{build_router_with_state, sample_provider, start_server, AppState};
|
||||
use crate::constants::{
|
||||
GATEWAY_HEADER, TRUSTED_ADMIN_SESSION_ID_HEADER, TRUSTED_ADMIN_USER_ID_HEADER,
|
||||
TRUSTED_ADMIN_USER_ROLE_HEADER,
|
||||
};
|
||||
use crate::data::GatewayDataState;
|
||||
|
||||
async fn assert_provider_strategy_route_returns_local_503(
|
||||
data_state: GatewayDataState,
|
||||
method: Method,
|
||||
path: &str,
|
||||
body: Option<serde_json::Value>,
|
||||
expected_message: &str,
|
||||
) {
|
||||
let upstream_hits = Arc::new(Mutex::new(0usize));
|
||||
let upstream_hits_clone = Arc::clone(&upstream_hits);
|
||||
let upstream = Router::new().route(
|
||||
"/{*path}",
|
||||
any(move |_request: Request| {
|
||||
let upstream_hits_inner = Arc::clone(&upstream_hits_clone);
|
||||
async move {
|
||||
*upstream_hits_inner.lock().expect("mutex should lock") += 1;
|
||||
(StatusCode::OK, Body::from("unexpected upstream hit"))
|
||||
}
|
||||
}),
|
||||
);
|
||||
|
||||
let (upstream_url, upstream_handle) = start_server(upstream).await;
|
||||
let gateway = build_router_with_state(
|
||||
AppState::new()
|
||||
.expect("gateway should build")
|
||||
.with_data_state_for_tests(data_state),
|
||||
);
|
||||
let (gateway_url, gateway_handle) = start_server(gateway).await;
|
||||
|
||||
let mut request = reqwest::Client::new()
|
||||
.request(method, format!("{gateway_url}{path}"))
|
||||
.header(crate::constants::GATEWAY_HEADER, "rust-phase3b")
|
||||
.header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
|
||||
.header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
|
||||
.header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123");
|
||||
if let Some(body) = body {
|
||||
request = request.json(&body);
|
||||
}
|
||||
let response = request.send().await.expect("request should succeed");
|
||||
|
||||
assert_eq!(response.status(), StatusCode::SERVICE_UNAVAILABLE);
|
||||
let payload: serde_json::Value = response.json().await.expect("json body should parse");
|
||||
assert_eq!(payload["error"]["message"], expected_message);
|
||||
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
|
||||
|
||||
gateway_handle.abort();
|
||||
upstream_handle.abort();
|
||||
}
|
||||
|
||||
const PROVIDER_STRATEGY_DATA_UNAVAILABLE_MESSAGE: &str = "Admin provider strategy data unavailable";
|
||||
const PROVIDER_STRATEGY_STATS_DATA_UNAVAILABLE_MESSAGE: &str =
|
||||
"Admin provider strategy stats data unavailable";
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_handles_admin_provider_strategy_list_locally_with_trusted_admin_principal() {
|
||||
let upstream_hits = Arc::new(Mutex::new(0usize));
|
||||
let upstream_hits_clone = Arc::clone(&upstream_hits);
|
||||
let upstream = Router::new().route(
|
||||
"/api/admin/provider-strategy/strategies",
|
||||
any(move |_request: Request| {
|
||||
let upstream_hits_inner = Arc::clone(&upstream_hits_clone);
|
||||
async move {
|
||||
*upstream_hits_inner.lock().expect("mutex should lock") += 1;
|
||||
(StatusCode::OK, Body::from("unexpected upstream hit"))
|
||||
}
|
||||
}),
|
||||
);
|
||||
|
||||
let (upstream_url, upstream_handle) = start_server(upstream).await;
|
||||
let gateway = build_router_with_state(AppState::new().expect("gateway should build"));
|
||||
let (gateway_url, gateway_handle) = start_server(gateway).await;
|
||||
|
||||
let response = reqwest::Client::new()
|
||||
.get(format!(
|
||||
"{gateway_url}/api/admin/provider-strategy/strategies"
|
||||
))
|
||||
.header(crate::constants::GATEWAY_HEADER, "rust-phase3b")
|
||||
.header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
|
||||
.header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
|
||||
.header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
|
||||
.send()
|
||||
.await
|
||||
.expect("request should succeed");
|
||||
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
let payload: serde_json::Value = response.json().await.expect("json body should parse");
|
||||
assert_eq!(payload["total"], 1);
|
||||
assert_eq!(payload["strategies"][0]["name"], "sticky_priority");
|
||||
assert_eq!(payload["strategies"][0]["priority"], 110);
|
||||
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
|
||||
|
||||
gateway_handle.abort();
|
||||
upstream_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_handles_admin_provider_strategy_billing_locally_returns_service_unavailable_when_disabled(
|
||||
) {
|
||||
assert_provider_strategy_route_returns_local_503(
|
||||
crate::data::GatewayDataState::disabled(),
|
||||
Method::PUT,
|
||||
"/api/admin/provider-strategy/providers/provider-openai/billing",
|
||||
Some(json!({
|
||||
"billing_type": "monthly_quota",
|
||||
"monthly_quota_usd": 100.0,
|
||||
"quota_reset_day": 30,
|
||||
"quota_last_reset_at": "2024-03-21T00:00:00Z",
|
||||
"quota_expires_at": "2024-04-21T00:00:00Z",
|
||||
"rpm_limit": 20,
|
||||
"provider_priority": 5
|
||||
})),
|
||||
PROVIDER_STRATEGY_DATA_UNAVAILABLE_MESSAGE,
|
||||
)
|
||||
.await;
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_handles_admin_provider_strategy_stats_locally_returns_service_unavailable_without_usage_reader(
|
||||
) {
|
||||
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
|
||||
vec![sample_provider("provider-openai", "openai", 10)],
|
||||
vec![],
|
||||
vec![],
|
||||
));
|
||||
let catalog_reader: Arc<dyn ProviderCatalogReadRepository> =
|
||||
provider_catalog_repository.clone();
|
||||
assert_provider_strategy_route_returns_local_503(
|
||||
crate::data::GatewayDataState::with_provider_catalog_reader_for_tests(
|
||||
catalog_reader,
|
||||
),
|
||||
Method::GET,
|
||||
"/api/admin/provider-strategy/providers/provider-openai/stats?hours=1",
|
||||
None,
|
||||
PROVIDER_STRATEGY_STATS_DATA_UNAVAILABLE_MESSAGE,
|
||||
)
|
||||
.await;
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_handles_admin_provider_strategy_quota_locally_returns_service_unavailable_without_writer(
|
||||
) {
|
||||
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
|
||||
vec![sample_provider("provider-openai", "openai", 10)],
|
||||
vec![],
|
||||
vec![],
|
||||
));
|
||||
let catalog_reader: Arc<dyn ProviderCatalogReadRepository> =
|
||||
provider_catalog_repository.clone();
|
||||
assert_provider_strategy_route_returns_local_503(
|
||||
crate::data::GatewayDataState::with_provider_catalog_reader_for_tests(
|
||||
catalog_reader,
|
||||
),
|
||||
Method::DELETE,
|
||||
"/api/admin/provider-strategy/providers/provider-openai/quota",
|
||||
None,
|
||||
PROVIDER_STRATEGY_DATA_UNAVAILABLE_MESSAGE,
|
||||
)
|
||||
.await;
|
||||
}
|
||||
async fn gateway_resets_admin_provider_strategy_quota_locally_with_trusted_admin_principal() {
|
||||
let upstream_hits = Arc::new(Mutex::new(0usize));
|
||||
let upstream_hits_clone = Arc::clone(&upstream_hits);
|
||||
let upstream = Router::new().route(
|
||||
"/api/admin/provider-strategy/providers/provider-openai/quota",
|
||||
any(move |_request: Request| {
|
||||
let upstream_hits_inner = Arc::clone(&upstream_hits_clone);
|
||||
async move {
|
||||
*upstream_hits_inner.lock().expect("mutex should lock") += 1;
|
||||
(StatusCode::OK, Body::from("unexpected upstream hit"))
|
||||
}
|
||||
}),
|
||||
);
|
||||
|
||||
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
|
||||
vec![
|
||||
sample_provider("provider-openai", "openai", 10).with_billing_fields(
|
||||
Some("monthly_quota".to_string()),
|
||||
Some(100.0),
|
||||
Some(12.5),
|
||||
Some(30),
|
||||
Some(1_711_000_000),
|
||||
Some(1_711_000_000 + 30 * 24 * 60 * 60),
|
||||
),
|
||||
],
|
||||
vec![],
|
||||
vec![],
|
||||
));
|
||||
|
||||
let (upstream_url, upstream_handle) = start_server(upstream).await;
|
||||
let gateway = build_router_with_state(
|
||||
AppState::new()
|
||||
.expect("gateway should build")
|
||||
.with_data_state_for_tests(
|
||||
GatewayDataState::with_provider_catalog_repository_for_tests(Arc::clone(
|
||||
&provider_catalog_repository,
|
||||
)),
|
||||
),
|
||||
);
|
||||
let (gateway_url, gateway_handle) = start_server(gateway).await;
|
||||
|
||||
let response = reqwest::Client::new()
|
||||
.delete(format!(
|
||||
"{gateway_url}/api/admin/provider-strategy/providers/provider-openai/quota"
|
||||
))
|
||||
.header(crate::constants::GATEWAY_HEADER, "rust-phase3b")
|
||||
.header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
|
||||
.header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
|
||||
.header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
|
||||
.send()
|
||||
.await
|
||||
.expect("request should succeed");
|
||||
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
let payload: serde_json::Value = response.json().await.expect("json body should parse");
|
||||
assert_eq!(payload["provider_name"], "openai");
|
||||
assert_eq!(payload["previous_used"], 12.5);
|
||||
assert_eq!(payload["current_used"], 0.0);
|
||||
|
||||
let updated = provider_catalog_repository
|
||||
.list_providers_by_ids(&["provider-openai".to_string()])
|
||||
.await
|
||||
.expect("provider query should succeed");
|
||||
assert_eq!(updated[0].monthly_used_usd, Some(0.0));
|
||||
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
|
||||
|
||||
gateway_handle.abort();
|
||||
upstream_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_updates_admin_provider_strategy_billing_locally_with_trusted_admin_principal() {
|
||||
let upstream_hits = Arc::new(Mutex::new(0usize));
|
||||
let upstream_hits_clone = Arc::clone(&upstream_hits);
|
||||
let upstream = Router::new().route(
|
||||
"/api/admin/provider-strategy/providers/provider-openai/billing",
|
||||
any(move |_request: Request| {
|
||||
let upstream_hits_inner = Arc::clone(&upstream_hits_clone);
|
||||
async move {
|
||||
*upstream_hits_inner.lock().expect("mutex should lock") += 1;
|
||||
(StatusCode::OK, Body::from("unexpected upstream hit"))
|
||||
}
|
||||
}),
|
||||
);
|
||||
|
||||
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
|
||||
vec![
|
||||
sample_provider("provider-openai", "openai", 7).with_billing_fields(
|
||||
Some("pay_as_you_go".to_string()),
|
||||
None,
|
||||
Some(12.5),
|
||||
Some(9),
|
||||
Some(1_711_000_000),
|
||||
None,
|
||||
),
|
||||
],
|
||||
vec![],
|
||||
vec![],
|
||||
));
|
||||
|
||||
let (upstream_url, upstream_handle) = start_server(upstream).await;
|
||||
let gateway = build_router_with_state(
|
||||
AppState::new()
|
||||
.expect("gateway should build")
|
||||
.with_data_state_for_tests(
|
||||
GatewayDataState::with_provider_catalog_repository_for_tests(Arc::clone(
|
||||
&provider_catalog_repository,
|
||||
)),
|
||||
),
|
||||
);
|
||||
let (gateway_url, gateway_handle) = start_server(gateway).await;
|
||||
|
||||
let response = reqwest::Client::new()
|
||||
.put(format!(
|
||||
"{gateway_url}/api/admin/provider-strategy/providers/provider-openai/billing"
|
||||
))
|
||||
.header(crate::constants::GATEWAY_HEADER, "rust-phase3b")
|
||||
.header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
|
||||
.header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
|
||||
.header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
|
||||
.json(&json!({
|
||||
"billing_type": "monthly_quota",
|
||||
"monthly_quota_usd": 50.0,
|
||||
"quota_last_reset_at": "2026-03-01T00:00:00Z",
|
||||
"quota_expires_at": "2026-04-01T00:00:00Z"
|
||||
}))
|
||||
.send()
|
||||
.await
|
||||
.expect("request should succeed");
|
||||
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
let payload: serde_json::Value = response.json().await.expect("json body should parse");
|
||||
assert_eq!(payload["provider"]["id"], "provider-openai");
|
||||
assert_eq!(payload["provider"]["billing_type"], "monthly_quota");
|
||||
assert_eq!(payload["provider"]["provider_priority"], 100);
|
||||
|
||||
let updated = provider_catalog_repository
|
||||
.list_providers_by_ids(&["provider-openai".to_string()])
|
||||
.await
|
||||
.expect("provider query should succeed");
|
||||
assert_eq!(updated[0].billing_type.as_deref(), Some("monthly_quota"));
|
||||
assert_eq!(updated[0].monthly_quota_usd, Some(50.0));
|
||||
assert_eq!(updated[0].monthly_used_usd, Some(12.5));
|
||||
assert_eq!(updated[0].quota_reset_day, Some(30));
|
||||
assert_eq!(updated[0].provider_priority, 100);
|
||||
assert_eq!(
|
||||
updated[0].quota_last_reset_at_unix_secs,
|
||||
Some(1_772_323_200)
|
||||
);
|
||||
assert_eq!(updated[0].quota_expires_at_unix_secs, Some(1_775_001_600));
|
||||
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
|
||||
|
||||
gateway_handle.abort();
|
||||
upstream_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_syncs_admin_provider_strategy_monthly_usage_from_reset_window_locally() {
|
||||
let upstream_hits = Arc::new(Mutex::new(0usize));
|
||||
let upstream_hits_clone = Arc::clone(&upstream_hits);
|
||||
let upstream = Router::new().route(
|
||||
"/api/admin/provider-strategy/providers/provider-openai/billing",
|
||||
any(move |_request: Request| {
|
||||
let upstream_hits_inner = Arc::clone(&upstream_hits_clone);
|
||||
async move {
|
||||
*upstream_hits_inner.lock().expect("mutex should lock") += 1;
|
||||
(StatusCode::OK, Body::from("unexpected upstream hit"))
|
||||
}
|
||||
}),
|
||||
);
|
||||
|
||||
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
|
||||
vec![
|
||||
sample_provider("provider-openai", "openai", 7).with_billing_fields(
|
||||
Some("pay_as_you_go".to_string()),
|
||||
None,
|
||||
Some(12.5),
|
||||
Some(9),
|
||||
Some(1_711_000_000),
|
||||
None,
|
||||
),
|
||||
],
|
||||
vec![],
|
||||
vec![],
|
||||
));
|
||||
let usage_repository = Arc::new(
|
||||
InMemoryUsageReadRepository::default().with_provider_usage_windows(vec![
|
||||
StoredProviderUsageWindow::new(
|
||||
"provider-openai".to_string(),
|
||||
1_772_236_799,
|
||||
8,
|
||||
7,
|
||||
1,
|
||||
100.0,
|
||||
9.99,
|
||||
)
|
||||
.expect("window should build"),
|
||||
StoredProviderUsageWindow::new(
|
||||
"provider-openai".to_string(),
|
||||
1_772_323_200,
|
||||
5,
|
||||
5,
|
||||
0,
|
||||
110.0,
|
||||
1.25,
|
||||
)
|
||||
.expect("window should build"),
|
||||
StoredProviderUsageWindow::new(
|
||||
"provider-openai".to_string(),
|
||||
1_772_409_600,
|
||||
6,
|
||||
5,
|
||||
1,
|
||||
150.0,
|
||||
0.75,
|
||||
)
|
||||
.expect("window should build"),
|
||||
]),
|
||||
);
|
||||
|
||||
let (upstream_url, upstream_handle) = start_server(upstream).await;
|
||||
let gateway = build_router_with_state(
|
||||
AppState::new()
|
||||
.expect("gateway should build")
|
||||
.with_data_state_for_tests(
|
||||
GatewayDataState::with_provider_catalog_and_usage_reader_for_tests(
|
||||
Arc::clone(&provider_catalog_repository),
|
||||
Arc::clone(&usage_repository),
|
||||
),
|
||||
),
|
||||
);
|
||||
let (gateway_url, gateway_handle) = start_server(gateway).await;
|
||||
|
||||
let response = reqwest::Client::new()
|
||||
.put(format!(
|
||||
"{gateway_url}/api/admin/provider-strategy/providers/provider-openai/billing"
|
||||
))
|
||||
.header(crate::constants::GATEWAY_HEADER, "rust-phase3b")
|
||||
.header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
|
||||
.header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
|
||||
.header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
|
||||
.json(&json!({
|
||||
"billing_type": "monthly_quota",
|
||||
"monthly_quota_usd": 50.0,
|
||||
"quota_last_reset_at": "2026-03-01T00:00:00Z",
|
||||
"quota_expires_at": "2026-04-01T00:00:00Z"
|
||||
}))
|
||||
.send()
|
||||
.await
|
||||
.expect("request should succeed");
|
||||
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
|
||||
let updated = provider_catalog_repository
|
||||
.list_providers_by_ids(&["provider-openai".to_string()])
|
||||
.await
|
||||
.expect("provider query should succeed");
|
||||
assert_eq!(updated[0].monthly_used_usd, Some(2.0));
|
||||
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
|
||||
|
||||
gateway_handle.abort();
|
||||
upstream_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_reads_admin_provider_strategy_stats_locally_with_trusted_admin_principal() {
|
||||
let upstream_hits = Arc::new(Mutex::new(0usize));
|
||||
let upstream_hits_clone = Arc::clone(&upstream_hits);
|
||||
let upstream = Router::new().route(
|
||||
"/api/admin/provider-strategy/providers/provider-openai/stats",
|
||||
any(move |_request: Request| {
|
||||
let upstream_hits_inner = Arc::clone(&upstream_hits_clone);
|
||||
async move {
|
||||
*upstream_hits_inner.lock().expect("mutex should lock") += 1;
|
||||
(StatusCode::OK, Body::from("unexpected upstream hit"))
|
||||
}
|
||||
}),
|
||||
);
|
||||
|
||||
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
|
||||
vec![
|
||||
sample_provider("provider-openai", "openai", 7).with_billing_fields(
|
||||
Some("monthly_quota".to_string()),
|
||||
Some(80.0),
|
||||
Some(12.5),
|
||||
Some(30),
|
||||
Some(1_711_000_000),
|
||||
Some(1_774_966_400),
|
||||
),
|
||||
],
|
||||
vec![],
|
||||
vec![],
|
||||
));
|
||||
let usage_repository = Arc::new(
|
||||
InMemoryUsageReadRepository::default().with_provider_usage_windows(vec![
|
||||
StoredProviderUsageWindow::new(
|
||||
"provider-openai".to_string(),
|
||||
1_700_000_000,
|
||||
10,
|
||||
9,
|
||||
1,
|
||||
120.0,
|
||||
1.25,
|
||||
)
|
||||
.expect("window should build"),
|
||||
StoredProviderUsageWindow::new(
|
||||
"provider-openai".to_string(),
|
||||
1_700_003_600,
|
||||
6,
|
||||
5,
|
||||
1,
|
||||
180.0,
|
||||
0.75,
|
||||
)
|
||||
.expect("window should build"),
|
||||
]),
|
||||
);
|
||||
|
||||
let (upstream_url, upstream_handle) = start_server(upstream).await;
|
||||
let gateway = build_router_with_state(
|
||||
AppState::new()
|
||||
.expect("gateway should build")
|
||||
.with_data_state_for_tests(
|
||||
GatewayDataState::with_provider_catalog_and_usage_reader_for_tests(
|
||||
Arc::clone(&provider_catalog_repository),
|
||||
Arc::clone(&usage_repository),
|
||||
),
|
||||
),
|
||||
);
|
||||
let (gateway_url, gateway_handle) = start_server(gateway).await;
|
||||
|
||||
let response = reqwest::Client::new()
|
||||
.get(format!(
|
||||
"{gateway_url}/api/admin/provider-strategy/providers/provider-openai/stats?hours=999999"
|
||||
))
|
||||
.header(crate::constants::GATEWAY_HEADER, "rust-phase3b")
|
||||
.header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
|
||||
.header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
|
||||
.header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
|
||||
.send()
|
||||
.await
|
||||
.expect("request should succeed");
|
||||
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
let payload: serde_json::Value = response.json().await.expect("json body should parse");
|
||||
assert_eq!(payload["provider_id"], "provider-openai");
|
||||
assert_eq!(payload["provider_name"], "openai");
|
||||
assert_eq!(payload["billing_info"]["billing_type"], "monthly_quota");
|
||||
assert_eq!(payload["billing_info"]["monthly_quota_usd"], 80.0);
|
||||
assert_eq!(payload["billing_info"]["monthly_used_usd"], 12.5);
|
||||
assert_eq!(payload["billing_info"]["quota_remaining_usd"], 67.5);
|
||||
assert_eq!(payload["usage_stats"]["total_requests"], 16);
|
||||
assert_eq!(payload["usage_stats"]["successful_requests"], 14);
|
||||
assert_eq!(payload["usage_stats"]["failed_requests"], 2);
|
||||
assert_eq!(payload["usage_stats"]["success_rate"], 0.875);
|
||||
assert_eq!(payload["usage_stats"]["avg_response_time_ms"], 150.0);
|
||||
assert_eq!(payload["usage_stats"]["total_cost_usd"], 2.0);
|
||||
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
|
||||
|
||||
gateway_handle.abort();
|
||||
upstream_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_rounds_admin_provider_strategy_stats_total_cost_locally() {
|
||||
let upstream_hits = Arc::new(Mutex::new(0usize));
|
||||
let upstream_hits_clone = Arc::clone(&upstream_hits);
|
||||
let upstream = Router::new().route(
|
||||
"/api/admin/provider-strategy/providers/provider-openai/stats",
|
||||
any(move |_request: Request| {
|
||||
let upstream_hits_inner = Arc::clone(&upstream_hits_clone);
|
||||
async move {
|
||||
*upstream_hits_inner.lock().expect("mutex should lock") += 1;
|
||||
(StatusCode::OK, Body::from("unexpected upstream hit"))
|
||||
}
|
||||
}),
|
||||
);
|
||||
|
||||
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
|
||||
vec![
|
||||
sample_provider("provider-openai", "openai", 7).with_billing_fields(
|
||||
Some("monthly_quota".to_string()),
|
||||
Some(80.0),
|
||||
Some(12.5),
|
||||
Some(30),
|
||||
Some(1_711_000_000),
|
||||
Some(1_774_966_400),
|
||||
),
|
||||
],
|
||||
vec![],
|
||||
vec![],
|
||||
));
|
||||
let usage_repository = Arc::new(
|
||||
InMemoryUsageReadRepository::default().with_provider_usage_windows(vec![
|
||||
StoredProviderUsageWindow::new(
|
||||
"provider-openai".to_string(),
|
||||
1_700_000_000,
|
||||
2,
|
||||
2,
|
||||
0,
|
||||
120.0,
|
||||
1.23456,
|
||||
)
|
||||
.expect("window should build"),
|
||||
]),
|
||||
);
|
||||
|
||||
let (upstream_url, upstream_handle) = start_server(upstream).await;
|
||||
let gateway = build_router_with_state(
|
||||
AppState::new()
|
||||
.expect("gateway should build")
|
||||
.with_data_state_for_tests(
|
||||
GatewayDataState::with_provider_catalog_and_usage_reader_for_tests(
|
||||
Arc::clone(&provider_catalog_repository),
|
||||
Arc::clone(&usage_repository),
|
||||
),
|
||||
),
|
||||
);
|
||||
let (gateway_url, gateway_handle) = start_server(gateway).await;
|
||||
|
||||
let response = reqwest::Client::new()
|
||||
.get(format!(
|
||||
"{gateway_url}/api/admin/provider-strategy/providers/provider-openai/stats?hours=999999"
|
||||
))
|
||||
.header(crate::constants::GATEWAY_HEADER, "rust-phase3b")
|
||||
.header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
|
||||
.header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
|
||||
.header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
|
||||
.send()
|
||||
.await
|
||||
.expect("request should succeed");
|
||||
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
let payload: serde_json::Value = response.json().await.expect("json body should parse");
|
||||
assert_eq!(payload["usage_stats"]["total_cost_usd"], 1.2346);
|
||||
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
|
||||
|
||||
gateway_handle.abort();
|
||||
upstream_handle.abort();
|
||||
}
|
||||
1604
apps/aether-gateway/src/tests/control/admin/providers.rs
Normal file
1604
apps/aether-gateway/src/tests/control/admin/providers.rs
Normal file
File diff suppressed because it is too large
Load Diff
216
apps/aether-gateway/src/tests/control/admin/proxy_nodes.rs
Normal file
216
apps/aether-gateway/src/tests/control/admin/proxy_nodes.rs
Normal file
@@ -0,0 +1,216 @@
|
||||
use std::sync::{Arc, Mutex};
|
||||
|
||||
use aether_data::repository::proxy_nodes::{InMemoryProxyNodeRepository, StoredProxyNodeEvent};
|
||||
use axum::body::Body;
|
||||
use axum::routing::any;
|
||||
use axum::{extract::Request, Router};
|
||||
use http::StatusCode;
|
||||
use serde_json::json;
|
||||
|
||||
use super::super::{build_router_with_state, sample_proxy_node, start_server, AppState};
|
||||
use crate::constants::{
|
||||
GATEWAY_HEADER, TRUSTED_ADMIN_SESSION_ID_HEADER, TRUSTED_ADMIN_USER_ID_HEADER,
|
||||
TRUSTED_ADMIN_USER_ROLE_HEADER,
|
||||
};
|
||||
use crate::data::GatewayDataState;
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_handles_admin_proxy_nodes_locally_with_trusted_admin_principal() {
|
||||
let upstream_hits = Arc::new(Mutex::new(0usize));
|
||||
let upstream_hits_clone = Arc::clone(&upstream_hits);
|
||||
let upstream = Router::new().route(
|
||||
"/api/admin/proxy-nodes",
|
||||
any(move |_request: Request| {
|
||||
let upstream_hits_inner = Arc::clone(&upstream_hits_clone);
|
||||
async move {
|
||||
*upstream_hits_inner.lock().expect("mutex should lock") += 1;
|
||||
(StatusCode::OK, Body::from("unexpected upstream hit"))
|
||||
}
|
||||
}),
|
||||
);
|
||||
|
||||
let mut manual_node = sample_proxy_node("proxy-node-manual");
|
||||
manual_node.name = "alpha-manual".to_string();
|
||||
manual_node.status = "online".to_string();
|
||||
manual_node.is_manual = true;
|
||||
manual_node.tunnel_mode = false;
|
||||
manual_node.tunnel_connected = false;
|
||||
manual_node.proxy_url = Some("http://proxy.example:8080".to_string());
|
||||
manual_node.proxy_username = Some("alice".to_string());
|
||||
manual_node.proxy_password = Some("supersecret".to_string());
|
||||
manual_node.last_heartbeat_at_unix_secs = None;
|
||||
manual_node.tunnel_connected_at_unix_secs = None;
|
||||
|
||||
let mut tunnel_node = sample_proxy_node("proxy-node-tunnel");
|
||||
tunnel_node.name = "zeta-tunnel".to_string();
|
||||
tunnel_node.status = "offline".to_string();
|
||||
|
||||
let proxy_node_repository = Arc::new(InMemoryProxyNodeRepository::seed(vec![
|
||||
tunnel_node,
|
||||
manual_node,
|
||||
]));
|
||||
|
||||
let (upstream_url, upstream_handle) = start_server(upstream).await;
|
||||
let gateway = build_router_with_state(
|
||||
AppState::new()
|
||||
.expect("gateway should build")
|
||||
.with_data_state_for_tests(GatewayDataState::with_proxy_node_repository_for_tests(
|
||||
proxy_node_repository,
|
||||
)),
|
||||
);
|
||||
let (gateway_url, gateway_handle) = start_server(gateway).await;
|
||||
|
||||
let response = reqwest::Client::new()
|
||||
.get(format!(
|
||||
"{gateway_url}/api/admin/proxy-nodes?status=online&skip=0&limit=10"
|
||||
))
|
||||
.header(crate::constants::GATEWAY_HEADER, "rust-phase3b")
|
||||
.header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
|
||||
.header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
|
||||
.header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
|
||||
.send()
|
||||
.await
|
||||
.expect("request should succeed");
|
||||
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
let payload: serde_json::Value = response.json().await.expect("json body should parse");
|
||||
assert_eq!(payload["total"], 1);
|
||||
assert_eq!(payload["skip"], 0);
|
||||
assert_eq!(payload["limit"], 10);
|
||||
|
||||
let items = payload["items"].as_array().expect("items should be array");
|
||||
assert_eq!(items.len(), 1);
|
||||
assert_eq!(items[0]["id"], "proxy-node-manual");
|
||||
assert_eq!(items[0]["name"], "alpha-manual");
|
||||
assert_eq!(items[0]["status"], "online");
|
||||
assert_eq!(items[0]["is_manual"], true);
|
||||
assert_eq!(items[0]["proxy_url"], "http://proxy.example:8080");
|
||||
assert_eq!(items[0]["proxy_username"], "alice");
|
||||
assert_eq!(items[0]["proxy_password"], "su****et");
|
||||
assert!(items[0]["created_at"].is_string());
|
||||
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
|
||||
|
||||
gateway_handle.abort();
|
||||
upstream_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_rejects_admin_proxy_nodes_unavailable_routes_locally() {
|
||||
let upstream_hits = Arc::new(Mutex::new(0usize));
|
||||
let upstream_hits_clone = Arc::clone(&upstream_hits);
|
||||
let upstream = Router::new().route(
|
||||
"/api/admin/proxy-nodes/register",
|
||||
any(move |_request: Request| {
|
||||
let upstream_hits_inner = Arc::clone(&upstream_hits_clone);
|
||||
async move {
|
||||
*upstream_hits_inner.lock().expect("mutex should lock") += 1;
|
||||
(StatusCode::OK, Body::from("unexpected upstream hit"))
|
||||
}
|
||||
}),
|
||||
);
|
||||
|
||||
let (upstream_url, upstream_handle) = start_server(upstream).await;
|
||||
let gateway = build_router_with_state(AppState::new().expect("gateway should build"));
|
||||
let (gateway_url, gateway_handle) = start_server(gateway).await;
|
||||
let client = reqwest::Client::new();
|
||||
|
||||
let register_response = client
|
||||
.post(format!("{gateway_url}/api/admin/proxy-nodes/register"))
|
||||
.header(crate::constants::GATEWAY_HEADER, "rust-phase3b")
|
||||
.header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
|
||||
.header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
|
||||
.header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
|
||||
.json(&json!({
|
||||
"name": "proxy-1",
|
||||
"ip": "1.1.1.1",
|
||||
"port": 8080
|
||||
}))
|
||||
.send()
|
||||
.await
|
||||
.expect("request should succeed");
|
||||
assert_eq!(register_response.status(), StatusCode::SERVICE_UNAVAILABLE);
|
||||
let register_payload: serde_json::Value = register_response
|
||||
.json()
|
||||
.await
|
||||
.expect("json body should parse");
|
||||
assert_eq!(
|
||||
register_payload["detail"],
|
||||
"Admin proxy nodes data unavailable"
|
||||
);
|
||||
|
||||
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
|
||||
|
||||
gateway_handle.abort();
|
||||
upstream_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_handles_admin_proxy_node_events_locally_with_trusted_admin_principal() {
|
||||
let upstream_hits = Arc::new(Mutex::new(0usize));
|
||||
let upstream_hits_clone = Arc::clone(&upstream_hits);
|
||||
let upstream = Router::new().route(
|
||||
"/api/admin/proxy-nodes/node-1/events",
|
||||
any(move |_request: Request| {
|
||||
let upstream_hits_inner = Arc::clone(&upstream_hits_clone);
|
||||
async move {
|
||||
*upstream_hits_inner.lock().expect("mutex should lock") += 1;
|
||||
(StatusCode::OK, Body::from("unexpected upstream hit"))
|
||||
}
|
||||
}),
|
||||
);
|
||||
|
||||
let proxy_node_repository = Arc::new(InMemoryProxyNodeRepository::seed_with_events(
|
||||
vec![sample_proxy_node("node-1")],
|
||||
vec![
|
||||
StoredProxyNodeEvent {
|
||||
id: 1,
|
||||
node_id: "node-1".to_string(),
|
||||
event_type: "connected".to_string(),
|
||||
detail: Some("older".to_string()),
|
||||
created_at_unix_secs: Some(1_710_000_000),
|
||||
},
|
||||
StoredProxyNodeEvent {
|
||||
id: 2,
|
||||
node_id: "node-1".to_string(),
|
||||
event_type: "disconnected".to_string(),
|
||||
detail: Some("newer".to_string()),
|
||||
created_at_unix_secs: Some(1_710_000_100),
|
||||
},
|
||||
],
|
||||
));
|
||||
|
||||
let (upstream_url, upstream_handle) = start_server(upstream).await;
|
||||
let gateway = build_router_with_state(
|
||||
AppState::new()
|
||||
.expect("gateway should build")
|
||||
.with_data_state_for_tests(GatewayDataState::with_proxy_node_repository_for_tests(
|
||||
proxy_node_repository,
|
||||
)),
|
||||
);
|
||||
let (gateway_url, gateway_handle) = start_server(gateway).await;
|
||||
|
||||
let response = reqwest::Client::new()
|
||||
.get(format!(
|
||||
"{gateway_url}/api/admin/proxy-nodes/node-1/events?limit=1"
|
||||
))
|
||||
.header(crate::constants::GATEWAY_HEADER, "rust-phase3b")
|
||||
.header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
|
||||
.header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
|
||||
.header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
|
||||
.send()
|
||||
.await
|
||||
.expect("request should succeed");
|
||||
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
let payload: serde_json::Value = response.json().await.expect("json body should parse");
|
||||
let items = payload["items"].as_array().expect("items should be array");
|
||||
assert_eq!(items.len(), 1);
|
||||
assert_eq!(items[0]["id"], 2);
|
||||
assert_eq!(items[0]["event_type"], "disconnected");
|
||||
assert_eq!(items[0]["detail"], "newer");
|
||||
assert_eq!(items[0]["created_at"], "2024-03-09T16:01:40Z");
|
||||
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
|
||||
|
||||
gateway_handle.abort();
|
||||
upstream_handle.abort();
|
||||
}
|
||||
406
apps/aether-gateway/src/tests/control/admin/security.rs
Normal file
406
apps/aether-gateway/src/tests/control/admin/security.rs
Normal file
@@ -0,0 +1,406 @@
|
||||
use std::sync::{Arc, Mutex};
|
||||
|
||||
use axum::body::{Body, Bytes};
|
||||
use axum::routing::any;
|
||||
use axum::{extract::Request, Router};
|
||||
use http::{HeaderMap, HeaderValue, StatusCode};
|
||||
use serde_json::json;
|
||||
|
||||
use super::super::{build_router_with_state, start_server, AppState};
|
||||
use crate::audit::AdminAuditEvent;
|
||||
use crate::constants::{
|
||||
GATEWAY_HEADER, TRUSTED_ADMIN_MANAGEMENT_TOKEN_ID_HEADER, TRUSTED_ADMIN_SESSION_ID_HEADER,
|
||||
TRUSTED_ADMIN_USER_ID_HEADER, TRUSTED_ADMIN_USER_ROLE_HEADER,
|
||||
};
|
||||
use crate::control::resolve_public_request_context;
|
||||
use crate::handlers::admin::maybe_build_local_admin_security_response;
|
||||
|
||||
async fn send_admin_security_request(
|
||||
gateway: Router,
|
||||
method: reqwest::Method,
|
||||
path: &str,
|
||||
body: Option<serde_json::Value>,
|
||||
) -> (StatusCode, serde_json::Value, usize) {
|
||||
let upstream_hits = Arc::new(Mutex::new(0usize));
|
||||
let upstream_hits_clone = Arc::clone(&upstream_hits);
|
||||
let upstream = Router::new().route(
|
||||
path,
|
||||
any(move |_request: Request| {
|
||||
let upstream_hits_inner = Arc::clone(&upstream_hits_clone);
|
||||
async move {
|
||||
*upstream_hits_inner.lock().expect("mutex should lock") += 1;
|
||||
(StatusCode::OK, Body::from("unexpected upstream hit"))
|
||||
}
|
||||
}),
|
||||
);
|
||||
|
||||
let (upstream_url, upstream_handle) = start_server(upstream).await;
|
||||
let (gateway_url, gateway_handle) = start_server(gateway).await;
|
||||
|
||||
let client = reqwest::Client::new();
|
||||
let mut request = client
|
||||
.request(method, format!("{gateway_url}{path}"))
|
||||
.header(crate::constants::GATEWAY_HEADER, "rust-phase3b")
|
||||
.header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
|
||||
.header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
|
||||
.header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123");
|
||||
if let Some(body) = body {
|
||||
request = request.json(&body);
|
||||
}
|
||||
|
||||
let response = request.send().await.expect("request should succeed");
|
||||
let status = response.status();
|
||||
let payload: serde_json::Value = response.json().await.expect("json body should parse");
|
||||
let upstream_count = *upstream_hits.lock().expect("mutex should lock");
|
||||
|
||||
gateway_handle.abort();
|
||||
upstream_handle.abort();
|
||||
|
||||
(status, payload, upstream_count)
|
||||
}
|
||||
|
||||
fn trusted_admin_headers() -> HeaderMap {
|
||||
let mut headers = HeaderMap::new();
|
||||
headers.insert(GATEWAY_HEADER, HeaderValue::from_static("rust-phase3b"));
|
||||
headers.insert(
|
||||
TRUSTED_ADMIN_USER_ID_HEADER,
|
||||
HeaderValue::from_static("admin-user-123"),
|
||||
);
|
||||
headers.insert(
|
||||
TRUSTED_ADMIN_USER_ROLE_HEADER,
|
||||
HeaderValue::from_static("admin"),
|
||||
);
|
||||
headers.insert(
|
||||
TRUSTED_ADMIN_SESSION_ID_HEADER,
|
||||
HeaderValue::from_static("session-123"),
|
||||
);
|
||||
headers.insert(
|
||||
TRUSTED_ADMIN_MANAGEMENT_TOKEN_ID_HEADER,
|
||||
HeaderValue::from_static("management-token-123"),
|
||||
);
|
||||
headers
|
||||
}
|
||||
|
||||
async fn local_admin_security_response(
|
||||
state: &AppState,
|
||||
method: http::Method,
|
||||
uri: &str,
|
||||
body: Option<serde_json::Value>,
|
||||
) -> axum::response::Response<Body> {
|
||||
let headers = trusted_admin_headers();
|
||||
let request_context = resolve_public_request_context(
|
||||
state,
|
||||
&method,
|
||||
&uri.parse().expect("uri should parse"),
|
||||
&headers,
|
||||
"trace-123",
|
||||
)
|
||||
.await
|
||||
.expect("request context should resolve");
|
||||
let body_bytes = body.map(|value| Bytes::from(value.to_string()));
|
||||
maybe_build_local_admin_security_response(state, &request_context, body_bytes.as_ref())
|
||||
.await
|
||||
.expect("local security response should build")
|
||||
.expect("security route should resolve locally")
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_handles_admin_security_blacklist_add_locally_with_trusted_admin_principal() {
|
||||
let gateway = build_router_with_state(AppState::new().expect("gateway should build"));
|
||||
|
||||
let (status, payload, upstream_count) = send_admin_security_request(
|
||||
gateway,
|
||||
reqwest::Method::POST,
|
||||
"/api/admin/security/ip/blacklist",
|
||||
Some(json!({ "ip_address": "1.2.3.4", "reason": "manual", "ttl": 60 })),
|
||||
)
|
||||
.await;
|
||||
|
||||
assert_eq!(status, StatusCode::OK);
|
||||
assert_eq!(payload["success"], true);
|
||||
assert_eq!(payload["message"], "IP 1.2.3.4 已加入黑名单");
|
||||
assert_eq!(payload["reason"], "manual");
|
||||
assert_eq!(payload["ttl"], 60);
|
||||
assert_eq!(upstream_count, 0);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn local_admin_security_blacklist_add_attaches_explicit_audit() {
|
||||
let state = AppState::new().expect("gateway should build");
|
||||
let response = local_admin_security_response(
|
||||
&state,
|
||||
http::Method::POST,
|
||||
"/api/admin/security/ip/blacklist",
|
||||
Some(json!({ "ip_address": "1.2.3.4", "reason": "manual", "ttl": 60 })),
|
||||
)
|
||||
.await;
|
||||
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
let audit = response
|
||||
.extensions()
|
||||
.get::<AdminAuditEvent>()
|
||||
.cloned()
|
||||
.expect("blacklist add should attach audit");
|
||||
assert_eq!(audit.event_name, "admin_security_blacklist_added");
|
||||
assert_eq!(audit.action, "add_security_blacklist_entry");
|
||||
assert_eq!(audit.target_type, "security_blacklist_entry");
|
||||
assert_eq!(audit.target_id, "1.2.3.4");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_handles_admin_security_blacklist_remove_locally_with_trusted_admin_principal() {
|
||||
let gateway = build_router_with_state(
|
||||
AppState::new()
|
||||
.expect("gateway should build")
|
||||
.with_admin_security_blacklist_for_tests([(
|
||||
"1.2.3.4".to_string(),
|
||||
"manual".to_string(),
|
||||
)]),
|
||||
);
|
||||
|
||||
let (status, payload, upstream_count) = send_admin_security_request(
|
||||
gateway,
|
||||
reqwest::Method::DELETE,
|
||||
"/api/admin/security/ip/blacklist/1.2.3.4",
|
||||
None,
|
||||
)
|
||||
.await;
|
||||
|
||||
assert_eq!(status, StatusCode::OK);
|
||||
assert_eq!(payload["success"], true);
|
||||
assert_eq!(payload["message"], "IP 1.2.3.4 已从黑名单移除");
|
||||
assert_eq!(upstream_count, 0);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_rejects_admin_security_blacklist_remove_without_ip_address() {
|
||||
let gateway = build_router_with_state(AppState::new().expect("gateway should build"));
|
||||
|
||||
let (status, payload, upstream_count) = send_admin_security_request(
|
||||
gateway,
|
||||
reqwest::Method::DELETE,
|
||||
"/api/admin/security/ip/blacklist/",
|
||||
None,
|
||||
)
|
||||
.await;
|
||||
|
||||
assert_eq!(status, StatusCode::BAD_REQUEST);
|
||||
assert_eq!(payload["detail"], "缺少 ip_address");
|
||||
assert_eq!(upstream_count, 0);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_handles_admin_security_blacklist_stats_locally_with_trusted_admin_principal() {
|
||||
let gateway = build_router_with_state(
|
||||
AppState::new()
|
||||
.expect("gateway should build")
|
||||
.with_admin_security_blacklist_for_tests([
|
||||
("1.2.3.4".to_string(), "manual".to_string()),
|
||||
("5.6.7.8".to_string(), "abuse".to_string()),
|
||||
]),
|
||||
);
|
||||
|
||||
let (status, payload, upstream_count) = send_admin_security_request(
|
||||
gateway,
|
||||
reqwest::Method::GET,
|
||||
"/api/admin/security/ip/blacklist/stats",
|
||||
None,
|
||||
)
|
||||
.await;
|
||||
|
||||
assert_eq!(status, StatusCode::OK);
|
||||
assert_eq!(payload["available"], true);
|
||||
assert_eq!(payload["total"], 2);
|
||||
assert_eq!(upstream_count, 0);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn local_admin_security_blacklist_list_attaches_explicit_audit() {
|
||||
let state = AppState::new()
|
||||
.expect("gateway should build")
|
||||
.with_admin_security_blacklist_for_tests([
|
||||
("5.6.7.8".to_string(), "abuse".to_string()),
|
||||
("1.2.3.4".to_string(), "manual".to_string()),
|
||||
]);
|
||||
let response = local_admin_security_response(
|
||||
&state,
|
||||
http::Method::GET,
|
||||
"/api/admin/security/ip/blacklist",
|
||||
None,
|
||||
)
|
||||
.await;
|
||||
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
let audit = response
|
||||
.extensions()
|
||||
.get::<AdminAuditEvent>()
|
||||
.cloned()
|
||||
.expect("blacklist list should attach audit");
|
||||
assert_eq!(audit.event_name, "admin_security_blacklist_viewed");
|
||||
assert_eq!(audit.action, "view_security_blacklist");
|
||||
assert_eq!(audit.target_type, "security_blacklist");
|
||||
assert_eq!(audit.target_id, "global");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_handles_admin_security_whitelist_add_locally_with_trusted_admin_principal() {
|
||||
let gateway = build_router_with_state(AppState::new().expect("gateway should build"));
|
||||
|
||||
let (status, payload, upstream_count) = send_admin_security_request(
|
||||
gateway,
|
||||
reqwest::Method::POST,
|
||||
"/api/admin/security/ip/whitelist",
|
||||
Some(json!({ "ip_address": "1.2.3.4" })),
|
||||
)
|
||||
.await;
|
||||
|
||||
assert_eq!(status, StatusCode::OK);
|
||||
assert_eq!(payload["success"], true);
|
||||
assert_eq!(payload["message"], "IP 1.2.3.4 已加入白名单");
|
||||
assert_eq!(upstream_count, 0);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn local_admin_security_whitelist_add_attaches_explicit_audit() {
|
||||
let state = AppState::new().expect("gateway should build");
|
||||
let response = local_admin_security_response(
|
||||
&state,
|
||||
http::Method::POST,
|
||||
"/api/admin/security/ip/whitelist",
|
||||
Some(json!({ "ip_address": "1.2.3.4" })),
|
||||
)
|
||||
.await;
|
||||
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
let audit = response
|
||||
.extensions()
|
||||
.get::<AdminAuditEvent>()
|
||||
.cloned()
|
||||
.expect("whitelist add should attach audit");
|
||||
assert_eq!(audit.event_name, "admin_security_whitelist_added");
|
||||
assert_eq!(audit.action, "add_security_whitelist_entry");
|
||||
assert_eq!(audit.target_type, "security_whitelist_entry");
|
||||
assert_eq!(audit.target_id, "1.2.3.4");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_handles_admin_security_whitelist_remove_locally_with_trusted_admin_principal() {
|
||||
let gateway = build_router_with_state(
|
||||
AppState::new()
|
||||
.expect("gateway should build")
|
||||
.with_admin_security_whitelist_for_tests(["1.2.3.4".to_string()]),
|
||||
);
|
||||
|
||||
let (status, payload, upstream_count) = send_admin_security_request(
|
||||
gateway,
|
||||
reqwest::Method::DELETE,
|
||||
"/api/admin/security/ip/whitelist/1.2.3.4",
|
||||
None,
|
||||
)
|
||||
.await;
|
||||
|
||||
assert_eq!(status, StatusCode::OK);
|
||||
assert_eq!(payload["success"], true);
|
||||
assert_eq!(payload["message"], "IP 1.2.3.4 已从白名单移除");
|
||||
assert_eq!(upstream_count, 0);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_rejects_admin_security_whitelist_remove_without_ip_address() {
|
||||
let gateway = build_router_with_state(AppState::new().expect("gateway should build"));
|
||||
|
||||
let (status, payload, upstream_count) = send_admin_security_request(
|
||||
gateway,
|
||||
reqwest::Method::DELETE,
|
||||
"/api/admin/security/ip/whitelist/",
|
||||
None,
|
||||
)
|
||||
.await;
|
||||
|
||||
assert_eq!(status, StatusCode::BAD_REQUEST);
|
||||
assert_eq!(payload["detail"], "缺少 ip_address");
|
||||
assert_eq!(upstream_count, 0);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_handles_admin_security_whitelist_list_locally_with_trusted_admin_principal() {
|
||||
let gateway = build_router_with_state(
|
||||
AppState::new()
|
||||
.expect("gateway should build")
|
||||
.with_admin_security_whitelist_for_tests([
|
||||
"10.0.0.0/24".to_string(),
|
||||
"1.2.3.4".to_string(),
|
||||
]),
|
||||
);
|
||||
|
||||
let (status, payload, upstream_count) = send_admin_security_request(
|
||||
gateway,
|
||||
reqwest::Method::GET,
|
||||
"/api/admin/security/ip/whitelist",
|
||||
None,
|
||||
)
|
||||
.await;
|
||||
|
||||
assert_eq!(status, StatusCode::OK);
|
||||
assert_eq!(payload["whitelist"], json!(["1.2.3.4", "10.0.0.0/24"]));
|
||||
assert_eq!(payload["total"], 2);
|
||||
assert_eq!(upstream_count, 0);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn local_admin_security_whitelist_list_attaches_explicit_audit() {
|
||||
let state = AppState::new()
|
||||
.expect("gateway should build")
|
||||
.with_admin_security_whitelist_for_tests([
|
||||
"10.0.0.0/24".to_string(),
|
||||
"1.2.3.4".to_string(),
|
||||
]);
|
||||
let response = local_admin_security_response(
|
||||
&state,
|
||||
http::Method::GET,
|
||||
"/api/admin/security/ip/whitelist",
|
||||
None,
|
||||
)
|
||||
.await;
|
||||
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
let audit = response
|
||||
.extensions()
|
||||
.get::<AdminAuditEvent>()
|
||||
.cloned()
|
||||
.expect("whitelist list should attach audit");
|
||||
assert_eq!(audit.event_name, "admin_security_whitelist_viewed");
|
||||
assert_eq!(audit.action, "view_security_whitelist");
|
||||
assert_eq!(audit.target_type, "security_whitelist");
|
||||
assert_eq!(audit.target_id, "global");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_handles_admin_security_blacklist_list_locally_with_trusted_admin_principal() {
|
||||
let gateway = build_router_with_state(
|
||||
AppState::new()
|
||||
.expect("gateway should build")
|
||||
.with_admin_security_blacklist_for_tests([
|
||||
("5.6.7.8".to_string(), "abuse".to_string()),
|
||||
("1.2.3.4".to_string(), "manual".to_string()),
|
||||
]),
|
||||
);
|
||||
|
||||
let (status, payload, upstream_count) = send_admin_security_request(
|
||||
gateway,
|
||||
reqwest::Method::GET,
|
||||
"/api/admin/security/ip/blacklist",
|
||||
None,
|
||||
)
|
||||
.await;
|
||||
|
||||
assert_eq!(status, StatusCode::OK);
|
||||
assert_eq!(payload["total"], 2);
|
||||
let items = payload["items"].as_array().expect("items array exists");
|
||||
assert_eq!(items.len(), 2);
|
||||
assert_eq!(items[0]["ip_address"], "1.2.3.4");
|
||||
assert_eq!(items[0]["reason"], "manual");
|
||||
assert_eq!(items[1]["ip_address"], "5.6.7.8");
|
||||
assert_eq!(items[1]["reason"], "abuse");
|
||||
assert_eq!(upstream_count, 0);
|
||||
}
|
||||
1300
apps/aether-gateway/src/tests/control/admin/stats.rs
Normal file
1300
apps/aether-gateway/src/tests/control/admin/stats.rs
Normal file
File diff suppressed because it is too large
Load Diff
1506
apps/aether-gateway/src/tests/control/admin/system.rs
Normal file
1506
apps/aether-gateway/src/tests/control/admin/system.rs
Normal file
File diff suppressed because it is too large
Load Diff
1681
apps/aether-gateway/src/tests/control/admin/usage.rs
Normal file
1681
apps/aether-gateway/src/tests/control/admin/usage.rs
Normal file
File diff suppressed because it is too large
Load Diff
1395
apps/aether-gateway/src/tests/control/admin/users.rs
Normal file
1395
apps/aether-gateway/src/tests/control/admin/users.rs
Normal file
File diff suppressed because it is too large
Load Diff
944
apps/aether-gateway/src/tests/control/admin/video_tasks.rs
Normal file
944
apps/aether-gateway/src/tests/control/admin/video_tasks.rs
Normal file
@@ -0,0 +1,944 @@
|
||||
use std::sync::{Arc, Mutex};
|
||||
|
||||
use aether_crypto::DEVELOPMENT_ENCRYPTION_KEY;
|
||||
use aether_data::repository::provider_catalog::{
|
||||
InMemoryProviderCatalogReadRepository, ProviderCatalogReadRepository,
|
||||
};
|
||||
use aether_data::repository::video_tasks::{
|
||||
InMemoryVideoTaskRepository, UpsertVideoTask, VideoTaskStatus, VideoTaskWriteRepository,
|
||||
};
|
||||
use axum::body::{to_bytes, Body, Bytes};
|
||||
use axum::response::Response;
|
||||
use axum::routing::{any, get, post};
|
||||
use axum::{extract::Request, Json, Router};
|
||||
use http::{HeaderMap, HeaderValue, StatusCode};
|
||||
use serde_json::json;
|
||||
|
||||
use super::super::{
|
||||
build_router_with_state, build_state_with_execution_runtime_override, sample_endpoint,
|
||||
sample_key, sample_provider, start_server, AppState,
|
||||
};
|
||||
use crate::audit::AdminAuditEvent;
|
||||
use crate::constants::{
|
||||
GATEWAY_HEADER, TRUSTED_ADMIN_MANAGEMENT_TOKEN_ID_HEADER, TRUSTED_ADMIN_SESSION_ID_HEADER,
|
||||
TRUSTED_ADMIN_USER_ID_HEADER, TRUSTED_ADMIN_USER_ROLE_HEADER,
|
||||
};
|
||||
use crate::control::resolve_public_request_context;
|
||||
use crate::data::GatewayDataState;
|
||||
use crate::handlers::admin::maybe_build_local_admin_video_tasks_response;
|
||||
|
||||
fn trusted_admin_headers() -> HeaderMap {
|
||||
let mut headers = HeaderMap::new();
|
||||
headers.insert(GATEWAY_HEADER, HeaderValue::from_static("rust-phase3b"));
|
||||
headers.insert(
|
||||
TRUSTED_ADMIN_USER_ID_HEADER,
|
||||
HeaderValue::from_static("admin-user-123"),
|
||||
);
|
||||
headers.insert(
|
||||
TRUSTED_ADMIN_USER_ROLE_HEADER,
|
||||
HeaderValue::from_static("admin"),
|
||||
);
|
||||
headers.insert(
|
||||
TRUSTED_ADMIN_SESSION_ID_HEADER,
|
||||
HeaderValue::from_static("session-123"),
|
||||
);
|
||||
headers.insert(
|
||||
TRUSTED_ADMIN_MANAGEMENT_TOKEN_ID_HEADER,
|
||||
HeaderValue::from_static("management-token-123"),
|
||||
);
|
||||
headers
|
||||
}
|
||||
|
||||
async fn local_admin_video_tasks_response(
|
||||
state: &AppState,
|
||||
method: http::Method,
|
||||
uri: &str,
|
||||
_body: Option<serde_json::Value>,
|
||||
) -> axum::response::Response<Body> {
|
||||
let headers = trusted_admin_headers();
|
||||
let request_context = resolve_public_request_context(
|
||||
state,
|
||||
&method,
|
||||
&uri.parse().expect("uri should parse"),
|
||||
&headers,
|
||||
"trace-123",
|
||||
)
|
||||
.await
|
||||
.expect("request context should resolve");
|
||||
maybe_build_local_admin_video_tasks_response(state, &request_context)
|
||||
.await
|
||||
.expect("local video tasks response should build")
|
||||
.expect("video tasks route should resolve locally")
|
||||
}
|
||||
|
||||
fn sample_admin_video_task(
|
||||
id: &str,
|
||||
status: VideoTaskStatus,
|
||||
created_at_unix_secs: u64,
|
||||
user_id: &str,
|
||||
username: &str,
|
||||
provider_id: &str,
|
||||
model: &str,
|
||||
prompt: &str,
|
||||
) -> UpsertVideoTask {
|
||||
UpsertVideoTask {
|
||||
id: id.to_string(),
|
||||
short_id: Some(format!("short-{id}")),
|
||||
request_id: format!("request-{id}"),
|
||||
user_id: Some(user_id.to_string()),
|
||||
api_key_id: Some(format!("api-key-{id}")),
|
||||
username: Some(username.to_string()),
|
||||
api_key_name: Some("primary".to_string()),
|
||||
external_task_id: Some(format!("ext-{id}")),
|
||||
provider_id: Some(provider_id.to_string()),
|
||||
endpoint_id: Some("endpoint-1".to_string()),
|
||||
key_id: Some("provider-key-1".to_string()),
|
||||
client_api_format: Some("openai:video".to_string()),
|
||||
provider_api_format: Some("openai:video".to_string()),
|
||||
format_converted: false,
|
||||
model: Some(model.to_string()),
|
||||
prompt: Some(prompt.to_string()),
|
||||
original_request_body: Some(json!({ "prompt": prompt })),
|
||||
duration_seconds: Some(4),
|
||||
resolution: Some("720p".to_string()),
|
||||
aspect_ratio: Some("16:9".to_string()),
|
||||
size: Some("1280x720".to_string()),
|
||||
status,
|
||||
progress_percent: if matches!(status, VideoTaskStatus::Completed) {
|
||||
100
|
||||
} else {
|
||||
50
|
||||
},
|
||||
progress_message: Some("ok".to_string()),
|
||||
retry_count: 0,
|
||||
poll_interval_seconds: 10,
|
||||
next_poll_at_unix_secs: None,
|
||||
poll_count: 1,
|
||||
max_poll_count: 360,
|
||||
created_at_unix_secs,
|
||||
submitted_at_unix_secs: Some(created_at_unix_secs),
|
||||
completed_at_unix_secs: if matches!(status, VideoTaskStatus::Completed) {
|
||||
Some(created_at_unix_secs + 30)
|
||||
} else {
|
||||
None
|
||||
},
|
||||
updated_at_unix_secs: created_at_unix_secs + 5,
|
||||
error_code: None,
|
||||
error_message: None,
|
||||
video_url: Some(format!("https://example.com/{id}.mp4")),
|
||||
request_metadata: None,
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_handles_admin_video_tasks_list_locally_with_trusted_admin_principal() {
|
||||
let upstream_hits = Arc::new(Mutex::new(0usize));
|
||||
let upstream_hits_clone = Arc::clone(&upstream_hits);
|
||||
let upstream = Router::new().route(
|
||||
"/api/admin/video-tasks",
|
||||
any(move |_request: Request| {
|
||||
let upstream_hits_inner = Arc::clone(&upstream_hits_clone);
|
||||
async move {
|
||||
*upstream_hits_inner.lock().expect("mutex should lock") += 1;
|
||||
(StatusCode::OK, Body::from("unexpected upstream hit"))
|
||||
}
|
||||
}),
|
||||
);
|
||||
|
||||
let repository = Arc::new(InMemoryVideoTaskRepository::default());
|
||||
repository
|
||||
.upsert(sample_admin_video_task(
|
||||
"task-completed",
|
||||
VideoTaskStatus::Completed,
|
||||
1_710_000_100,
|
||||
"user-1",
|
||||
"alice",
|
||||
"provider-openai",
|
||||
"gpt-video",
|
||||
&"x".repeat(120),
|
||||
))
|
||||
.await
|
||||
.expect("task should upsert");
|
||||
repository
|
||||
.upsert(sample_admin_video_task(
|
||||
"task-processing",
|
||||
VideoTaskStatus::Processing,
|
||||
1_710_000_200,
|
||||
"user-2",
|
||||
"bob",
|
||||
"provider-anthropic",
|
||||
"claude-video",
|
||||
"short prompt",
|
||||
))
|
||||
.await
|
||||
.expect("task should upsert");
|
||||
|
||||
let provider_catalog_repository: Arc<dyn ProviderCatalogReadRepository> =
|
||||
Arc::new(InMemoryProviderCatalogReadRepository::seed(
|
||||
vec![
|
||||
sample_provider("provider-openai", "OpenAI", 10),
|
||||
sample_provider("provider-anthropic", "Anthropic", 20),
|
||||
],
|
||||
vec![],
|
||||
vec![],
|
||||
));
|
||||
let data_state = GatewayDataState::with_video_task_repository_and_provider_transport_for_tests(
|
||||
Arc::clone(&repository),
|
||||
provider_catalog_repository,
|
||||
DEVELOPMENT_ENCRYPTION_KEY,
|
||||
);
|
||||
|
||||
let (upstream_url, upstream_handle) = start_server(upstream).await;
|
||||
let gateway = build_router_with_state(
|
||||
AppState::new()
|
||||
.expect("gateway should build")
|
||||
.with_data_state_for_tests(data_state),
|
||||
);
|
||||
let (gateway_url, gateway_handle) = start_server(gateway).await;
|
||||
|
||||
let response = reqwest::Client::new()
|
||||
.get(format!(
|
||||
"{gateway_url}/api/admin/video-tasks?status=completed&page=1&page_size=20"
|
||||
))
|
||||
.header(crate::constants::GATEWAY_HEADER, "rust-phase3b")
|
||||
.header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
|
||||
.header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
|
||||
.header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
|
||||
.send()
|
||||
.await
|
||||
.expect("request should succeed");
|
||||
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
let payload: serde_json::Value = response.json().await.expect("json body should parse");
|
||||
assert_eq!(payload["total"], json!(1));
|
||||
assert_eq!(payload["page"], json!(1));
|
||||
assert_eq!(payload["page_size"], json!(20));
|
||||
assert_eq!(payload["pages"], json!(1));
|
||||
assert_eq!(payload["items"].as_array().map(Vec::len), Some(1));
|
||||
assert_eq!(payload["items"][0]["id"], "task-completed");
|
||||
assert_eq!(payload["items"][0]["username"], "alice");
|
||||
assert_eq!(payload["items"][0]["provider_name"], "OpenAI");
|
||||
assert_eq!(payload["items"][0]["status"], "completed");
|
||||
assert!(payload["items"][0]["prompt"]
|
||||
.as_str()
|
||||
.is_some_and(|value| value.ends_with("...")));
|
||||
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
|
||||
|
||||
gateway_handle.abort();
|
||||
upstream_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_handles_admin_video_tasks_stats_locally_with_trusted_admin_principal() {
|
||||
let upstream_hits = Arc::new(Mutex::new(0usize));
|
||||
let upstream_hits_clone = Arc::clone(&upstream_hits);
|
||||
let upstream = Router::new().route(
|
||||
"/api/admin/video-tasks/stats",
|
||||
any(move |_request: Request| {
|
||||
let upstream_hits_inner = Arc::clone(&upstream_hits_clone);
|
||||
async move {
|
||||
*upstream_hits_inner.lock().expect("mutex should lock") += 1;
|
||||
(StatusCode::OK, Body::from("unexpected upstream hit"))
|
||||
}
|
||||
}),
|
||||
);
|
||||
|
||||
let repository = Arc::new(InMemoryVideoTaskRepository::default());
|
||||
repository
|
||||
.upsert(sample_admin_video_task(
|
||||
"task-processing",
|
||||
VideoTaskStatus::Processing,
|
||||
1_710_000_000,
|
||||
"user-1",
|
||||
"alice",
|
||||
"provider-openai",
|
||||
"gpt-video",
|
||||
"prompt one",
|
||||
))
|
||||
.await
|
||||
.expect("task should upsert");
|
||||
repository
|
||||
.upsert(sample_admin_video_task(
|
||||
"task-completed",
|
||||
VideoTaskStatus::Completed,
|
||||
1_710_000_100,
|
||||
"user-2",
|
||||
"bob",
|
||||
"provider-openai",
|
||||
"gpt-video",
|
||||
"prompt two",
|
||||
))
|
||||
.await
|
||||
.expect("task should upsert");
|
||||
repository
|
||||
.upsert(sample_admin_video_task(
|
||||
"task-failed",
|
||||
VideoTaskStatus::Failed,
|
||||
1_710_000_200,
|
||||
"user-2",
|
||||
"bob",
|
||||
"provider-openai",
|
||||
"veo-2",
|
||||
"prompt three",
|
||||
))
|
||||
.await
|
||||
.expect("task should upsert");
|
||||
|
||||
let (upstream_url, upstream_handle) = start_server(upstream).await;
|
||||
let gateway = build_router_with_state(
|
||||
AppState::new()
|
||||
.expect("gateway should build")
|
||||
.with_video_task_data_repository_for_tests(repository),
|
||||
);
|
||||
let (gateway_url, gateway_handle) = start_server(gateway).await;
|
||||
|
||||
let response = reqwest::Client::new()
|
||||
.get(format!("{gateway_url}/api/admin/video-tasks/stats"))
|
||||
.header(crate::constants::GATEWAY_HEADER, "rust-phase3b")
|
||||
.header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
|
||||
.header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
|
||||
.header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
|
||||
.send()
|
||||
.await
|
||||
.expect("request should succeed");
|
||||
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
let payload: serde_json::Value = response.json().await.expect("json body should parse");
|
||||
assert_eq!(payload["total"], json!(3));
|
||||
assert_eq!(payload["today_count"], json!(0));
|
||||
assert_eq!(payload["active_users"], json!(2));
|
||||
assert_eq!(payload["processing_count"], json!(1));
|
||||
assert_eq!(payload["by_status"]["processing"], json!(1));
|
||||
assert_eq!(payload["by_status"]["completed"], json!(1));
|
||||
assert_eq!(payload["by_status"]["failed"], json!(1));
|
||||
assert_eq!(payload["by_model"]["gpt-video"], json!(2));
|
||||
assert_eq!(payload["by_model"]["veo-2"], json!(1));
|
||||
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
|
||||
|
||||
gateway_handle.abort();
|
||||
upstream_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_handles_admin_video_task_detail_locally_with_trusted_admin_principal() {
|
||||
let upstream_hits = Arc::new(Mutex::new(0usize));
|
||||
let upstream_hits_clone = Arc::clone(&upstream_hits);
|
||||
let upstream = Router::new().route(
|
||||
"/api/admin/video-tasks/task-detail",
|
||||
any(move |_request: Request| {
|
||||
let upstream_hits_inner = Arc::clone(&upstream_hits_clone);
|
||||
async move {
|
||||
*upstream_hits_inner.lock().expect("mutex should lock") += 1;
|
||||
(StatusCode::OK, Body::from("unexpected upstream hit"))
|
||||
}
|
||||
}),
|
||||
);
|
||||
|
||||
let repository = Arc::new(InMemoryVideoTaskRepository::default());
|
||||
repository
|
||||
.upsert(sample_admin_video_task(
|
||||
"task-detail",
|
||||
VideoTaskStatus::Completed,
|
||||
1_710_000_300,
|
||||
"user-9",
|
||||
"charlie",
|
||||
"provider-openai",
|
||||
"gpt-video",
|
||||
"detail prompt",
|
||||
))
|
||||
.await
|
||||
.expect("task should upsert");
|
||||
|
||||
let provider_catalog_repository: Arc<dyn ProviderCatalogReadRepository> =
|
||||
Arc::new(InMemoryProviderCatalogReadRepository::seed(
|
||||
vec![sample_provider("provider-openai", "OpenAI", 10)],
|
||||
vec![sample_endpoint(
|
||||
"endpoint-1",
|
||||
"provider-openai",
|
||||
"openai:video",
|
||||
"https://api.openai.example",
|
||||
)],
|
||||
vec![],
|
||||
));
|
||||
let data_state = GatewayDataState::with_video_task_repository_and_provider_transport_for_tests(
|
||||
Arc::clone(&repository),
|
||||
provider_catalog_repository,
|
||||
DEVELOPMENT_ENCRYPTION_KEY,
|
||||
);
|
||||
|
||||
let (upstream_url, upstream_handle) = start_server(upstream).await;
|
||||
let gateway = build_router_with_state(
|
||||
AppState::new()
|
||||
.expect("gateway should build")
|
||||
.with_data_state_for_tests(data_state),
|
||||
);
|
||||
let (gateway_url, gateway_handle) = start_server(gateway).await;
|
||||
|
||||
let response = reqwest::Client::new()
|
||||
.get(format!("{gateway_url}/api/admin/video-tasks/task-detail"))
|
||||
.header(crate::constants::GATEWAY_HEADER, "rust-phase3b")
|
||||
.header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
|
||||
.header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
|
||||
.header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
|
||||
.send()
|
||||
.await
|
||||
.expect("request should succeed");
|
||||
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
let payload: serde_json::Value = response.json().await.expect("json body should parse");
|
||||
assert_eq!(payload["id"], "task-detail");
|
||||
assert_eq!(payload["username"], "charlie");
|
||||
assert_eq!(payload["provider_name"], "OpenAI");
|
||||
assert_eq!(payload["endpoint"]["id"], "endpoint-1");
|
||||
assert_eq!(payload["endpoint"]["api_format"], "openai:video");
|
||||
assert_eq!(payload["status"], "completed");
|
||||
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
|
||||
|
||||
gateway_handle.abort();
|
||||
upstream_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn local_admin_video_task_detail_attaches_explicit_audit() {
|
||||
let repository = Arc::new(InMemoryVideoTaskRepository::default());
|
||||
repository
|
||||
.upsert(sample_admin_video_task(
|
||||
"task-detail-audit",
|
||||
VideoTaskStatus::Completed,
|
||||
1_710_000_350,
|
||||
"user-3",
|
||||
"dana",
|
||||
"provider-openai",
|
||||
"gpt-video",
|
||||
"detail audit prompt",
|
||||
))
|
||||
.await
|
||||
.expect("task should upsert");
|
||||
|
||||
let provider_catalog_repository: Arc<dyn ProviderCatalogReadRepository> =
|
||||
Arc::new(InMemoryProviderCatalogReadRepository::seed(
|
||||
vec![sample_provider("provider-openai", "OpenAI", 10)],
|
||||
vec![sample_endpoint(
|
||||
"endpoint-1",
|
||||
"provider-openai",
|
||||
"openai:video",
|
||||
"https://api.openai.example",
|
||||
)],
|
||||
vec![],
|
||||
));
|
||||
let data_state = GatewayDataState::with_video_task_repository_and_provider_transport_for_tests(
|
||||
Arc::clone(&repository),
|
||||
provider_catalog_repository,
|
||||
DEVELOPMENT_ENCRYPTION_KEY,
|
||||
);
|
||||
let state = AppState::new()
|
||||
.expect("gateway state should build")
|
||||
.with_data_state_for_tests(data_state);
|
||||
|
||||
let response = local_admin_video_tasks_response(
|
||||
&state,
|
||||
http::Method::GET,
|
||||
"/api/admin/video-tasks/task-detail-audit",
|
||||
None,
|
||||
)
|
||||
.await;
|
||||
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
let audit = response
|
||||
.extensions()
|
||||
.get::<AdminAuditEvent>()
|
||||
.cloned()
|
||||
.expect("video task detail should attach audit");
|
||||
assert_eq!(audit.event_name, "admin_video_task_detail_viewed");
|
||||
assert_eq!(audit.action, "view_video_task_detail");
|
||||
assert_eq!(audit.target_type, "video_task");
|
||||
assert_eq!(audit.target_id, "task-detail-audit");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_cancels_admin_video_task_locally_with_trusted_admin_principal() {
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
struct SeenExecutionRuntimeSyncRequest {
|
||||
method: String,
|
||||
url: String,
|
||||
authorization: String,
|
||||
}
|
||||
|
||||
let upstream_hits = Arc::new(Mutex::new(0usize));
|
||||
let upstream_hits_clone = Arc::clone(&upstream_hits);
|
||||
let upstream = Router::new().route(
|
||||
"/api/admin/video-tasks/task-openai-cancel/cancel",
|
||||
any(move |_request: Request| {
|
||||
let upstream_hits_inner = Arc::clone(&upstream_hits_clone);
|
||||
async move {
|
||||
*upstream_hits_inner.lock().expect("mutex should lock") += 1;
|
||||
(StatusCode::OK, Body::from("unexpected upstream hit"))
|
||||
}
|
||||
}),
|
||||
);
|
||||
|
||||
let seen_execution_runtime = Arc::new(Mutex::new(None::<SeenExecutionRuntimeSyncRequest>));
|
||||
let seen_execution_runtime_clone = Arc::clone(&seen_execution_runtime);
|
||||
let execution_runtime = Router::new().route(
|
||||
"/v1/execute/sync",
|
||||
any(move |request: Request| {
|
||||
let seen_execution_runtime_inner = Arc::clone(&seen_execution_runtime_clone);
|
||||
async move {
|
||||
let (_parts, body) = request.into_parts();
|
||||
let raw_body = to_bytes(body, usize::MAX).await.expect("body should read");
|
||||
let payload: serde_json::Value =
|
||||
serde_json::from_slice(&raw_body).expect("payload should parse");
|
||||
*seen_execution_runtime_inner
|
||||
.lock()
|
||||
.expect("mutex should lock") = Some(SeenExecutionRuntimeSyncRequest {
|
||||
method: payload
|
||||
.get("method")
|
||||
.and_then(|value| value.as_str())
|
||||
.unwrap_or_default()
|
||||
.to_string(),
|
||||
url: payload
|
||||
.get("url")
|
||||
.and_then(|value| value.as_str())
|
||||
.unwrap_or_default()
|
||||
.to_string(),
|
||||
authorization: payload
|
||||
.get("headers")
|
||||
.and_then(|value| value.get("authorization"))
|
||||
.and_then(|value| value.as_str())
|
||||
.unwrap_or_default()
|
||||
.to_string(),
|
||||
});
|
||||
Json(json!({
|
||||
"request_id": "cancel-openai-task-1",
|
||||
"status_code": 204,
|
||||
"headers": {},
|
||||
"telemetry": {
|
||||
"elapsed_ms": 11
|
||||
}
|
||||
}))
|
||||
}
|
||||
}),
|
||||
);
|
||||
|
||||
let repository = Arc::new(InMemoryVideoTaskRepository::default());
|
||||
let mut task = sample_admin_video_task(
|
||||
"task-openai-cancel",
|
||||
VideoTaskStatus::Processing,
|
||||
1_710_000_400,
|
||||
"user-1",
|
||||
"alice",
|
||||
"provider-openai",
|
||||
"sora-2",
|
||||
"cancel prompt",
|
||||
);
|
||||
task.external_task_id = Some("ext-video-task-123".to_string());
|
||||
task.request_metadata = Some(json!({
|
||||
"rust_local_snapshot": {
|
||||
"OpenAi": {
|
||||
"local_task_id": "task-openai-cancel",
|
||||
"upstream_task_id": "ext-video-task-123",
|
||||
"created_at_unix_secs": 1710000400,
|
||||
"user_id": "user-1",
|
||||
"api_key_id": "api-key-task-openai-cancel",
|
||||
"model": "sora-2",
|
||||
"prompt": "cancel prompt",
|
||||
"size": "1280x720",
|
||||
"seconds": "4",
|
||||
"remixed_from_video_id": null,
|
||||
"status": "Processing",
|
||||
"progress_percent": 50,
|
||||
"completed_at_unix_secs": null,
|
||||
"expires_at_unix_secs": null,
|
||||
"error_code": null,
|
||||
"error_message": null,
|
||||
"video_url": null,
|
||||
"persistence": {
|
||||
"request_id": "request-task-openai-cancel",
|
||||
"username": "alice",
|
||||
"api_key_name": "primary",
|
||||
"client_api_format": "openai:video",
|
||||
"provider_api_format": "openai:video",
|
||||
"original_request_body": {
|
||||
"prompt": "cancel prompt"
|
||||
},
|
||||
"format_converted": false
|
||||
},
|
||||
"transport": {
|
||||
"upstream_base_url": "https://api.openai.example",
|
||||
"provider_name": "openai-video",
|
||||
"provider_id": "provider-openai",
|
||||
"endpoint_id": "endpoint-1",
|
||||
"key_id": "provider-key-1",
|
||||
"headers": {
|
||||
"authorization": "Bearer sk-upstream-openai-video",
|
||||
"content-type": "application/json"
|
||||
},
|
||||
"content_type": "application/json",
|
||||
"model_name": "sora-2-upstream",
|
||||
"proxy": null,
|
||||
"tls_profile": null,
|
||||
"timeouts": null
|
||||
}
|
||||
}
|
||||
}
|
||||
}));
|
||||
repository
|
||||
.upsert(task)
|
||||
.await
|
||||
.expect("upsert should succeed");
|
||||
|
||||
let (upstream_url, upstream_handle) = start_server(upstream).await;
|
||||
let (execution_runtime_url, execution_runtime_handle) = start_server(execution_runtime).await;
|
||||
let gateway = build_router_with_state(
|
||||
build_state_with_execution_runtime_override(execution_runtime_url)
|
||||
.with_video_task_data_repository_for_tests(Arc::clone(&repository)),
|
||||
);
|
||||
let (gateway_url, gateway_handle) = start_server(gateway).await;
|
||||
|
||||
let response = reqwest::Client::new()
|
||||
.post(format!(
|
||||
"{gateway_url}/api/admin/video-tasks/task-openai-cancel/cancel"
|
||||
))
|
||||
.header(crate::constants::GATEWAY_HEADER, "rust-phase3b")
|
||||
.header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
|
||||
.header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
|
||||
.header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
|
||||
.send()
|
||||
.await
|
||||
.expect("request should succeed");
|
||||
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
assert_eq!(
|
||||
response
|
||||
.json::<serde_json::Value>()
|
||||
.await
|
||||
.expect("body should parse"),
|
||||
json!({
|
||||
"id": "task-openai-cancel",
|
||||
"status": "cancelled",
|
||||
"message": "Task cancelled successfully",
|
||||
})
|
||||
);
|
||||
|
||||
let detail = reqwest::Client::new()
|
||||
.get(format!(
|
||||
"{gateway_url}/api/admin/video-tasks/task-openai-cancel"
|
||||
))
|
||||
.header(crate::constants::GATEWAY_HEADER, "rust-phase3b")
|
||||
.header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
|
||||
.header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
|
||||
.header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
|
||||
.send()
|
||||
.await
|
||||
.expect("detail request should succeed");
|
||||
assert_eq!(detail.status(), StatusCode::OK);
|
||||
let detail_json: serde_json::Value = detail.json().await.expect("detail should parse");
|
||||
assert_eq!(detail_json["status"], "cancelled");
|
||||
assert_eq!(
|
||||
detail_json["request_metadata"]["rust_local_snapshot"]["OpenAi"]["status"],
|
||||
"Cancelled"
|
||||
);
|
||||
|
||||
let seen_execution_runtime_request = seen_execution_runtime
|
||||
.lock()
|
||||
.expect("mutex should lock")
|
||||
.clone()
|
||||
.expect("execution runtime should be called");
|
||||
assert_eq!(seen_execution_runtime_request.method, "DELETE");
|
||||
assert_eq!(
|
||||
seen_execution_runtime_request.url,
|
||||
"https://api.openai.example/v1/videos/ext-video-task-123"
|
||||
);
|
||||
assert_eq!(
|
||||
seen_execution_runtime_request.authorization,
|
||||
"Bearer sk-upstream-openai-video"
|
||||
);
|
||||
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
|
||||
|
||||
gateway_handle.abort();
|
||||
execution_runtime_handle.abort();
|
||||
upstream_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn local_admin_video_task_cancel_attaches_explicit_audit() {
|
||||
let repository = Arc::new(InMemoryVideoTaskRepository::default());
|
||||
let mut task = sample_admin_video_task(
|
||||
"task-cancel-audit",
|
||||
VideoTaskStatus::Processing,
|
||||
1_710_000_450,
|
||||
"user-4",
|
||||
"erin",
|
||||
"provider-openai",
|
||||
"gpt-video",
|
||||
"cancel audit prompt",
|
||||
);
|
||||
task.client_api_format = None;
|
||||
task.provider_api_format = None;
|
||||
task.request_metadata = Some(json!({}));
|
||||
repository.upsert(task).await.expect("task should upsert");
|
||||
|
||||
let state = AppState::new()
|
||||
.expect("gateway state should build")
|
||||
.with_video_task_data_repository_for_tests(repository);
|
||||
|
||||
let response = local_admin_video_tasks_response(
|
||||
&state,
|
||||
http::Method::POST,
|
||||
"/api/admin/video-tasks/task-cancel-audit/cancel",
|
||||
None,
|
||||
)
|
||||
.await;
|
||||
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
let audit = response
|
||||
.extensions()
|
||||
.get::<AdminAuditEvent>()
|
||||
.cloned()
|
||||
.expect("video task cancel should attach audit");
|
||||
assert_eq!(audit.event_name, "admin_video_task_cancelled");
|
||||
assert_eq!(audit.action, "cancel_video_task");
|
||||
assert_eq!(audit.target_type, "video_task");
|
||||
assert_eq!(audit.target_id, "task-cancel-audit");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_redirects_admin_video_task_video_locally_with_trusted_admin_principal() {
|
||||
let upstream_hits = Arc::new(Mutex::new(0usize));
|
||||
let upstream_hits_clone = Arc::clone(&upstream_hits);
|
||||
let upstream = Router::new().route(
|
||||
"/api/admin/video-tasks/task-redirect/video",
|
||||
any(move |_request: Request| {
|
||||
let upstream_hits_inner = Arc::clone(&upstream_hits_clone);
|
||||
async move {
|
||||
*upstream_hits_inner.lock().expect("mutex should lock") += 1;
|
||||
(StatusCode::OK, Body::from("unexpected upstream hit"))
|
||||
}
|
||||
}),
|
||||
);
|
||||
|
||||
let repository = Arc::new(InMemoryVideoTaskRepository::default());
|
||||
repository
|
||||
.upsert(sample_admin_video_task(
|
||||
"task-redirect",
|
||||
VideoTaskStatus::Completed,
|
||||
1_710_000_500,
|
||||
"user-1",
|
||||
"alice",
|
||||
"provider-openai",
|
||||
"gpt-video",
|
||||
"redirect prompt",
|
||||
))
|
||||
.await
|
||||
.expect("task should upsert");
|
||||
|
||||
let (upstream_url, upstream_handle) = start_server(upstream).await;
|
||||
let gateway = build_router_with_state(
|
||||
AppState::new()
|
||||
.expect("gateway state should build")
|
||||
.with_video_task_data_repository_for_tests(repository),
|
||||
);
|
||||
let (gateway_url, gateway_handle) = start_server(gateway).await;
|
||||
|
||||
let client = reqwest::Client::builder()
|
||||
.redirect(reqwest::redirect::Policy::none())
|
||||
.build()
|
||||
.expect("client should build");
|
||||
let response = client
|
||||
.get(format!(
|
||||
"{gateway_url}/api/admin/video-tasks/task-redirect/video"
|
||||
))
|
||||
.header(crate::constants::GATEWAY_HEADER, "rust-phase3b")
|
||||
.header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
|
||||
.header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
|
||||
.header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
|
||||
.send()
|
||||
.await
|
||||
.expect("request should succeed");
|
||||
|
||||
assert_eq!(response.status(), StatusCode::TEMPORARY_REDIRECT);
|
||||
assert_eq!(
|
||||
response
|
||||
.headers()
|
||||
.get(http::header::LOCATION)
|
||||
.and_then(|value| value.to_str().ok()),
|
||||
Some("https://example.com/task-redirect.mp4")
|
||||
);
|
||||
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
|
||||
|
||||
gateway_handle.abort();
|
||||
upstream_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn local_admin_video_task_video_redirect_attaches_explicit_audit() {
|
||||
let repository = Arc::new(InMemoryVideoTaskRepository::default());
|
||||
repository
|
||||
.upsert(sample_admin_video_task(
|
||||
"task-video-audit",
|
||||
VideoTaskStatus::Completed,
|
||||
1_710_000_550,
|
||||
"user-5",
|
||||
"frank",
|
||||
"provider-openai",
|
||||
"gpt-video",
|
||||
"video audit prompt",
|
||||
))
|
||||
.await
|
||||
.expect("task should upsert");
|
||||
|
||||
let state = AppState::new()
|
||||
.expect("gateway state should build")
|
||||
.with_video_task_data_repository_for_tests(repository);
|
||||
|
||||
let response = local_admin_video_tasks_response(
|
||||
&state,
|
||||
http::Method::GET,
|
||||
"/api/admin/video-tasks/task-video-audit/video",
|
||||
None,
|
||||
)
|
||||
.await;
|
||||
|
||||
assert_eq!(response.status(), StatusCode::TEMPORARY_REDIRECT);
|
||||
let audit = response
|
||||
.extensions()
|
||||
.get::<AdminAuditEvent>()
|
||||
.cloned()
|
||||
.expect("video task video should attach audit");
|
||||
assert_eq!(audit.event_name, "admin_video_task_video_viewed");
|
||||
assert_eq!(audit.action, "view_video_task_video");
|
||||
assert_eq!(audit.target_type, "video_task_video");
|
||||
assert_eq!(audit.target_id, "task-video-audit");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_proxies_admin_video_task_video_locally_with_trusted_admin_principal() {
|
||||
let upstream_hits = Arc::new(Mutex::new(0usize));
|
||||
let upstream_hits_clone = Arc::clone(&upstream_hits);
|
||||
let seen_api_key = Arc::new(Mutex::new(None::<String>));
|
||||
let seen_api_key_clone = Arc::clone(&seen_api_key);
|
||||
let upstream = Router::new()
|
||||
.route(
|
||||
"/api/admin/video-tasks/task-proxy/video",
|
||||
any(move |_request: Request| {
|
||||
let upstream_hits_inner = Arc::clone(&upstream_hits_clone);
|
||||
async move {
|
||||
*upstream_hits_inner.lock().expect("mutex should lock") += 1;
|
||||
(StatusCode::OK, Body::from("unexpected upstream hit"))
|
||||
}
|
||||
}),
|
||||
)
|
||||
.route(
|
||||
"/generativelanguage.googleapis.com/v1beta/files/video-task-123:download",
|
||||
any(move |request: Request| {
|
||||
let seen_api_key_inner = Arc::clone(&seen_api_key_clone);
|
||||
async move {
|
||||
*seen_api_key_inner.lock().expect("mutex should lock") = request
|
||||
.headers()
|
||||
.get("x-goog-api-key")
|
||||
.and_then(|value| value.to_str().ok())
|
||||
.map(ToOwned::to_owned);
|
||||
let mut response = Response::builder()
|
||||
.status(StatusCode::OK)
|
||||
.body(Body::from("proxied-video-bytes"))
|
||||
.expect("response should build");
|
||||
response.headers_mut().insert(
|
||||
http::header::CONTENT_TYPE,
|
||||
HeaderValue::from_static("video/mp4"),
|
||||
);
|
||||
response
|
||||
}
|
||||
}),
|
||||
);
|
||||
|
||||
let (upstream_url, upstream_handle) = start_server(upstream).await;
|
||||
let repository = Arc::new(InMemoryVideoTaskRepository::default());
|
||||
let mut task = sample_admin_video_task(
|
||||
"task-proxy",
|
||||
VideoTaskStatus::Completed,
|
||||
1_710_000_600,
|
||||
"user-1",
|
||||
"alice",
|
||||
"provider-gemini",
|
||||
"veo-3",
|
||||
"proxy prompt",
|
||||
);
|
||||
task.provider_id = Some("provider-gemini".to_string());
|
||||
task.endpoint_id = Some("endpoint-gemini".to_string());
|
||||
task.key_id = Some("key-gemini".to_string());
|
||||
task.client_api_format = Some("gemini:video".to_string());
|
||||
task.provider_api_format = Some("gemini:video".to_string());
|
||||
task.video_url = Some(format!(
|
||||
"{upstream_url}/generativelanguage.googleapis.com/v1beta/files/video-task-123:download"
|
||||
));
|
||||
repository.upsert(task).await.expect("task should upsert");
|
||||
|
||||
let provider_catalog_repository: Arc<dyn ProviderCatalogReadRepository> =
|
||||
Arc::new(InMemoryProviderCatalogReadRepository::seed(
|
||||
vec![sample_provider("provider-gemini", "Gemini", 10)],
|
||||
vec![sample_endpoint(
|
||||
"endpoint-gemini",
|
||||
"provider-gemini",
|
||||
"gemini:video",
|
||||
"https://generativelanguage.googleapis.com",
|
||||
)],
|
||||
vec![sample_key(
|
||||
"key-gemini",
|
||||
"provider-gemini",
|
||||
"gemini:video",
|
||||
"gemini-upstream-secret",
|
||||
)],
|
||||
));
|
||||
let data_state = GatewayDataState::with_video_task_repository_and_provider_transport_for_tests(
|
||||
Arc::clone(&repository),
|
||||
provider_catalog_repository,
|
||||
DEVELOPMENT_ENCRYPTION_KEY,
|
||||
);
|
||||
|
||||
let gateway = build_router_with_state(
|
||||
AppState::new()
|
||||
.expect("gateway state should build")
|
||||
.with_data_state_for_tests(data_state),
|
||||
);
|
||||
let (gateway_url, gateway_handle) = start_server(gateway).await;
|
||||
|
||||
let response = reqwest::Client::new()
|
||||
.get(format!(
|
||||
"{gateway_url}/api/admin/video-tasks/task-proxy/video"
|
||||
))
|
||||
.header(crate::constants::GATEWAY_HEADER, "rust-phase3b")
|
||||
.header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
|
||||
.header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
|
||||
.header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
|
||||
.send()
|
||||
.await
|
||||
.expect("request should succeed");
|
||||
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
assert_eq!(
|
||||
response
|
||||
.headers()
|
||||
.get(http::header::CONTENT_TYPE)
|
||||
.and_then(|value| value.to_str().ok()),
|
||||
Some("video/mp4")
|
||||
);
|
||||
assert_eq!(
|
||||
response
|
||||
.headers()
|
||||
.get(http::header::CONTENT_DISPOSITION)
|
||||
.and_then(|value| value.to_str().ok()),
|
||||
Some("inline; filename=\"video_task-proxy.mp4\"")
|
||||
);
|
||||
assert_eq!(
|
||||
response.bytes().await.expect("body should read"),
|
||||
Bytes::from_static(b"proxied-video-bytes")
|
||||
);
|
||||
assert_eq!(
|
||||
seen_api_key.lock().expect("mutex should lock").as_deref(),
|
||||
Some("gemini-upstream-secret")
|
||||
);
|
||||
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
|
||||
|
||||
gateway_handle.abort();
|
||||
upstream_handle.abort();
|
||||
}
|
||||
1600
apps/aether-gateway/src/tests/control/admin/wallets.rs
Normal file
1600
apps/aether-gateway/src/tests/control/admin/wallets.rs
Normal file
File diff suppressed because it is too large
Load Diff
569
apps/aether-gateway/src/tests/control/helpers.rs
Normal file
569
apps/aether-gateway/src/tests/control/helpers.rs
Normal file
@@ -0,0 +1,569 @@
|
||||
use aether_crypto::{encrypt_python_fernet_plaintext, DEVELOPMENT_ENCRYPTION_KEY};
|
||||
use chrono::{DateTime, Utc};
|
||||
use hmac::Mac;
|
||||
use serde_json::{json, Map, Value};
|
||||
use sha2::{Digest, Sha256};
|
||||
|
||||
use super::{
|
||||
RequestCandidateStatus, StoredAdminGlobalModel, StoredAdminProviderModel,
|
||||
StoredAuthApiKeySnapshot, StoredLdapModuleConfig, StoredManagementToken,
|
||||
StoredManagementTokenUserSummary, StoredManagementTokenWithUser, StoredOAuthProviderConfig,
|
||||
StoredOAuthProviderModuleConfig, StoredProviderActiveGlobalModel,
|
||||
StoredProviderCatalogEndpoint, StoredProviderCatalogKey, StoredProviderCatalogProvider,
|
||||
StoredProviderModelStats, StoredProviderQuotaSnapshot, StoredProxyNode,
|
||||
StoredPublicGlobalModel, StoredRequestCandidate,
|
||||
};
|
||||
use crate::AppState;
|
||||
|
||||
pub(super) fn sample_currently_usable_auth_snapshot(
|
||||
api_key_id: &str,
|
||||
user_id: &str,
|
||||
) -> StoredAuthApiKeySnapshot {
|
||||
StoredAuthApiKeySnapshot::new(
|
||||
user_id.to_string(),
|
||||
"alice".to_string(),
|
||||
Some("alice@example.com".to_string()),
|
||||
"user".to_string(),
|
||||
"local".to_string(),
|
||||
true,
|
||||
false,
|
||||
Some(serde_json::json!(["openai"])),
|
||||
Some(serde_json::json!(["openai:chat"])),
|
||||
Some(serde_json::json!(["gpt-5"])),
|
||||
api_key_id.to_string(),
|
||||
Some("default".to_string()),
|
||||
true,
|
||||
false,
|
||||
false,
|
||||
Some(60),
|
||||
Some(5),
|
||||
Some(4_102_444_800),
|
||||
Some(serde_json::json!(["openai"])),
|
||||
Some(serde_json::json!(["openai:chat"])),
|
||||
Some(serde_json::json!(["gpt-5"])),
|
||||
)
|
||||
.expect("auth snapshot should build")
|
||||
}
|
||||
|
||||
pub(super) fn sample_expired_auth_snapshot(
|
||||
api_key_id: &str,
|
||||
user_id: &str,
|
||||
) -> StoredAuthApiKeySnapshot {
|
||||
let mut snapshot = sample_currently_usable_auth_snapshot(api_key_id, user_id);
|
||||
snapshot.api_key_expires_at_unix_secs = Some(1);
|
||||
snapshot
|
||||
}
|
||||
|
||||
pub(super) fn sample_locked_auth_snapshot(
|
||||
api_key_id: &str,
|
||||
user_id: &str,
|
||||
) -> StoredAuthApiKeySnapshot {
|
||||
let mut snapshot = sample_currently_usable_auth_snapshot(api_key_id, user_id);
|
||||
snapshot.api_key_is_locked = true;
|
||||
snapshot
|
||||
}
|
||||
|
||||
pub(super) fn hash_api_key(value: &str) -> String {
|
||||
let mut hasher = Sha256::new();
|
||||
hasher.update(value.as_bytes());
|
||||
format!("{:x}", hasher.finalize())
|
||||
}
|
||||
|
||||
pub(super) fn test_auth_secret() -> String {
|
||||
std::env::var("JWT_SECRET_KEY")
|
||||
.ok()
|
||||
.map(|value| value.trim().to_string())
|
||||
.filter(|value| !value.is_empty())
|
||||
.unwrap_or_else(|| "aether-rust-dev-jwt-secret".to_string())
|
||||
}
|
||||
|
||||
pub(super) fn build_test_auth_token(
|
||||
token_type: &str,
|
||||
mut payload: serde_json::Map<String, serde_json::Value>,
|
||||
expires_at: chrono::DateTime<chrono::Utc>,
|
||||
) -> String {
|
||||
use base64::Engine as _;
|
||||
|
||||
let header = json!({ "alg": "HS256", "typ": "JWT" });
|
||||
payload.insert("exp".to_string(), json!(expires_at.timestamp()));
|
||||
payload.insert("type".to_string(), json!(token_type));
|
||||
let header_segment = base64::engine::general_purpose::URL_SAFE_NO_PAD.encode(
|
||||
serde_json::to_vec(&header)
|
||||
.expect("jwt header should serialize")
|
||||
.as_slice(),
|
||||
);
|
||||
let payload_segment = base64::engine::general_purpose::URL_SAFE_NO_PAD.encode(
|
||||
serde_json::to_vec(&payload)
|
||||
.expect("jwt payload should serialize")
|
||||
.as_slice(),
|
||||
);
|
||||
let signing_input = format!("{header_segment}.{payload_segment}");
|
||||
let mut mac = hmac::Hmac::<sha2::Sha256>::new_from_slice(test_auth_secret().as_bytes())
|
||||
.expect("jwt secret should build");
|
||||
mac.update(signing_input.as_bytes());
|
||||
let signature = mac.finalize().into_bytes();
|
||||
format!(
|
||||
"{header_segment}.{payload_segment}.{}",
|
||||
base64::engine::general_purpose::URL_SAFE_NO_PAD.encode(signature.as_slice())
|
||||
)
|
||||
}
|
||||
|
||||
pub(super) async fn issue_test_admin_access_token(
|
||||
state: &AppState,
|
||||
client_device_id: &str,
|
||||
) -> String {
|
||||
let user = state
|
||||
.create_local_auth_user_with_settings(
|
||||
Some("admin@example.com".to_string()),
|
||||
true,
|
||||
"admin".to_string(),
|
||||
"hash".to_string(),
|
||||
"admin".to_string(),
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
)
|
||||
.await
|
||||
.expect("admin user should be created")
|
||||
.expect("admin user should exist");
|
||||
let now = chrono::Utc::now();
|
||||
let session_id = "session-admin-token".to_string();
|
||||
let refresh_token = "refresh-session-admin-token".to_string();
|
||||
let session = crate::data::state::StoredUserSessionRecord::new(
|
||||
session_id.clone(),
|
||||
user.id.clone(),
|
||||
client_device_id.to_string(),
|
||||
None,
|
||||
crate::data::state::StoredUserSessionRecord::hash_refresh_token(
|
||||
&refresh_token,
|
||||
),
|
||||
None,
|
||||
None,
|
||||
Some(now),
|
||||
Some(now + chrono::Duration::days(7)),
|
||||
None,
|
||||
None,
|
||||
Some("127.0.0.1".to_string()),
|
||||
Some("admin-test".to_string()),
|
||||
Some(now),
|
||||
Some(now),
|
||||
)
|
||||
.expect("session should build");
|
||||
state
|
||||
.create_user_session(session)
|
||||
.await
|
||||
.expect("session should persist")
|
||||
.expect("session should exist");
|
||||
build_test_auth_token(
|
||||
"access",
|
||||
serde_json::Map::from_iter([
|
||||
("user_id".to_string(), json!(user.id)),
|
||||
("role".to_string(), json!("admin")),
|
||||
(
|
||||
"created_at".to_string(),
|
||||
json!(user.created_at.map(|value| value.to_rfc3339())),
|
||||
),
|
||||
("session_id".to_string(), json!(session_id)),
|
||||
]),
|
||||
now + chrono::Duration::hours(12),
|
||||
)
|
||||
}
|
||||
|
||||
pub(super) fn sample_provider(
|
||||
id: &str,
|
||||
name: &str,
|
||||
priority: i32,
|
||||
) -> StoredProviderCatalogProvider {
|
||||
StoredProviderCatalogProvider::new(
|
||||
id.to_string(),
|
||||
name.to_string(),
|
||||
Some("https://example.com".to_string()),
|
||||
"custom".to_string(),
|
||||
)
|
||||
.expect("provider should build")
|
||||
.with_routing_fields(priority)
|
||||
}
|
||||
|
||||
pub(super) fn sample_proxy_node(node_id: &str) -> StoredProxyNode {
|
||||
StoredProxyNode::new(
|
||||
node_id.to_string(),
|
||||
"proxy-node-1".to_string(),
|
||||
"127.0.0.1".to_string(),
|
||||
0,
|
||||
false,
|
||||
"offline".to_string(),
|
||||
30,
|
||||
0,
|
||||
0,
|
||||
0,
|
||||
0,
|
||||
0,
|
||||
true,
|
||||
false,
|
||||
7,
|
||||
)
|
||||
.expect("proxy node should build")
|
||||
.with_runtime_fields(
|
||||
Some("test".to_string()),
|
||||
Some("system".to_string()),
|
||||
Some(1_710_000_000),
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
Some(1_710_000_010),
|
||||
Some(json!({
|
||||
"upgrade_to": "1.2.3",
|
||||
"allowed_ports": [443],
|
||||
})),
|
||||
Some(1_709_000_000),
|
||||
Some(1_710_000_100),
|
||||
)
|
||||
}
|
||||
|
||||
pub(super) fn sample_provider_quota(provider_id: &str) -> StoredProviderQuotaSnapshot {
|
||||
StoredProviderQuotaSnapshot::new(
|
||||
provider_id.to_string(),
|
||||
"monthly_quota".to_string(),
|
||||
Some(100.0),
|
||||
12.5,
|
||||
Some(30),
|
||||
Some(1_711_000_000),
|
||||
Some(1_711_000_000 + 30 * 24 * 60 * 60),
|
||||
true,
|
||||
)
|
||||
.expect("provider quota should build")
|
||||
}
|
||||
|
||||
pub(super) fn sample_provider_model_stats(
|
||||
provider_id: &str,
|
||||
total_models: i64,
|
||||
active_models: i64,
|
||||
) -> StoredProviderModelStats {
|
||||
StoredProviderModelStats::new(provider_id.to_string(), total_models, active_models)
|
||||
.expect("provider model stats should build")
|
||||
}
|
||||
|
||||
pub(super) fn sample_provider_active_global_model(
|
||||
provider_id: &str,
|
||||
global_model_id: &str,
|
||||
) -> StoredProviderActiveGlobalModel {
|
||||
StoredProviderActiveGlobalModel::new(provider_id.to_string(), global_model_id.to_string())
|
||||
.expect("provider active global model should build")
|
||||
}
|
||||
|
||||
pub(super) fn sample_management_token(
|
||||
token_id: &str,
|
||||
user_id: &str,
|
||||
username: &str,
|
||||
is_active: bool,
|
||||
) -> StoredManagementTokenWithUser {
|
||||
let token = StoredManagementToken::new(
|
||||
token_id.to_string(),
|
||||
user_id.to_string(),
|
||||
format!("{username}-token"),
|
||||
)
|
||||
.expect("management token should build")
|
||||
.with_display_fields(
|
||||
Some(format!("{username} token")),
|
||||
Some("ae_test".to_string()),
|
||||
Some(json!(["127.0.0.1"])),
|
||||
)
|
||||
.with_runtime_fields(
|
||||
Some(4_102_444_800),
|
||||
Some(1_711_000_000),
|
||||
Some("127.0.0.1".to_string()),
|
||||
7,
|
||||
is_active,
|
||||
)
|
||||
.with_timestamps(Some(1_710_000_000), Some(1_711_000_100));
|
||||
let user = StoredManagementTokenUserSummary::new(
|
||||
user_id.to_string(),
|
||||
Some(format!("{username}@example.com")),
|
||||
username.to_string(),
|
||||
"admin".to_string(),
|
||||
)
|
||||
.expect("management token user should build");
|
||||
StoredManagementTokenWithUser::new(token, user)
|
||||
}
|
||||
|
||||
pub(super) fn sample_admin_provider_model(
|
||||
id: &str,
|
||||
provider_id: &str,
|
||||
global_model_id: &str,
|
||||
provider_model_name: &str,
|
||||
) -> StoredAdminProviderModel {
|
||||
StoredAdminProviderModel::new(
|
||||
id.to_string(),
|
||||
provider_id.to_string(),
|
||||
global_model_id.to_string(),
|
||||
provider_model_name.to_string(),
|
||||
Some(json!([{"name": format!("{provider_model_name}-alias"), "priority": 1}])),
|
||||
Some(0.02),
|
||||
Some(json!({
|
||||
"tiers": [{
|
||||
"up_to": null,
|
||||
"input_price_per_1m": 3.0,
|
||||
"output_price_per_1m": 15.0,
|
||||
}]
|
||||
})),
|
||||
Some(true),
|
||||
Some(true),
|
||||
None,
|
||||
Some(false),
|
||||
Some(false),
|
||||
true,
|
||||
true,
|
||||
Some(json!({"billing": {"mode": "local"}, "provider_hint": provider_model_name})),
|
||||
Some(1_711_000_000),
|
||||
Some(1_711_000_100),
|
||||
Some("gpt-5".to_string()),
|
||||
Some("GPT 5".to_string()),
|
||||
Some(0.03),
|
||||
Some(json!({
|
||||
"tiers": [{
|
||||
"up_to": null,
|
||||
"input_price_per_1m": 4.0,
|
||||
"output_price_per_1m": 20.0,
|
||||
}]
|
||||
})),
|
||||
Some(json!({"streaming": true, "vision": false, "billing": {"currency": "USD"}})),
|
||||
)
|
||||
.expect("admin provider model should build")
|
||||
}
|
||||
|
||||
pub(super) fn sample_admin_global_model(
|
||||
id: &str,
|
||||
name: &str,
|
||||
display_name: &str,
|
||||
) -> StoredAdminGlobalModel {
|
||||
StoredAdminGlobalModel::new(
|
||||
id.to_string(),
|
||||
name.to_string(),
|
||||
display_name.to_string(),
|
||||
true,
|
||||
Some(0.03),
|
||||
Some(json!({
|
||||
"tiers": [{
|
||||
"up_to": null,
|
||||
"input_price_per_1m": 4.0,
|
||||
"output_price_per_1m": 20.0,
|
||||
}]
|
||||
})),
|
||||
Some(json!(["streaming", "vision"])),
|
||||
Some(json!({"streaming": true, "vision": false, "billing": {"currency": "USD"}})),
|
||||
Some(1_711_000_000),
|
||||
Some(1_711_000_100),
|
||||
)
|
||||
.expect("admin global model should build")
|
||||
}
|
||||
|
||||
pub(super) fn sample_public_global_model_with_mappings(
|
||||
id: &str,
|
||||
name: &str,
|
||||
display_name: &str,
|
||||
mappings: &[&str],
|
||||
) -> StoredPublicGlobalModel {
|
||||
StoredPublicGlobalModel::new(
|
||||
id.to_string(),
|
||||
name.to_string(),
|
||||
Some(display_name.to_string()),
|
||||
true,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
Some(json!({ "model_mappings": mappings })),
|
||||
0,
|
||||
)
|
||||
.expect("public global model should build")
|
||||
}
|
||||
|
||||
pub(super) fn sample_oauth_module_provider(
|
||||
provider_type: &str,
|
||||
display_name: &str,
|
||||
) -> StoredOAuthProviderModuleConfig {
|
||||
StoredOAuthProviderModuleConfig::new(
|
||||
provider_type.to_string(),
|
||||
display_name.to_string(),
|
||||
"client-id".to_string(),
|
||||
Some("encrypted-secret".to_string()),
|
||||
"https://example.com/oauth/callback".to_string(),
|
||||
)
|
||||
.expect("oauth module provider should build")
|
||||
}
|
||||
|
||||
pub(super) fn sample_ldap_module_config() -> StoredLdapModuleConfig {
|
||||
StoredLdapModuleConfig {
|
||||
server_url: "ldaps://ldap.example.com".to_string(),
|
||||
bind_dn: "cn=admin,dc=example,dc=com".to_string(),
|
||||
bind_password_encrypted: Some("encrypted-password".to_string()),
|
||||
base_dn: "dc=example,dc=com".to_string(),
|
||||
user_search_filter: Some("(uid={username})".to_string()),
|
||||
username_attr: Some("uid".to_string()),
|
||||
email_attr: Some("mail".to_string()),
|
||||
display_name_attr: Some("displayName".to_string()),
|
||||
is_enabled: true,
|
||||
is_exclusive: false,
|
||||
use_starttls: true,
|
||||
connect_timeout: Some(10),
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) fn sample_oauth_provider_config(provider_type: &str) -> StoredOAuthProviderConfig {
|
||||
StoredOAuthProviderConfig::new(
|
||||
provider_type.to_string(),
|
||||
"Linux Do".to_string(),
|
||||
"client-id".to_string(),
|
||||
"https://backend.example.com/oauth/callback".to_string(),
|
||||
"https://frontend.example.com/auth/callback".to_string(),
|
||||
)
|
||||
.expect("oauth provider config should build")
|
||||
.with_config_fields(
|
||||
Some(
|
||||
encrypt_python_fernet_plaintext(DEVELOPMENT_ENCRYPTION_KEY, "secret-value")
|
||||
.expect("secret should encrypt"),
|
||||
),
|
||||
Some("https://connect.linux.do/oauth2/authorize".to_string()),
|
||||
Some("https://connect.linux.do/oauth2/token".to_string()),
|
||||
Some("https://connect.linux.do/api/user".to_string()),
|
||||
Some(vec!["openid".to_string()]),
|
||||
Some(json!({"email": "email"})),
|
||||
Some(json!({"team": true})),
|
||||
true,
|
||||
)
|
||||
}
|
||||
|
||||
pub(super) fn sample_endpoint(
|
||||
id: &str,
|
||||
provider_id: &str,
|
||||
api_format: &str,
|
||||
base_url: &str,
|
||||
) -> StoredProviderCatalogEndpoint {
|
||||
StoredProviderCatalogEndpoint::new(
|
||||
id.to_string(),
|
||||
provider_id.to_string(),
|
||||
api_format.to_string(),
|
||||
None,
|
||||
None,
|
||||
true,
|
||||
)
|
||||
.expect("endpoint should build")
|
||||
.with_health_score(0.9)
|
||||
.with_transport_fields(
|
||||
base_url.to_string(),
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
)
|
||||
.expect("endpoint transport should build")
|
||||
}
|
||||
|
||||
pub(super) fn sample_key(
|
||||
id: &str,
|
||||
provider_id: &str,
|
||||
api_format: &str,
|
||||
secret: &str,
|
||||
) -> StoredProviderCatalogKey {
|
||||
let encrypted_api_key = encrypt_python_fernet_plaintext(DEVELOPMENT_ENCRYPTION_KEY, secret)
|
||||
.expect("api key ciphertext should build");
|
||||
StoredProviderCatalogKey::new(
|
||||
id.to_string(),
|
||||
provider_id.to_string(),
|
||||
"default".to_string(),
|
||||
"api_key".to_string(),
|
||||
None,
|
||||
true,
|
||||
)
|
||||
.expect("key should build")
|
||||
.with_transport_fields(
|
||||
Some(json!([api_format])),
|
||||
encrypted_api_key,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
)
|
||||
.expect("key transport should build")
|
||||
}
|
||||
|
||||
pub(super) fn sample_request_candidate(
|
||||
id: &str,
|
||||
request_id: &str,
|
||||
endpoint_id: &str,
|
||||
status: RequestCandidateStatus,
|
||||
created_at_unix_secs: i64,
|
||||
finished_at_unix_secs: Option<i64>,
|
||||
) -> StoredRequestCandidate {
|
||||
StoredRequestCandidate::new(
|
||||
id.to_string(),
|
||||
request_id.to_string(),
|
||||
Some("user-1".to_string()),
|
||||
Some("api-key-1".to_string()),
|
||||
Some("alice".to_string()),
|
||||
Some("default".to_string()),
|
||||
0,
|
||||
0,
|
||||
Some("provider-1".to_string()),
|
||||
Some(endpoint_id.to_string()),
|
||||
Some("key-1".to_string()),
|
||||
status,
|
||||
None,
|
||||
false,
|
||||
Some(200),
|
||||
matches!(status, RequestCandidateStatus::Failed).then_some("rate_limit".to_string()),
|
||||
None,
|
||||
Some(120),
|
||||
Some(1),
|
||||
None,
|
||||
None,
|
||||
created_at_unix_secs,
|
||||
Some(created_at_unix_secs),
|
||||
finished_at_unix_secs,
|
||||
)
|
||||
.expect("request candidate should build")
|
||||
}
|
||||
|
||||
pub(super) fn sample_recent_key_rpm_candidate(
|
||||
id: &str,
|
||||
request_id: &str,
|
||||
endpoint_id: &str,
|
||||
key_id: &str,
|
||||
now_unix_secs: i64,
|
||||
concurrent_requests: i32,
|
||||
) -> StoredRequestCandidate {
|
||||
StoredRequestCandidate::new(
|
||||
id.to_string(),
|
||||
request_id.to_string(),
|
||||
Some("user-1".to_string()),
|
||||
Some("api-key-1".to_string()),
|
||||
Some("alice".to_string()),
|
||||
Some("default".to_string()),
|
||||
0,
|
||||
0,
|
||||
Some("provider-1".to_string()),
|
||||
Some(endpoint_id.to_string()),
|
||||
Some(key_id.to_string()),
|
||||
RequestCandidateStatus::Success,
|
||||
None,
|
||||
false,
|
||||
Some(200),
|
||||
None,
|
||||
None,
|
||||
Some(120),
|
||||
Some(concurrent_requests),
|
||||
None,
|
||||
None,
|
||||
now_unix_secs - 10,
|
||||
Some(now_unix_secs - 10),
|
||||
Some(now_unix_secs - 8),
|
||||
)
|
||||
.expect("request candidate should build")
|
||||
}
|
||||
312
apps/aether-gateway/src/tests/control/internal.rs
Normal file
312
apps/aether-gateway/src/tests/control/internal.rs
Normal file
@@ -0,0 +1,312 @@
|
||||
use std::sync::{Arc, Mutex};
|
||||
|
||||
use axum::body::Body;
|
||||
use axum::routing::{any, post};
|
||||
use axum::{extract::Request, Json, Router};
|
||||
use http::header::HeaderValue;
|
||||
use http::StatusCode;
|
||||
use serde_json::json;
|
||||
|
||||
use super::{
|
||||
build_router_with_state, sample_proxy_node, start_server, AppState, GatewayDataState,
|
||||
InMemoryProxyNodeRepository, TRACE_ID_HEADER,
|
||||
};
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_handles_internal_tunnel_heartbeat_locally_with_loopback() {
|
||||
let upstream_hits = Arc::new(Mutex::new(0usize));
|
||||
let upstream_hits_clone = Arc::clone(&upstream_hits);
|
||||
let upstream = Router::new().route(
|
||||
"/api/internal/tunnel/heartbeat",
|
||||
any(move |_request: Request| {
|
||||
let upstream_hits_inner = Arc::clone(&upstream_hits_clone);
|
||||
async move {
|
||||
*upstream_hits_inner.lock().expect("mutex should lock") += 1;
|
||||
(StatusCode::OK, Body::from("unexpected upstream hit"))
|
||||
}
|
||||
}),
|
||||
);
|
||||
|
||||
let repository = Arc::new(InMemoryProxyNodeRepository::seed(vec![sample_proxy_node(
|
||||
"node-123",
|
||||
)]));
|
||||
|
||||
let (upstream_url, upstream_handle) = start_server(upstream).await;
|
||||
let gateway = build_router_with_state(
|
||||
AppState::new()
|
||||
.expect("gateway should build")
|
||||
.with_data_state_for_tests(GatewayDataState::with_proxy_node_repository_for_tests(
|
||||
Arc::clone(&repository),
|
||||
)),
|
||||
);
|
||||
let (gateway_url, gateway_handle) = start_server(gateway).await;
|
||||
|
||||
let response = reqwest::Client::new()
|
||||
.post(format!("{gateway_url}/api/internal/tunnel/heartbeat"))
|
||||
.json(&json!({
|
||||
"node_id": "node-123",
|
||||
"heartbeat_interval": 45,
|
||||
"active_connections": 5,
|
||||
"total_requests": 9,
|
||||
"avg_latency_ms": 12.5,
|
||||
"failed_requests": 1,
|
||||
"dns_failures": 2,
|
||||
"stream_errors": 3,
|
||||
"proxy_metadata": {"arch": "arm64"},
|
||||
"proxy_version": "2.0.0",
|
||||
}))
|
||||
.send()
|
||||
.await
|
||||
.expect("request should succeed");
|
||||
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
let payload: serde_json::Value = response.json().await.expect("json body should parse");
|
||||
assert_eq!(payload["config_version"], 7);
|
||||
assert_eq!(payload["upgrade_to"], "1.2.3");
|
||||
assert_eq!(payload["remote_config"]["allowed_ports"][0], 443);
|
||||
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
|
||||
|
||||
gateway_handle.abort();
|
||||
upstream_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_handles_internal_tunnel_node_status_locally_with_loopback() {
|
||||
let upstream_hits = Arc::new(Mutex::new(0usize));
|
||||
let upstream_hits_clone = Arc::clone(&upstream_hits);
|
||||
let upstream = Router::new().route(
|
||||
"/api/internal/tunnel/node-status",
|
||||
any(move |_request: Request| {
|
||||
let upstream_hits_inner = Arc::clone(&upstream_hits_clone);
|
||||
async move {
|
||||
*upstream_hits_inner.lock().expect("mutex should lock") += 1;
|
||||
(StatusCode::OK, Body::from("unexpected upstream hit"))
|
||||
}
|
||||
}),
|
||||
);
|
||||
|
||||
let repository = Arc::new(InMemoryProxyNodeRepository::seed(vec![sample_proxy_node(
|
||||
"node-123",
|
||||
)]));
|
||||
|
||||
let (upstream_url, upstream_handle) = start_server(upstream).await;
|
||||
let gateway = build_router_with_state(
|
||||
AppState::new()
|
||||
.expect("gateway should build")
|
||||
.with_data_state_for_tests(GatewayDataState::with_proxy_node_repository_for_tests(
|
||||
Arc::clone(&repository),
|
||||
)),
|
||||
);
|
||||
let (gateway_url, gateway_handle) = start_server(gateway).await;
|
||||
|
||||
let response = reqwest::Client::new()
|
||||
.post(format!("{gateway_url}/api/internal/tunnel/node-status"))
|
||||
.json(&json!({
|
||||
"node_id": "node-123",
|
||||
"connected": true,
|
||||
"conn_count": 4,
|
||||
}))
|
||||
.send()
|
||||
.await
|
||||
.expect("request should succeed");
|
||||
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
let payload: serde_json::Value = response.json().await.expect("json body should parse");
|
||||
assert_eq!(payload["updated"], json!(true));
|
||||
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
|
||||
|
||||
gateway_handle.abort();
|
||||
upstream_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_owns_proxy_tunnel_path_without_proxying_upstream() {
|
||||
let upstream_hits = Arc::new(Mutex::new(0usize));
|
||||
let upstream_hits_clone = Arc::clone(&upstream_hits);
|
||||
let upstream = Router::new().route(
|
||||
"/api/internal/proxy-tunnel",
|
||||
any(move |_request: Request| {
|
||||
let upstream_hits_inner = Arc::clone(&upstream_hits_clone);
|
||||
async move {
|
||||
*upstream_hits_inner.lock().expect("mutex should lock") += 1;
|
||||
(StatusCode::OK, Body::from("unexpected upstream hit"))
|
||||
}
|
||||
}),
|
||||
);
|
||||
|
||||
let (upstream_url, upstream_handle) = start_server(upstream).await;
|
||||
let gateway = build_router_with_state(AppState::new().expect("gateway should build"));
|
||||
let (gateway_url, gateway_handle) = start_server(gateway).await;
|
||||
|
||||
let response = reqwest::Client::new()
|
||||
.get(format!("{gateway_url}/api/internal/proxy-tunnel"))
|
||||
.header("x-node-id", "node-123")
|
||||
.send()
|
||||
.await
|
||||
.expect("request should succeed");
|
||||
|
||||
assert_eq!(response.status(), StatusCode::BAD_REQUEST);
|
||||
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
|
||||
|
||||
gateway_handle.abort();
|
||||
upstream_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_handles_internal_tunnel_relay_locally_without_proxying_upstream() {
|
||||
let upstream_hits = Arc::new(Mutex::new(0usize));
|
||||
let upstream_hits_clone = Arc::clone(&upstream_hits);
|
||||
let upstream = Router::new().route(
|
||||
"/api/internal/tunnel/relay/node-123",
|
||||
any(move |_request: Request| {
|
||||
let upstream_hits_inner = Arc::clone(&upstream_hits_clone);
|
||||
async move {
|
||||
*upstream_hits_inner.lock().expect("mutex should lock") += 1;
|
||||
(StatusCode::OK, Body::from("unexpected upstream hit"))
|
||||
}
|
||||
}),
|
||||
);
|
||||
|
||||
let (upstream_url, upstream_handle) = start_server(upstream).await;
|
||||
let gateway = build_router_with_state(AppState::new().expect("gateway should build"));
|
||||
let (gateway_url, gateway_handle) = start_server(gateway).await;
|
||||
|
||||
let response = reqwest::Client::new()
|
||||
.post(format!("{gateway_url}/api/internal/tunnel/relay/node-123"))
|
||||
.body(Vec::<u8>::new())
|
||||
.send()
|
||||
.await
|
||||
.expect("request should succeed");
|
||||
|
||||
assert_eq!(response.status(), StatusCode::BAD_REQUEST);
|
||||
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
|
||||
|
||||
gateway_handle.abort();
|
||||
upstream_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_forwards_tunnel_relay_to_attachment_owner() {
|
||||
let owner_hits = Arc::new(Mutex::new(0usize));
|
||||
let owner_hits_clone = Arc::clone(&owner_hits);
|
||||
let owner = Router::new().route(
|
||||
"/api/internal/tunnel/relay/node-123",
|
||||
post(move |headers: axum::http::HeaderMap, body: Body| {
|
||||
let owner_hits_inner = Arc::clone(&owner_hits_clone);
|
||||
async move {
|
||||
*owner_hits_inner.lock().expect("mutex should lock") += 1;
|
||||
assert_eq!(
|
||||
headers
|
||||
.get(aether_contracts::tunnel::TUNNEL_RELAY_FORWARDED_BY_HEADER)
|
||||
.and_then(|value| value.to_str().ok()),
|
||||
Some("gateway-a")
|
||||
);
|
||||
let body = axum::body::to_bytes(body, usize::MAX)
|
||||
.await
|
||||
.expect("body should read");
|
||||
let mut response = axum::http::Response::builder()
|
||||
.status(StatusCode::OK)
|
||||
.body(Body::from(body))
|
||||
.expect("response should build");
|
||||
response.headers_mut().insert(
|
||||
http::header::CONTENT_TYPE,
|
||||
HeaderValue::from_static("application/octet-stream"),
|
||||
);
|
||||
response
|
||||
}
|
||||
}),
|
||||
);
|
||||
|
||||
let (owner_url, owner_handle) = start_server(owner).await;
|
||||
let data_state = GatewayDataState::disabled().with_system_config_values_for_tests(vec![(
|
||||
"tunnel.attachments.node-123".to_string(),
|
||||
json!({
|
||||
"gateway_instance_id": "gateway-b",
|
||||
"relay_base_url": owner_url,
|
||||
"conn_count": 1,
|
||||
"observed_at_unix_secs": 4_102_444_800u64,
|
||||
}),
|
||||
)]);
|
||||
let gateway = build_router_with_state(
|
||||
AppState::new()
|
||||
.expect("gateway should build")
|
||||
.with_data_state_for_tests(data_state)
|
||||
.with_tunnel_identity_for_tests("gateway-a", Some("http://gateway-a.internal")),
|
||||
);
|
||||
let (gateway_url, gateway_handle) = start_server(gateway).await;
|
||||
|
||||
let response = reqwest::Client::new()
|
||||
.post(format!("{gateway_url}/api/internal/tunnel/relay/node-123"))
|
||||
.header(TRACE_ID_HEADER, "trace-owner-forward")
|
||||
.body("relay-envelope")
|
||||
.send()
|
||||
.await
|
||||
.expect("request should succeed");
|
||||
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
assert_eq!(
|
||||
response
|
||||
.headers()
|
||||
.get(TRACE_ID_HEADER)
|
||||
.and_then(|value| value.to_str().ok()),
|
||||
Some("trace-owner-forward")
|
||||
);
|
||||
assert_eq!(
|
||||
response.text().await.expect("body should read"),
|
||||
"relay-envelope"
|
||||
);
|
||||
assert_eq!(*owner_hits.lock().expect("mutex should lock"), 1);
|
||||
|
||||
gateway_handle.abort();
|
||||
owner_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_does_not_forward_tunnel_relay_twice() {
|
||||
let owner_hits = Arc::new(Mutex::new(0usize));
|
||||
let owner_hits_clone = Arc::clone(&owner_hits);
|
||||
let owner = Router::new().route(
|
||||
"/api/internal/tunnel/relay/node-123",
|
||||
post(move |_request: Request| {
|
||||
let owner_hits_inner = Arc::clone(&owner_hits_clone);
|
||||
async move {
|
||||
*owner_hits_inner.lock().expect("mutex should lock") += 1;
|
||||
(StatusCode::OK, Body::from("unexpected owner hit"))
|
||||
}
|
||||
}),
|
||||
);
|
||||
|
||||
let (owner_url, owner_handle) = start_server(owner).await;
|
||||
let data_state = GatewayDataState::disabled().with_system_config_values_for_tests(vec![(
|
||||
"tunnel.attachments.node-123".to_string(),
|
||||
json!({
|
||||
"gateway_instance_id": "gateway-b",
|
||||
"relay_base_url": owner_url,
|
||||
"conn_count": 1,
|
||||
"observed_at_unix_secs": 4_102_444_800u64,
|
||||
}),
|
||||
)]);
|
||||
let gateway = build_router_with_state(
|
||||
AppState::new()
|
||||
.expect("gateway should build")
|
||||
.with_data_state_for_tests(data_state)
|
||||
.with_tunnel_identity_for_tests("gateway-a", Some("http://gateway-a.internal")),
|
||||
);
|
||||
let (gateway_url, gateway_handle) = start_server(gateway).await;
|
||||
|
||||
let response = reqwest::Client::new()
|
||||
.post(format!("{gateway_url}/api/internal/tunnel/relay/node-123"))
|
||||
.header(
|
||||
aether_contracts::tunnel::TUNNEL_RELAY_FORWARDED_BY_HEADER,
|
||||
"gateway-z",
|
||||
)
|
||||
.send()
|
||||
.await
|
||||
.expect("request should succeed");
|
||||
|
||||
assert_eq!(response.status(), StatusCode::SERVICE_UNAVAILABLE);
|
||||
assert_eq!(*owner_hits.lock().expect("mutex should lock"), 0);
|
||||
|
||||
gateway_handle.abort();
|
||||
owner_handle.abort();
|
||||
}
|
||||
61
apps/aether-gateway/src/tests/control/mod.rs
Normal file
61
apps/aether-gateway/src/tests/control/mod.rs
Normal file
@@ -0,0 +1,61 @@
|
||||
use std::collections::BTreeMap;
|
||||
use std::sync::{Arc, Mutex};
|
||||
use std::time::{SystemTime, UNIX_EPOCH};
|
||||
|
||||
use aether_crypto::{
|
||||
decrypt_python_fernet_ciphertext, encrypt_python_fernet_plaintext, DEVELOPMENT_ENCRYPTION_KEY,
|
||||
};
|
||||
use aether_data::repository::auth::{
|
||||
InMemoryAuthApiKeySnapshotRepository, StoredAuthApiKeySnapshot,
|
||||
};
|
||||
use aether_data::repository::auth_modules::{
|
||||
InMemoryAuthModuleReadRepository, StoredLdapModuleConfig, StoredOAuthProviderModuleConfig,
|
||||
};
|
||||
use aether_data::repository::candidates::{
|
||||
InMemoryRequestCandidateRepository, RequestCandidateStatus, StoredRequestCandidate,
|
||||
};
|
||||
use aether_data::repository::global_models::{
|
||||
GlobalModelReadRepository, InMemoryGlobalModelReadRepository, StoredAdminGlobalModel,
|
||||
StoredAdminProviderModel, StoredProviderActiveGlobalModel, StoredProviderModelStats,
|
||||
StoredPublicGlobalModel,
|
||||
};
|
||||
use aether_data::repository::management_tokens::{
|
||||
InMemoryManagementTokenRepository, ManagementTokenReadRepository, StoredManagementToken,
|
||||
StoredManagementTokenUserSummary, StoredManagementTokenWithUser,
|
||||
};
|
||||
use aether_data::repository::oauth_providers::{
|
||||
InMemoryOAuthProviderRepository, OAuthProviderReadRepository, StoredOAuthProviderConfig,
|
||||
};
|
||||
use aether_data::repository::provider_catalog::{
|
||||
InMemoryProviderCatalogReadRepository, ProviderCatalogReadRepository,
|
||||
StoredProviderCatalogEndpoint, StoredProviderCatalogKey, StoredProviderCatalogProvider,
|
||||
};
|
||||
use aether_data::repository::proxy_nodes::{
|
||||
InMemoryProxyNodeRepository, ProxyNodeReadRepository, StoredProxyNode, StoredProxyNodeEvent,
|
||||
};
|
||||
use aether_data::repository::quota::{
|
||||
InMemoryProviderQuotaRepository, StoredProviderQuotaSnapshot,
|
||||
};
|
||||
use aether_data::repository::wallet::InMemoryWalletRepository;
|
||||
use axum::body::{to_bytes, Body, Bytes};
|
||||
use axum::response::Response;
|
||||
use axum::routing::{any, post};
|
||||
use axum::{extract::Request, Json, Router};
|
||||
use http::header::{HeaderName, HeaderValue};
|
||||
use http::{HeaderMap, StatusCode};
|
||||
use serde_json::json;
|
||||
|
||||
mod admin;
|
||||
mod helpers;
|
||||
mod internal;
|
||||
mod proxy;
|
||||
|
||||
use super::{
|
||||
build_router, build_router_with_execution_runtime_override, build_router_with_state,
|
||||
build_state_with_execution_runtime_override, start_server, wait_until, AppState,
|
||||
FrontdoorCorsConfig, FrontdoorUserRpmConfig, GatewayFallbackMetricKind, GatewayFallbackReason,
|
||||
Infallible, UsageRuntimeConfig, VideoTaskTruthSourceMode,
|
||||
};
|
||||
use crate::constants::*;
|
||||
use crate::data::GatewayDataState;
|
||||
use helpers::*;
|
||||
801
apps/aether-gateway/src/tests/control/proxy/local_denials.rs
Normal file
801
apps/aether-gateway/src/tests/control/proxy/local_denials.rs
Normal file
@@ -0,0 +1,801 @@
|
||||
use std::sync::{Arc, Mutex};
|
||||
|
||||
use axum::body::Body;
|
||||
use axum::routing::any;
|
||||
use axum::{extract::Request, Json, Router};
|
||||
use http::StatusCode;
|
||||
use serde_json::json;
|
||||
|
||||
use super::super::{
|
||||
build_router_with_state, hash_api_key, sample_currently_usable_auth_snapshot,
|
||||
sample_expired_auth_snapshot, sample_locked_auth_snapshot, start_server, AppState,
|
||||
GatewayDataState, InMemoryAuthApiKeySnapshotRepository, InMemoryWalletRepository,
|
||||
};
|
||||
use crate::constants::{
|
||||
CONTROL_ROUTE_CLASS_HEADER, EXECUTION_PATH_HEADER, EXECUTION_PATH_LOCAL_AUTH_DENIED,
|
||||
GATEWAY_HEADER, TRACE_ID_HEADER, TRUSTED_AUTH_ACCESS_ALLOWED_HEADER,
|
||||
TRUSTED_AUTH_API_KEY_ID_HEADER, TRUSTED_AUTH_BALANCE_HEADER, TRUSTED_AUTH_USER_ID_HEADER,
|
||||
};
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_locally_denies_explicit_trusted_balance_failure_without_hitting_control_or_upstream(
|
||||
) {
|
||||
let auth_context_hits = Arc::new(Mutex::new(0usize));
|
||||
let auth_context_hits_clone = Arc::clone(&auth_context_hits);
|
||||
let public_hits = Arc::new(Mutex::new(0usize));
|
||||
let public_hits_clone = Arc::clone(&public_hits);
|
||||
|
||||
let upstream = Router::new()
|
||||
.route(
|
||||
"/api/internal/gateway/auth-context",
|
||||
any(move |_request: Request| {
|
||||
let auth_context_hits_inner = Arc::clone(&auth_context_hits_clone);
|
||||
async move {
|
||||
*auth_context_hits_inner.lock().expect("mutex should lock") += 1;
|
||||
Json(json!({
|
||||
"auth_context": {
|
||||
"user_id": "user-from-control",
|
||||
"api_key_id": "key-from-control",
|
||||
"balance_remaining": 99.0,
|
||||
"access_allowed": true
|
||||
}
|
||||
}))
|
||||
}
|
||||
}),
|
||||
)
|
||||
.route(
|
||||
"/v1/chat/completions",
|
||||
any(move |_request: Request| {
|
||||
let public_hits_inner = Arc::clone(&public_hits_clone);
|
||||
async move {
|
||||
*public_hits_inner.lock().expect("mutex should lock") += 1;
|
||||
(StatusCode::OK, Body::from("unexpected upstream hit"))
|
||||
}
|
||||
}),
|
||||
);
|
||||
|
||||
let repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed(vec![(
|
||||
Some("hash-1".to_string()),
|
||||
sample_currently_usable_auth_snapshot("key-123", "user-123"),
|
||||
)]));
|
||||
let (upstream_url, upstream_handle) = start_server(upstream).await;
|
||||
let gateway = build_router_with_state(
|
||||
AppState::new()
|
||||
.expect("gateway state should build")
|
||||
.with_auth_api_key_data_reader_for_tests(repository),
|
||||
);
|
||||
let (gateway_url, gateway_handle) = start_server(gateway).await;
|
||||
|
||||
let response = reqwest::Client::new()
|
||||
.post(format!("{gateway_url}/v1/chat/completions"))
|
||||
.header(http::header::CONTENT_TYPE, "application/json")
|
||||
.header(TRACE_ID_HEADER, "trace-control-balance-denied-1")
|
||||
.header(crate::constants::GATEWAY_HEADER, "rust-phase3b")
|
||||
.header(TRUSTED_AUTH_USER_ID_HEADER, "user-123")
|
||||
.header(TRUSTED_AUTH_API_KEY_ID_HEADER, "key-123")
|
||||
.header(TRUSTED_AUTH_BALANCE_HEADER, "0")
|
||||
.header(TRUSTED_AUTH_ACCESS_ALLOWED_HEADER, "false")
|
||||
.body("{\"model\":\"gpt-5\",\"messages\":[]}")
|
||||
.send()
|
||||
.await
|
||||
.expect("request should succeed");
|
||||
|
||||
assert_eq!(response.status(), StatusCode::TOO_MANY_REQUESTS);
|
||||
assert_eq!(
|
||||
response
|
||||
.headers()
|
||||
.get(EXECUTION_PATH_HEADER)
|
||||
.and_then(|value| value.to_str().ok()),
|
||||
Some(EXECUTION_PATH_LOCAL_AUTH_DENIED)
|
||||
);
|
||||
assert_eq!(
|
||||
response
|
||||
.headers()
|
||||
.get(CONTROL_ROUTE_CLASS_HEADER)
|
||||
.and_then(|value| value.to_str().ok()),
|
||||
Some("ai_public")
|
||||
);
|
||||
let payload: serde_json::Value = response.json().await.expect("response json should parse");
|
||||
assert_eq!(payload["error"]["type"], "balance_exceeded");
|
||||
assert_eq!(payload["error"]["message"], "余额不足(剩余: $0.00)");
|
||||
assert_eq!(payload["error"]["details"]["balance_type"], "USD");
|
||||
assert_eq!(payload["error"]["details"]["remaining"], 0.0);
|
||||
|
||||
assert_eq!(*auth_context_hits.lock().expect("mutex should lock"), 0);
|
||||
assert_eq!(*public_hits.lock().expect("mutex should lock"), 0);
|
||||
|
||||
gateway_handle.abort();
|
||||
upstream_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_locally_denies_invalid_trusted_snapshot_without_hitting_control_or_upstream() {
|
||||
let auth_context_hits = Arc::new(Mutex::new(0usize));
|
||||
let auth_context_hits_clone = Arc::clone(&auth_context_hits);
|
||||
let public_hits = Arc::new(Mutex::new(0usize));
|
||||
let public_hits_clone = Arc::clone(&public_hits);
|
||||
|
||||
let upstream = Router::new()
|
||||
.route(
|
||||
"/api/internal/gateway/auth-context",
|
||||
any(move |_request: Request| {
|
||||
let auth_context_hits_inner = Arc::clone(&auth_context_hits_clone);
|
||||
async move {
|
||||
*auth_context_hits_inner.lock().expect("mutex should lock") += 1;
|
||||
Json(json!({
|
||||
"auth_context": {
|
||||
"user_id": "user-from-control",
|
||||
"api_key_id": "key-from-control",
|
||||
"access_allowed": true
|
||||
}
|
||||
}))
|
||||
}
|
||||
}),
|
||||
)
|
||||
.route(
|
||||
"/v1/chat/completions",
|
||||
any(move |_request: Request| {
|
||||
let public_hits_inner = Arc::clone(&public_hits_clone);
|
||||
async move {
|
||||
*public_hits_inner.lock().expect("mutex should lock") += 1;
|
||||
(StatusCode::OK, Body::from("unexpected upstream hit"))
|
||||
}
|
||||
}),
|
||||
);
|
||||
|
||||
let repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed(vec![(
|
||||
Some("hash-1".to_string()),
|
||||
sample_expired_auth_snapshot("key-123", "user-123"),
|
||||
)]));
|
||||
let (upstream_url, upstream_handle) = start_server(upstream).await;
|
||||
let gateway = build_router_with_state(
|
||||
AppState::new()
|
||||
.expect("gateway state should build")
|
||||
.with_auth_api_key_data_reader_for_tests(repository),
|
||||
);
|
||||
let (gateway_url, gateway_handle) = start_server(gateway).await;
|
||||
|
||||
let response = reqwest::Client::new()
|
||||
.post(format!("{gateway_url}/v1/chat/completions"))
|
||||
.header(http::header::CONTENT_TYPE, "application/json")
|
||||
.header(TRACE_ID_HEADER, "trace-control-invalid-trusted-1")
|
||||
.header(crate::constants::GATEWAY_HEADER, "rust-phase3b")
|
||||
.header(TRUSTED_AUTH_USER_ID_HEADER, "user-123")
|
||||
.header(TRUSTED_AUTH_API_KEY_ID_HEADER, "key-123")
|
||||
.body("{\"model\":\"gpt-5\",\"messages\":[]}")
|
||||
.send()
|
||||
.await
|
||||
.expect("request should succeed");
|
||||
|
||||
assert_eq!(response.status(), StatusCode::UNAUTHORIZED);
|
||||
assert_eq!(
|
||||
response
|
||||
.headers()
|
||||
.get(EXECUTION_PATH_HEADER)
|
||||
.and_then(|value| value.to_str().ok()),
|
||||
Some(EXECUTION_PATH_LOCAL_AUTH_DENIED)
|
||||
);
|
||||
let payload: serde_json::Value = response.json().await.expect("response json should parse");
|
||||
assert_eq!(payload["error"]["type"], "http_error");
|
||||
assert_eq!(payload["error"]["message"], "无效的API密钥");
|
||||
|
||||
assert_eq!(*auth_context_hits.lock().expect("mutex should lock"), 0);
|
||||
assert_eq!(*public_hits.lock().expect("mutex should lock"), 0);
|
||||
|
||||
gateway_handle.abort();
|
||||
upstream_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_locally_denies_missing_wallet_without_hitting_control_or_upstream() {
|
||||
let auth_context_hits = Arc::new(Mutex::new(0usize));
|
||||
let auth_context_hits_clone = Arc::clone(&auth_context_hits);
|
||||
let public_hits = Arc::new(Mutex::new(0usize));
|
||||
let public_hits_clone = Arc::clone(&public_hits);
|
||||
|
||||
let upstream = Router::new()
|
||||
.route(
|
||||
"/api/internal/gateway/auth-context",
|
||||
any(move |_request: Request| {
|
||||
let auth_context_hits_inner = Arc::clone(&auth_context_hits_clone);
|
||||
async move {
|
||||
*auth_context_hits_inner.lock().expect("mutex should lock") += 1;
|
||||
Json(json!({"auth_context": null}))
|
||||
}
|
||||
}),
|
||||
)
|
||||
.route(
|
||||
"/v1/chat/completions",
|
||||
any(move |_request: Request| {
|
||||
let public_hits_inner = Arc::clone(&public_hits_clone);
|
||||
async move {
|
||||
*public_hits_inner.lock().expect("mutex should lock") += 1;
|
||||
(StatusCode::OK, Body::from("unexpected upstream hit"))
|
||||
}
|
||||
}),
|
||||
);
|
||||
|
||||
let auth_repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed(vec![(
|
||||
Some("hash-1".to_string()),
|
||||
sample_currently_usable_auth_snapshot("key-123", "user-123"),
|
||||
)]));
|
||||
let wallet_repository = Arc::new(InMemoryWalletRepository::default());
|
||||
let (upstream_url, upstream_handle) = start_server(upstream).await;
|
||||
let data_state =
|
||||
GatewayDataState::with_auth_and_wallet_for_tests(auth_repository, wallet_repository);
|
||||
let gateway = build_router_with_state(
|
||||
AppState::new()
|
||||
.expect("gateway state should build")
|
||||
.with_data_state_for_tests(data_state),
|
||||
);
|
||||
let (gateway_url, gateway_handle) = start_server(gateway).await;
|
||||
|
||||
let response = reqwest::Client::new()
|
||||
.post(format!("{gateway_url}/v1/chat/completions"))
|
||||
.header(http::header::CONTENT_TYPE, "application/json")
|
||||
.header(TRACE_ID_HEADER, "trace-control-wallet-missing-1")
|
||||
.header(crate::constants::GATEWAY_HEADER, "rust-phase3b")
|
||||
.header(TRUSTED_AUTH_USER_ID_HEADER, "user-123")
|
||||
.header(TRUSTED_AUTH_API_KEY_ID_HEADER, "key-123")
|
||||
.body("{\"model\":\"gpt-5\",\"messages\":[]}")
|
||||
.send()
|
||||
.await
|
||||
.expect("request should succeed");
|
||||
|
||||
assert_eq!(response.status(), StatusCode::FORBIDDEN);
|
||||
assert_eq!(
|
||||
response
|
||||
.headers()
|
||||
.get(EXECUTION_PATH_HEADER)
|
||||
.and_then(|value| value.to_str().ok()),
|
||||
Some(EXECUTION_PATH_LOCAL_AUTH_DENIED)
|
||||
);
|
||||
let payload: serde_json::Value = response.json().await.expect("response json should parse");
|
||||
assert_eq!(payload["error"]["type"], "http_error");
|
||||
assert_eq!(payload["error"]["message"], "钱包不可用");
|
||||
|
||||
assert_eq!(*auth_context_hits.lock().expect("mutex should lock"), 0);
|
||||
assert_eq!(*public_hits.lock().expect("mutex should lock"), 0);
|
||||
|
||||
gateway_handle.abort();
|
||||
upstream_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_locally_denies_invalid_bearer_api_key_without_hitting_control_or_upstream() {
|
||||
let auth_context_hits = Arc::new(Mutex::new(0usize));
|
||||
let auth_context_hits_clone = Arc::clone(&auth_context_hits);
|
||||
let public_hits = Arc::new(Mutex::new(0usize));
|
||||
let public_hits_clone = Arc::clone(&public_hits);
|
||||
|
||||
let upstream = Router::new()
|
||||
.route(
|
||||
"/api/internal/gateway/auth-context",
|
||||
any(move |_request: Request| {
|
||||
let auth_context_hits_inner = Arc::clone(&auth_context_hits_clone);
|
||||
async move {
|
||||
*auth_context_hits_inner.lock().expect("mutex should lock") += 1;
|
||||
Json(json!({"auth_context": null}))
|
||||
}
|
||||
}),
|
||||
)
|
||||
.route(
|
||||
"/v1/chat/completions",
|
||||
any(move |_request: Request| {
|
||||
let public_hits_inner = Arc::clone(&public_hits_clone);
|
||||
async move {
|
||||
*public_hits_inner.lock().expect("mutex should lock") += 1;
|
||||
(StatusCode::OK, Body::from("unexpected upstream hit"))
|
||||
}
|
||||
}),
|
||||
);
|
||||
|
||||
let repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed(vec![(
|
||||
Some(hash_api_key("sk-other")),
|
||||
sample_currently_usable_auth_snapshot("key-123", "user-123"),
|
||||
)]));
|
||||
let (upstream_url, upstream_handle) = start_server(upstream).await;
|
||||
let gateway = build_router_with_state(
|
||||
AppState::new()
|
||||
.expect("gateway state should build")
|
||||
.with_auth_api_key_data_reader_for_tests(repository),
|
||||
);
|
||||
let (gateway_url, gateway_handle) = start_server(gateway).await;
|
||||
|
||||
let response = reqwest::Client::new()
|
||||
.post(format!("{gateway_url}/v1/chat/completions"))
|
||||
.header(http::header::CONTENT_TYPE, "application/json")
|
||||
.header(http::header::AUTHORIZATION, "Bearer sk-missing")
|
||||
.header(TRACE_ID_HEADER, "trace-control-invalid-bearer-1")
|
||||
.body("{\"model\":\"gpt-5\",\"messages\":[]}")
|
||||
.send()
|
||||
.await
|
||||
.expect("request should succeed");
|
||||
|
||||
assert_eq!(response.status(), StatusCode::UNAUTHORIZED);
|
||||
assert_eq!(
|
||||
response
|
||||
.headers()
|
||||
.get(EXECUTION_PATH_HEADER)
|
||||
.and_then(|value| value.to_str().ok()),
|
||||
Some(EXECUTION_PATH_LOCAL_AUTH_DENIED)
|
||||
);
|
||||
let payload: serde_json::Value = response.json().await.expect("response json should parse");
|
||||
assert_eq!(payload["error"]["type"], "http_error");
|
||||
assert_eq!(payload["error"]["message"], "无效的API密钥");
|
||||
assert_eq!(*auth_context_hits.lock().expect("mutex should lock"), 0);
|
||||
assert_eq!(*public_hits.lock().expect("mutex should lock"), 0);
|
||||
|
||||
gateway_handle.abort();
|
||||
upstream_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_locally_denies_admin_proxy_without_admin_principal_and_without_hitting_upstream() {
|
||||
let upstream_hits = Arc::new(Mutex::new(0usize));
|
||||
let upstream_hits_clone = Arc::clone(&upstream_hits);
|
||||
|
||||
let upstream = Router::new().route(
|
||||
"/api/admin/endpoints/health/summary",
|
||||
any(move |_request: Request| {
|
||||
let upstream_hits_inner = Arc::clone(&upstream_hits_clone);
|
||||
async move {
|
||||
*upstream_hits_inner.lock().expect("mutex should lock") += 1;
|
||||
(StatusCode::OK, Body::from("unexpected upstream hit"))
|
||||
}
|
||||
}),
|
||||
);
|
||||
|
||||
let (upstream_url, upstream_handle) = start_server(upstream).await;
|
||||
let gateway = build_router_with_state(AppState::new().expect("gateway should build"));
|
||||
let (gateway_url, gateway_handle) = start_server(gateway).await;
|
||||
|
||||
let response = reqwest::Client::new()
|
||||
.get(format!("{gateway_url}/api/admin/endpoints/health/summary"))
|
||||
.send()
|
||||
.await
|
||||
.expect("request should succeed");
|
||||
|
||||
assert_eq!(response.status(), StatusCode::UNAUTHORIZED);
|
||||
let payload: serde_json::Value = response.json().await.expect("response json should parse");
|
||||
assert_eq!(payload["detail"], "admin authentication required");
|
||||
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
|
||||
|
||||
gateway_handle.abort();
|
||||
upstream_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_rejects_unclassified_oauth_public_route_as_local_not_found_without_hitting_upstream(
|
||||
) {
|
||||
let upstream_hits = Arc::new(Mutex::new(0usize));
|
||||
let upstream_hits_clone = Arc::clone(&upstream_hits);
|
||||
|
||||
let upstream = Router::new().route(
|
||||
"/api/oauth/linuxdo/unsupported",
|
||||
any(move |_request: Request| {
|
||||
let upstream_hits_inner = Arc::clone(&upstream_hits_clone);
|
||||
async move {
|
||||
*upstream_hits_inner.lock().expect("mutex should lock") += 1;
|
||||
(StatusCode::OK, Body::from("unexpected upstream hit"))
|
||||
}
|
||||
}),
|
||||
);
|
||||
|
||||
let (upstream_url, upstream_handle) = start_server(upstream).await;
|
||||
let gateway = build_router_with_state(AppState::new().expect("gateway should build"));
|
||||
let (gateway_url, gateway_handle) = start_server(gateway).await;
|
||||
|
||||
let response = reqwest::Client::new()
|
||||
.get(format!("{gateway_url}/api/oauth/linuxdo/unsupported"))
|
||||
.send()
|
||||
.await
|
||||
.expect("request should succeed");
|
||||
|
||||
assert_eq!(response.status(), StatusCode::NOT_FOUND);
|
||||
let payload: serde_json::Value = response.json().await.expect("response json should parse");
|
||||
assert_eq!(payload["error"]["type"], "http_error");
|
||||
assert_eq!(payload["error"]["message"], "Route not found");
|
||||
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
|
||||
|
||||
gateway_handle.abort();
|
||||
upstream_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_locally_rejects_unclassified_admin_route_without_hitting_upstream() {
|
||||
let upstream_hits = Arc::new(Mutex::new(0usize));
|
||||
let upstream_hits_clone = Arc::clone(&upstream_hits);
|
||||
|
||||
let upstream = Router::new().route(
|
||||
"/api/admin/unsupported",
|
||||
any(move |_request: Request| {
|
||||
let upstream_hits_inner = Arc::clone(&upstream_hits_clone);
|
||||
async move {
|
||||
*upstream_hits_inner.lock().expect("mutex should lock") += 1;
|
||||
(StatusCode::OK, Body::from("unexpected upstream hit"))
|
||||
}
|
||||
}),
|
||||
);
|
||||
|
||||
let (upstream_url, upstream_handle) = start_server(upstream).await;
|
||||
let gateway = build_router_with_state(AppState::new().expect("gateway should build"));
|
||||
let (gateway_url, gateway_handle) = start_server(gateway).await;
|
||||
|
||||
let response = reqwest::Client::new()
|
||||
.get(format!("{gateway_url}/api/admin/unsupported"))
|
||||
.send()
|
||||
.await
|
||||
.expect("request should succeed");
|
||||
|
||||
assert_eq!(response.status(), StatusCode::NOT_IMPLEMENTED);
|
||||
let payload: serde_json::Value = response.json().await.expect("response json should parse");
|
||||
assert_eq!(
|
||||
payload["detail"],
|
||||
"admin proxy route not implemented in rust frontdoor"
|
||||
);
|
||||
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
|
||||
|
||||
gateway_handle.abort();
|
||||
upstream_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_locally_denies_disallowed_claude_api_format_without_hitting_control_or_upstream() {
|
||||
let auth_context_hits = Arc::new(Mutex::new(0usize));
|
||||
let auth_context_hits_clone = Arc::clone(&auth_context_hits);
|
||||
let public_hits = Arc::new(Mutex::new(0usize));
|
||||
let public_hits_clone = Arc::clone(&public_hits);
|
||||
|
||||
let upstream = Router::new()
|
||||
.route(
|
||||
"/api/internal/gateway/auth-context",
|
||||
any(move |_request: Request| {
|
||||
let auth_context_hits_inner = Arc::clone(&auth_context_hits_clone);
|
||||
async move {
|
||||
*auth_context_hits_inner.lock().expect("mutex should lock") += 1;
|
||||
Json(json!({"auth_context": null}))
|
||||
}
|
||||
}),
|
||||
)
|
||||
.route(
|
||||
"/v1/messages",
|
||||
any(move |_request: Request| {
|
||||
let public_hits_inner = Arc::clone(&public_hits_clone);
|
||||
async move {
|
||||
*public_hits_inner.lock().expect("mutex should lock") += 1;
|
||||
(StatusCode::OK, Body::from("unexpected upstream hit"))
|
||||
}
|
||||
}),
|
||||
);
|
||||
|
||||
let mut snapshot = sample_currently_usable_auth_snapshot("key-claude-123", "user-claude-123");
|
||||
snapshot.api_key_allowed_providers = Some(vec!["claude".to_string()]);
|
||||
snapshot.user_allowed_providers = Some(vec!["claude".to_string()]);
|
||||
snapshot.api_key_allowed_api_formats = Some(vec!["openai:chat".to_string()]);
|
||||
let repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed(vec![(
|
||||
Some(hash_api_key("sk-claude-auth-123")),
|
||||
snapshot,
|
||||
)]));
|
||||
let (upstream_url, upstream_handle) = start_server(upstream).await;
|
||||
let gateway = build_router_with_state(
|
||||
AppState::new()
|
||||
.expect("gateway state should build")
|
||||
.with_auth_api_key_data_reader_for_tests(repository),
|
||||
);
|
||||
let (gateway_url, gateway_handle) = start_server(gateway).await;
|
||||
|
||||
let response = reqwest::Client::new()
|
||||
.post(format!("{gateway_url}/v1/messages"))
|
||||
.header(http::header::CONTENT_TYPE, "application/json")
|
||||
.header("x-api-key", "sk-claude-auth-123")
|
||||
.header(TRACE_ID_HEADER, "trace-control-claude-format-1")
|
||||
.body("{\"model\":\"claude-3-7-sonnet\",\"messages\":[]}")
|
||||
.send()
|
||||
.await
|
||||
.expect("request should succeed");
|
||||
|
||||
assert_eq!(response.status(), StatusCode::FORBIDDEN);
|
||||
let payload: serde_json::Value = response.json().await.expect("response json should parse");
|
||||
assert_eq!(payload["error"]["type"], "http_error");
|
||||
assert_eq!(
|
||||
payload["error"]["message"],
|
||||
"当前密钥不允许访问 claude:chat 格式"
|
||||
);
|
||||
assert_eq!(*auth_context_hits.lock().expect("mutex should lock"), 0);
|
||||
assert_eq!(*public_hits.lock().expect("mutex should lock"), 0);
|
||||
|
||||
gateway_handle.abort();
|
||||
upstream_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_locally_denies_disallowed_provider_without_hitting_control_or_upstream() {
|
||||
let auth_context_hits = Arc::new(Mutex::new(0usize));
|
||||
let auth_context_hits_clone = Arc::clone(&auth_context_hits);
|
||||
let public_hits = Arc::new(Mutex::new(0usize));
|
||||
let public_hits_clone = Arc::clone(&public_hits);
|
||||
|
||||
let upstream = Router::new()
|
||||
.route(
|
||||
"/api/internal/gateway/auth-context",
|
||||
any(move |_request: Request| {
|
||||
let auth_context_hits_inner = Arc::clone(&auth_context_hits_clone);
|
||||
async move {
|
||||
*auth_context_hits_inner.lock().expect("mutex should lock") += 1;
|
||||
Json(json!({"auth_context": null}))
|
||||
}
|
||||
}),
|
||||
)
|
||||
.route(
|
||||
"/v1/messages",
|
||||
any(move |_request: Request| {
|
||||
let public_hits_inner = Arc::clone(&public_hits_clone);
|
||||
async move {
|
||||
*public_hits_inner.lock().expect("mutex should lock") += 1;
|
||||
(StatusCode::OK, Body::from("unexpected upstream hit"))
|
||||
}
|
||||
}),
|
||||
);
|
||||
|
||||
let mut snapshot =
|
||||
sample_currently_usable_auth_snapshot("key-claude-allow-123", "user-claude-allow-123");
|
||||
snapshot.api_key_allowed_api_formats = None;
|
||||
snapshot.user_allowed_api_formats = None;
|
||||
snapshot.api_key_allowed_providers = Some(vec!["openai".to_string()]);
|
||||
let repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed(vec![(
|
||||
Some(hash_api_key("sk-claude-provider-123")),
|
||||
snapshot,
|
||||
)]));
|
||||
let (upstream_url, upstream_handle) = start_server(upstream).await;
|
||||
let gateway = build_router_with_state(
|
||||
AppState::new()
|
||||
.expect("gateway state should build")
|
||||
.with_auth_api_key_data_reader_for_tests(repository),
|
||||
);
|
||||
let (gateway_url, gateway_handle) = start_server(gateway).await;
|
||||
|
||||
let response = reqwest::Client::new()
|
||||
.post(format!("{gateway_url}/v1/messages"))
|
||||
.header(http::header::CONTENT_TYPE, "application/json")
|
||||
.header("x-api-key", "sk-claude-provider-123")
|
||||
.header(TRACE_ID_HEADER, "trace-control-provider-guard-1")
|
||||
.body("{\"model\":\"claude-3-7-sonnet\",\"messages\":[]}")
|
||||
.send()
|
||||
.await
|
||||
.expect("request should succeed");
|
||||
|
||||
assert_eq!(response.status(), StatusCode::FORBIDDEN);
|
||||
let payload: serde_json::Value = response.json().await.expect("response json should parse");
|
||||
assert_eq!(payload["error"]["type"], "http_error");
|
||||
assert_eq!(
|
||||
payload["error"]["message"],
|
||||
"当前密钥不允许访问 claude 提供商"
|
||||
);
|
||||
assert_eq!(*auth_context_hits.lock().expect("mutex should lock"), 0);
|
||||
assert_eq!(*public_hits.lock().expect("mutex should lock"), 0);
|
||||
|
||||
gateway_handle.abort();
|
||||
upstream_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_locally_denies_disallowed_gemini_model_without_hitting_control_or_upstream() {
|
||||
let auth_context_hits = Arc::new(Mutex::new(0usize));
|
||||
let auth_context_hits_clone = Arc::clone(&auth_context_hits);
|
||||
let public_hits = Arc::new(Mutex::new(0usize));
|
||||
let public_hits_clone = Arc::clone(&public_hits);
|
||||
|
||||
let upstream = Router::new()
|
||||
.route(
|
||||
"/api/internal/gateway/auth-context",
|
||||
any(move |_request: Request| {
|
||||
let auth_context_hits_inner = Arc::clone(&auth_context_hits_clone);
|
||||
async move {
|
||||
*auth_context_hits_inner.lock().expect("mutex should lock") += 1;
|
||||
Json(json!({"auth_context": null}))
|
||||
}
|
||||
}),
|
||||
)
|
||||
.route(
|
||||
"/v1beta/models/{*path}",
|
||||
any(move |_request: Request| {
|
||||
let public_hits_inner = Arc::clone(&public_hits_clone);
|
||||
async move {
|
||||
*public_hits_inner.lock().expect("mutex should lock") += 1;
|
||||
(StatusCode::OK, Body::from("unexpected upstream hit"))
|
||||
}
|
||||
}),
|
||||
);
|
||||
|
||||
let mut snapshot = sample_currently_usable_auth_snapshot("key-gemini-123", "user-gemini-123");
|
||||
snapshot.api_key_allowed_providers = Some(vec!["gemini".to_string()]);
|
||||
snapshot.user_allowed_providers = Some(vec!["gemini".to_string()]);
|
||||
snapshot.api_key_allowed_api_formats = Some(vec!["gemini:chat".to_string()]);
|
||||
snapshot.api_key_allowed_models = Some(vec!["gemini-1.5-pro".to_string()]);
|
||||
let repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed(vec![(
|
||||
Some(hash_api_key("gemini-client-key-123")),
|
||||
snapshot,
|
||||
)]));
|
||||
let (upstream_url, upstream_handle) = start_server(upstream).await;
|
||||
let gateway = build_router_with_state(
|
||||
AppState::new()
|
||||
.expect("gateway state should build")
|
||||
.with_auth_api_key_data_reader_for_tests(repository),
|
||||
);
|
||||
let (gateway_url, gateway_handle) = start_server(gateway).await;
|
||||
|
||||
let response = reqwest::Client::new()
|
||||
.post(format!(
|
||||
"{gateway_url}/v1beta/models/gemini-2.5-pro:generateContent"
|
||||
))
|
||||
.header(http::header::CONTENT_TYPE, "application/json")
|
||||
.header("x-goog-api-key", "gemini-client-key-123")
|
||||
.header(TRACE_ID_HEADER, "trace-control-gemini-model-1")
|
||||
.body("{\"contents\":[]}")
|
||||
.send()
|
||||
.await
|
||||
.expect("request should succeed");
|
||||
|
||||
assert_eq!(response.status(), StatusCode::FORBIDDEN);
|
||||
let payload: serde_json::Value = response.json().await.expect("response json should parse");
|
||||
assert_eq!(payload["error"]["type"], "http_error");
|
||||
assert_eq!(
|
||||
payload["error"]["message"],
|
||||
"当前密钥不允许访问模型 gemini-2.5-pro"
|
||||
);
|
||||
assert_eq!(*auth_context_hits.lock().expect("mutex should lock"), 0);
|
||||
assert_eq!(*public_hits.lock().expect("mutex should lock"), 0);
|
||||
|
||||
gateway_handle.abort();
|
||||
upstream_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_locally_denies_locked_trusted_snapshot_without_hitting_control_or_upstream() {
|
||||
let auth_context_hits = Arc::new(Mutex::new(0usize));
|
||||
let auth_context_hits_clone = Arc::clone(&auth_context_hits);
|
||||
let public_hits = Arc::new(Mutex::new(0usize));
|
||||
let public_hits_clone = Arc::clone(&public_hits);
|
||||
|
||||
let upstream = Router::new()
|
||||
.route(
|
||||
"/api/internal/gateway/auth-context",
|
||||
any(move |_request: Request| {
|
||||
let auth_context_hits_inner = Arc::clone(&auth_context_hits_clone);
|
||||
async move {
|
||||
*auth_context_hits_inner.lock().expect("mutex should lock") += 1;
|
||||
Json(json!({"auth_context": null}))
|
||||
}
|
||||
}),
|
||||
)
|
||||
.route(
|
||||
"/v1/chat/completions",
|
||||
any(move |_request: Request| {
|
||||
let public_hits_inner = Arc::clone(&public_hits_clone);
|
||||
async move {
|
||||
*public_hits_inner.lock().expect("mutex should lock") += 1;
|
||||
(StatusCode::OK, Body::from("unexpected upstream hit"))
|
||||
}
|
||||
}),
|
||||
);
|
||||
|
||||
let repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed(vec![(
|
||||
Some("hash-1".to_string()),
|
||||
sample_locked_auth_snapshot("key-locked-123", "user-locked-123"),
|
||||
)]));
|
||||
let (upstream_url, upstream_handle) = start_server(upstream).await;
|
||||
let gateway = build_router_with_state(
|
||||
AppState::new()
|
||||
.expect("gateway state should build")
|
||||
.with_auth_api_key_data_reader_for_tests(repository),
|
||||
);
|
||||
let (gateway_url, gateway_handle) = start_server(gateway).await;
|
||||
|
||||
let response = reqwest::Client::new()
|
||||
.post(format!("{gateway_url}/v1/chat/completions"))
|
||||
.header(http::header::CONTENT_TYPE, "application/json")
|
||||
.header(TRACE_ID_HEADER, "trace-control-locked-trusted-1")
|
||||
.header(crate::constants::GATEWAY_HEADER, "rust-phase3b")
|
||||
.header(TRUSTED_AUTH_USER_ID_HEADER, "user-locked-123")
|
||||
.header(TRUSTED_AUTH_API_KEY_ID_HEADER, "key-locked-123")
|
||||
.body("{\"model\":\"gpt-5\",\"messages\":[]}")
|
||||
.send()
|
||||
.await
|
||||
.expect("request should succeed");
|
||||
|
||||
assert_eq!(response.status(), StatusCode::FORBIDDEN);
|
||||
assert_eq!(
|
||||
response
|
||||
.headers()
|
||||
.get(EXECUTION_PATH_HEADER)
|
||||
.and_then(|value| value.to_str().ok()),
|
||||
Some(EXECUTION_PATH_LOCAL_AUTH_DENIED)
|
||||
);
|
||||
let payload: serde_json::Value = response.json().await.expect("response json should parse");
|
||||
assert_eq!(payload["error"]["type"], "http_error");
|
||||
assert_eq!(
|
||||
payload["error"]["message"],
|
||||
"该密钥已被管理员锁定,请联系管理员"
|
||||
);
|
||||
assert_eq!(*auth_context_hits.lock().expect("mutex should lock"), 0);
|
||||
assert_eq!(*public_hits.lock().expect("mutex should lock"), 0);
|
||||
|
||||
gateway_handle.abort();
|
||||
upstream_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_locally_denies_disallowed_openai_model_without_hitting_control_or_upstream() {
|
||||
let auth_context_hits = Arc::new(Mutex::new(0usize));
|
||||
let auth_context_hits_clone = Arc::clone(&auth_context_hits);
|
||||
let public_hits = Arc::new(Mutex::new(0usize));
|
||||
let public_hits_clone = Arc::clone(&public_hits);
|
||||
|
||||
let upstream = Router::new()
|
||||
.route(
|
||||
"/api/internal/gateway/auth-context",
|
||||
any(move |_request: Request| {
|
||||
let auth_context_hits_inner = Arc::clone(&auth_context_hits_clone);
|
||||
async move {
|
||||
*auth_context_hits_inner.lock().expect("mutex should lock") += 1;
|
||||
Json(json!({"auth_context": null}))
|
||||
}
|
||||
}),
|
||||
)
|
||||
.route(
|
||||
"/v1/chat/completions",
|
||||
any(move |_request: Request| {
|
||||
let public_hits_inner = Arc::clone(&public_hits_clone);
|
||||
async move {
|
||||
*public_hits_inner.lock().expect("mutex should lock") += 1;
|
||||
(StatusCode::OK, Body::from("unexpected upstream hit"))
|
||||
}
|
||||
}),
|
||||
);
|
||||
|
||||
let mut snapshot =
|
||||
sample_currently_usable_auth_snapshot("key-openai-model-123", "user-openai-model-123");
|
||||
snapshot.api_key_allowed_models = Some(vec!["gpt-4.1".to_string()]);
|
||||
let repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed(vec![(
|
||||
Some(hash_api_key("sk-openai-model-guard-123")),
|
||||
snapshot,
|
||||
)]));
|
||||
let (upstream_url, upstream_handle) = start_server(upstream).await;
|
||||
let gateway = build_router_with_state(
|
||||
AppState::new()
|
||||
.expect("gateway state should build")
|
||||
.with_auth_api_key_data_reader_for_tests(repository),
|
||||
);
|
||||
let (gateway_url, gateway_handle) = start_server(gateway).await;
|
||||
|
||||
let response = reqwest::Client::new()
|
||||
.post(format!("{gateway_url}/v1/chat/completions"))
|
||||
.header(http::header::CONTENT_TYPE, "application/json")
|
||||
.header(
|
||||
http::header::AUTHORIZATION,
|
||||
"Bearer sk-openai-model-guard-123",
|
||||
)
|
||||
.header(TRACE_ID_HEADER, "trace-control-openai-model-guard-1")
|
||||
.body("{\"model\":\"gpt-5\",\"messages\":[]}")
|
||||
.send()
|
||||
.await
|
||||
.expect("request should succeed");
|
||||
|
||||
assert_eq!(response.status(), StatusCode::FORBIDDEN);
|
||||
assert_eq!(
|
||||
response
|
||||
.headers()
|
||||
.get(EXECUTION_PATH_HEADER)
|
||||
.and_then(|value| value.to_str().ok()),
|
||||
Some(EXECUTION_PATH_LOCAL_AUTH_DENIED)
|
||||
);
|
||||
let payload: serde_json::Value = response.json().await.expect("response json should parse");
|
||||
assert_eq!(payload["error"]["type"], "http_error");
|
||||
assert_eq!(payload["error"]["message"], "当前密钥不允许访问模型 gpt-5");
|
||||
assert_eq!(*auth_context_hits.lock().expect("mutex should lock"), 0);
|
||||
assert_eq!(*public_hits.lock().expect("mutex should lock"), 0);
|
||||
|
||||
gateway_handle.abort();
|
||||
upstream_handle.abort();
|
||||
}
|
||||
2
apps/aether-gateway/src/tests/control/proxy/mod.rs
Normal file
2
apps/aether-gateway/src/tests/control/proxy/mod.rs
Normal file
@@ -0,0 +1,2 @@
|
||||
mod local_denials;
|
||||
mod routing;
|
||||
152
apps/aether-gateway/src/tests/control/proxy/routing.rs
Normal file
152
apps/aether-gateway/src/tests/control/proxy/routing.rs
Normal file
@@ -0,0 +1,152 @@
|
||||
use std::sync::{Arc, Mutex};
|
||||
|
||||
use axum::body::Body;
|
||||
use axum::routing::any;
|
||||
use axum::{extract::Request, Router};
|
||||
use http::StatusCode;
|
||||
|
||||
use super::super::{
|
||||
build_router, build_router_with_state, hash_api_key, sample_currently_usable_auth_snapshot,
|
||||
start_server, AppState, InMemoryAuthApiKeySnapshotRepository,
|
||||
};
|
||||
use crate::constants::{
|
||||
CONTROL_ENDPOINT_SIGNATURE_HEADER, CONTROL_EXECUTION_RUNTIME_HEADER,
|
||||
CONTROL_ROUTE_CLASS_HEADER, CONTROL_ROUTE_FAMILY_HEADER, CONTROL_ROUTE_KIND_HEADER,
|
||||
GATEWAY_HEADER, TRACE_ID_HEADER, TRUSTED_ADMIN_SESSION_ID_HEADER, TRUSTED_ADMIN_USER_ID_HEADER,
|
||||
TRUSTED_ADMIN_USER_ROLE_HEADER,
|
||||
};
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_rejects_spoofed_admin_principal_headers_without_gateway_marker_locally() {
|
||||
let upstream_hits = Arc::new(Mutex::new(0usize));
|
||||
let upstream_hits_clone = Arc::clone(&upstream_hits);
|
||||
let upstream = Router::new().route(
|
||||
"/api/admin/endpoints/health/api-formats",
|
||||
any(move |_request: Request| {
|
||||
let upstream_hits_inner = Arc::clone(&upstream_hits_clone);
|
||||
async move {
|
||||
*upstream_hits_inner.lock().expect("mutex should lock") += 1;
|
||||
(StatusCode::OK, Body::from("unexpected upstream hit"))
|
||||
}
|
||||
}),
|
||||
);
|
||||
|
||||
let (upstream_url, upstream_handle) = start_server(upstream).await;
|
||||
let gateway = build_router().expect("gateway should build");
|
||||
let (gateway_url, gateway_handle) = start_server(gateway).await;
|
||||
|
||||
let response = reqwest::Client::new()
|
||||
.get(format!(
|
||||
"{gateway_url}/api/admin/endpoints/health/api-formats"
|
||||
))
|
||||
.header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
|
||||
.header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
|
||||
.header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
|
||||
.send()
|
||||
.await
|
||||
.expect("request should succeed");
|
||||
|
||||
assert_eq!(response.status(), StatusCode::UNAUTHORIZED);
|
||||
let payload: serde_json::Value = response.json().await.expect("body should parse");
|
||||
assert_eq!(payload["detail"], "admin authentication required");
|
||||
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
|
||||
|
||||
gateway_handle.abort();
|
||||
upstream_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_handles_models_support_routes_locally_with_public_support_control_headers() {
|
||||
let upstream_hits = Arc::new(Mutex::new(0usize));
|
||||
let upstream_hits_clone = Arc::clone(&upstream_hits);
|
||||
let upstream = Router::new().route(
|
||||
"/v1/models",
|
||||
any(move |_request: Request| {
|
||||
let upstream_hits_inner = Arc::clone(&upstream_hits_clone);
|
||||
async move {
|
||||
*upstream_hits_inner.lock().expect("mutex should lock") += 1;
|
||||
(
|
||||
StatusCode::OK,
|
||||
[(GATEWAY_HEADER, "fallback-probe")],
|
||||
Body::from("{\"object\":\"list\",\"data\":[]}"),
|
||||
)
|
||||
}
|
||||
}),
|
||||
);
|
||||
|
||||
let auth_repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed(vec![(
|
||||
Some(hash_api_key("sk-test-models")),
|
||||
sample_currently_usable_auth_snapshot("key-models-1", "user-models-1"),
|
||||
)]));
|
||||
|
||||
let (upstream_url, upstream_handle) = start_server(upstream).await;
|
||||
let gateway = build_router_with_state(
|
||||
AppState::new()
|
||||
.expect("gateway should build")
|
||||
.with_auth_api_key_data_reader_for_tests(auth_repository),
|
||||
);
|
||||
let (gateway_url, gateway_handle) = start_server(gateway).await;
|
||||
|
||||
let response = reqwest::Client::new()
|
||||
.get(format!("{gateway_url}/v1/models?limit=20"))
|
||||
.header(http::header::AUTHORIZATION, "Bearer sk-test-models")
|
||||
.header(TRACE_ID_HEADER, "trace-models-123")
|
||||
.send()
|
||||
.await
|
||||
.expect("request should succeed");
|
||||
|
||||
assert_eq!(response.status(), StatusCode::NOT_IMPLEMENTED);
|
||||
assert_eq!(
|
||||
response
|
||||
.headers()
|
||||
.get(CONTROL_ROUTE_CLASS_HEADER)
|
||||
.and_then(|value| value.to_str().ok()),
|
||||
Some("public_support")
|
||||
);
|
||||
assert_eq!(
|
||||
response
|
||||
.headers()
|
||||
.get(CONTROL_EXECUTION_RUNTIME_HEADER)
|
||||
.and_then(|value| value.to_str().ok()),
|
||||
Some("false")
|
||||
);
|
||||
assert_eq!(
|
||||
response
|
||||
.headers()
|
||||
.get(CONTROL_ROUTE_FAMILY_HEADER)
|
||||
.and_then(|value| value.to_str().ok()),
|
||||
Some("models")
|
||||
);
|
||||
assert_eq!(
|
||||
response
|
||||
.headers()
|
||||
.get(CONTROL_ROUTE_KIND_HEADER)
|
||||
.and_then(|value| value.to_str().ok()),
|
||||
Some("list")
|
||||
);
|
||||
assert_eq!(
|
||||
response
|
||||
.headers()
|
||||
.get(CONTROL_ENDPOINT_SIGNATURE_HEADER)
|
||||
.and_then(|value| value.to_str().ok()),
|
||||
Some("openai:chat")
|
||||
);
|
||||
assert_eq!(
|
||||
response
|
||||
.headers()
|
||||
.get(TRACE_ID_HEADER)
|
||||
.and_then(|value| value.to_str().ok()),
|
||||
Some("trace-models-123")
|
||||
);
|
||||
let payload: serde_json::Value = response.json().await.expect("body should parse");
|
||||
assert_eq!(
|
||||
payload["detail"],
|
||||
"public support route not implemented in rust frontdoor"
|
||||
);
|
||||
assert_eq!(payload["route_family"], "models");
|
||||
assert_eq!(payload["route_kind"], "list");
|
||||
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
|
||||
|
||||
gateway_handle.abort();
|
||||
upstream_handle.abort();
|
||||
}
|
||||
Reference in New Issue
Block a user