mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-02 09:20:22 +08:00
Merge remote-tracking branch 'origin/pr-484'
This commit is contained in:
@@ -11,6 +11,9 @@ use crate::tests::{
|
||||
EXECUTION_PATH_EXECUTION_RUNTIME_STREAM, EXECUTION_PATH_EXECUTION_RUNTIME_SYNC,
|
||||
EXECUTION_PATH_HEADER,
|
||||
};
|
||||
use aether_data::repository::billing::InMemoryBillingReadRepository;
|
||||
use aether_data::repository::usage::InMemoryUsageReadRepository;
|
||||
use aether_data::repository::wallet::{InMemoryWalletRepository, StoredWalletSnapshot};
|
||||
use base64::Engine as _;
|
||||
|
||||
#[tokio::test]
|
||||
@@ -1154,6 +1157,118 @@ async fn gateway_handles_internal_gateway_decision_sync_locally_with_supplied_au
|
||||
upstream_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_internal_decision_sync_revalidates_supplied_auth_context_wallet() {
|
||||
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("proxied"))
|
||||
}
|
||||
}),
|
||||
);
|
||||
|
||||
let auth_repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed(vec![(
|
||||
None,
|
||||
unrestricted_models_snapshot("api-key-empty-wallet", "user-empty-wallet"),
|
||||
)]));
|
||||
let candidate_repository =
|
||||
Arc::new(InMemoryMinimalCandidateSelectionReadRepository::seed(vec![
|
||||
sample_models_candidate_row("provider-1", "openai", "openai:chat", "gpt-5", 10),
|
||||
]));
|
||||
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
|
||||
vec![sample_provider("provider-1", "openai", 10)],
|
||||
vec![sample_endpoint(
|
||||
"endpoint-provider-1",
|
||||
"provider-1",
|
||||
"openai:chat",
|
||||
"https://api.openai.example",
|
||||
)],
|
||||
vec![sample_key(
|
||||
"key-provider-1",
|
||||
"provider-1",
|
||||
"openai:chat",
|
||||
"sk-upstream-openai",
|
||||
)],
|
||||
));
|
||||
let request_candidate_repository = Arc::new(InMemoryRequestCandidateRepository::default());
|
||||
let usage_repository = Arc::new(InMemoryUsageReadRepository::default());
|
||||
let billing_repository = Arc::new(InMemoryBillingReadRepository::seed(Vec::new()));
|
||||
let wallet_repository = Arc::new(InMemoryWalletRepository::seed(vec![
|
||||
StoredWalletSnapshot::new(
|
||||
"wallet-empty".to_string(),
|
||||
Some("user-empty-wallet".to_string()),
|
||||
None,
|
||||
0.0,
|
||||
0.0,
|
||||
"finite".to_string(),
|
||||
"USD".to_string(),
|
||||
"active".to_string(),
|
||||
0.0,
|
||||
0.0,
|
||||
0.0,
|
||||
0.0,
|
||||
100,
|
||||
)
|
||||
.expect("wallet 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(
|
||||
crate::data::GatewayDataState::with_auth_candidate_selection_provider_catalog_request_candidates_usage_billing_and_wallet_for_tests(
|
||||
auth_repository,
|
||||
candidate_repository,
|
||||
provider_catalog_repository,
|
||||
request_candidate_repository,
|
||||
usage_repository,
|
||||
billing_repository,
|
||||
wallet_repository,
|
||||
DEVELOPMENT_ENCRYPTION_KEY,
|
||||
),
|
||||
),
|
||||
);
|
||||
let (gateway_url, gateway_handle) = start_server(gateway).await;
|
||||
|
||||
let response = reqwest::Client::new()
|
||||
.post(format!("{gateway_url}/api/internal/gateway/decision-sync"))
|
||||
.json(&json!({
|
||||
"trace_id": "trace-internal-decision-sync-empty-wallet",
|
||||
"method": "POST",
|
||||
"path": "/v1/chat/completions",
|
||||
"headers": {
|
||||
"content-type": "application/json",
|
||||
},
|
||||
"body_json": {
|
||||
"model": "gpt-5",
|
||||
"messages": [],
|
||||
},
|
||||
"auth_context": {
|
||||
"user_id": "user-empty-wallet",
|
||||
"api_key_id": "api-key-empty-wallet",
|
||||
"access_allowed": 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["action"], "fallback_plan");
|
||||
assert_eq!(payload["auth_context"], serde_json::Value::Null);
|
||||
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
|
||||
|
||||
gateway_handle.abort();
|
||||
upstream_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_returns_internal_gateway_decision_sync_fallback_with_resolved_auth_context() {
|
||||
let upstream_hits = Arc::new(Mutex::new(0usize));
|
||||
|
||||
@@ -4133,13 +4133,7 @@ async fn gateway_handles_wallet_balance_locally_without_proxying_upstream() {
|
||||
#[tokio::test]
|
||||
async fn gateway_handles_wallet_today_cost_locally_without_proxying_upstream() {
|
||||
let auth_now = Utc::now();
|
||||
let usage_now = chrono::DateTime::<chrono::Utc>::from_naive_utc_and_offset(
|
||||
auth_now
|
||||
.date_naive()
|
||||
.and_hms_opt(12, 0, 0)
|
||||
.expect("midday should be valid"),
|
||||
chrono::Utc,
|
||||
);
|
||||
let usage_now = auth_now - chrono::Duration::minutes(30);
|
||||
let user = sample_auth_user(auth_now);
|
||||
let access_token = build_test_auth_token(
|
||||
"access",
|
||||
@@ -4211,6 +4205,72 @@ async fn gateway_handles_wallet_today_cost_locally_without_proxying_upstream() {
|
||||
upstream_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_wallet_flow_today_entry_uses_live_settled_usage() {
|
||||
let auth_now = Utc::now();
|
||||
let user = sample_auth_user(auth_now);
|
||||
let access_token = build_test_auth_token(
|
||||
"access",
|
||||
serde_json::Map::from_iter([
|
||||
("user_id".to_string(), json!(user.id)),
|
||||
("role".to_string(), json!(user.role)),
|
||||
(
|
||||
"created_at".to_string(),
|
||||
json!(user.created_at.map(|value| value.to_rfc3339())),
|
||||
),
|
||||
("session_id".to_string(), json!("session-wallet-flow-today")),
|
||||
]),
|
||||
auth_now + chrono::Duration::hours(1),
|
||||
);
|
||||
let usage_repository = Arc::new(InMemoryUsageReadRepository::seed(vec![
|
||||
sample_user_usage_audit(
|
||||
"usage-wallet-flow-today",
|
||||
"req-wallet-flow-today",
|
||||
"user-auth-1",
|
||||
"gpt-4.1",
|
||||
"OpenAI",
|
||||
"completed",
|
||||
auth_now - chrono::Duration::minutes(5),
|
||||
),
|
||||
]));
|
||||
let (gateway_url, upstream_hits, gateway_handle, upstream_handle) =
|
||||
start_auth_gateway_with_usage_state(
|
||||
user,
|
||||
sample_auth_wallet("user-auth-1", auth_now),
|
||||
[sample_auth_session(
|
||||
"user-auth-1",
|
||||
"session-wallet-flow-today",
|
||||
"device-wallet-flow-today",
|
||||
"refresh-token-placeholder",
|
||||
auth_now,
|
||||
)],
|
||||
usage_repository,
|
||||
)
|
||||
.await;
|
||||
|
||||
let response = reqwest::Client::new()
|
||||
.get(format!("{gateway_url}/api/wallet/flow?limit=20&offset=0"))
|
||||
.header("authorization", format!("Bearer {access_token}"))
|
||||
.header("x-client-device-id", "device-wallet-flow-today")
|
||||
.header("user-agent", "AetherTest/1.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["today_entry"]["is_today"], true);
|
||||
assert_eq!(payload["today_entry"]["timezone"], "Asia/Shanghai");
|
||||
assert_eq!(payload["today_entry"]["total_requests"], 1);
|
||||
assert_eq!(payload["today_entry"]["input_tokens"], 120);
|
||||
assert_eq!(payload["today_entry"]["cache_read_tokens"], 15);
|
||||
assert_eq!(payload["today_entry"]["total_cost"], 1.25);
|
||||
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_wallet_today_cost_without_usage_reader() {
|
||||
let now = Utc::now();
|
||||
|
||||
Reference in New Issue
Block a user