i/Aether
1
0
Fork 0
mirror of https://github.com/fawney19/Aether.git synced 2026-09-10 05:00:19 +08:00
Code Issues Packages Projects Releases Wiki Activity

Assert usage server timing over HTTP

Browse Source
This commit is contained in:
Novick Yuan
2026-05-28 15:01:59 +08:00
parent 35cf44b38e
commit b108dc5ea6
3 changed files with 76 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files
apps/aether-gateway/src/tests/control/admin/usage.rs
+19
View File
@@ -33,6 +33,7 @@ use crate::constants::{
};
use crate::control::resolve_public_request_context;
use crate::data::GatewayDataState;
use crate::tests::{assert_usage_server_now_header_between, unix_epoch_millis_for_tests};
const ADMIN_USAGE_DATA_UNAVAILABLE_DETAIL: &str = "Admin usage data unavailable";
const DAY_1_UNIX_SECS: i64 = 1_711_000_000;
@@ -1014,6 +1015,7 @@ async fn gateway_handles_admin_usage_active_locally_with_trusted_admin_principal
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let client_send_unix_ms = unix_epoch_millis_for_tests();
let response =
admin_request(reqwest::Client::new().get(format!(
"{gateway_url}/api/admin/usage/active?start_date=2024-03-21&end_date=2024-03-22&tz_offset_minutes=0"
@@ -1021,9 +1023,17 @@ async fn gateway_handles_admin_usage_active_locally_with_trusted_admin_principal
.send()
.await
.expect("request should succeed");
let client_receive_unix_ms = unix_epoch_millis_for_tests();
assert_eq!(response.status(), StatusCode::OK);
let response_headers = response.headers().clone();
assert_usage_server_now_header_between(
&response_headers,
client_send_unix_ms,
client_receive_unix_ms,
);
let payload: serde_json::Value = response.json().await.expect("json body should parse");
assert!(payload.get("server_now_unix_ms").is_none());
assert_eq!(payload["requests"].as_array().expect("array").len(), 1);
assert_eq!(payload["requests"][0]["id"], "usage-pending");
assert_eq!(payload["requests"][0]["effective_input_tokens"], 5);
@@ -1233,15 +1243,24 @@ async fn gateway_handles_admin_usage_records_locally_with_trusted_admin_principa
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let client_send_unix_ms = unix_epoch_millis_for_tests();
let response = admin_request(reqwest::Client::new().get(format!(
"{gateway_url}/api/admin/usage/records?start_date=2024-03-21&end_date=2024-03-22&tz_offset_minutes=0&status=failed&provider=Anthropic&limit=10&offset=0"
)))
.send()
.await
.expect("request should succeed");
let client_receive_unix_ms = unix_epoch_millis_for_tests();
assert_eq!(response.status(), StatusCode::OK);
let response_headers = response.headers().clone();
assert_usage_server_now_header_between(
&response_headers,
client_send_unix_ms,
client_receive_unix_ms,
);
let payload: serde_json::Value = response.json().await.expect("json body should parse");
assert!(payload.get("server_now_unix_ms").is_none());
assert_eq!(payload["total"], 1);
assert_eq!(payload["records"][0]["id"], "usage-b");
assert_eq!(payload["records"][0]["username"], "bob");
apps/aether-gateway/src/tests/frontdoor/public_support.rs
+22 -4
View File
@@ -11,10 +11,10 @@ use super::{
};
use crate::data::GatewayDataState;
use crate::tests::{
any, build_router, build_router_with_state, json, start_server, to_bytes, AppState, Arc, Body,
Json, Mutex, Request, Router, StatusCode, CONTROL_ROUTE_FAMILY_HEADER,
CONTROL_ROUTE_KIND_HEADER, TRUSTED_ADMIN_SESSION_ID_HEADER, TRUSTED_ADMIN_USER_ID_HEADER,
TRUSTED_ADMIN_USER_ROLE_HEADER,
any, assert_usage_server_now_header_between, build_router, build_router_with_state, json,
start_server, to_bytes, unix_epoch_millis_for_tests, AppState, Arc, Body, Json, Mutex, Request,
Router, StatusCode, CONTROL_ROUTE_FAMILY_HEADER, CONTROL_ROUTE_KIND_HEADER,
TRUSTED_ADMIN_SESSION_ID_HEADER, TRUSTED_ADMIN_USER_ID_HEADER, TRUSTED_ADMIN_USER_ROLE_HEADER,
};
use aether_crypto::{encrypt_python_fernet_plaintext, DEVELOPMENT_ENCRYPTION_KEY};
use aether_data::repository::announcements::{AnnouncementListQuery, AnnouncementReadRepository};
@@ -5345,6 +5345,7 @@ async fn gateway_handles_users_me_usage_locally_without_proxying_upstream() {
})
.await;
let client_send_unix_ms = unix_epoch_millis_for_tests();
let response = reqwest::Client::new()
.get(format!(
"{gateway_url}/api/users/me/usage?limit=10&offset=0&search=renamed-key"
@@ -5355,9 +5356,17 @@ async fn gateway_handles_users_me_usage_locally_without_proxying_upstream() {
.send()
.await
.expect("request should succeed");
let client_receive_unix_ms = unix_epoch_millis_for_tests();
assert_eq!(response.status(), StatusCode::OK);
let response_headers = response.headers().clone();
assert_usage_server_now_header_between(
&response_headers,
client_send_unix_ms,
client_receive_unix_ms,
);
let payload: serde_json::Value = response.json().await.expect("json body should parse");
assert!(payload.get("server_now_unix_ms").is_none());
assert_eq!(payload["total_requests"], 2);
assert_eq!(payload["total_input_tokens"], 240);
assert_eq!(payload["pagination"]["total"], 3);
@@ -5667,6 +5676,7 @@ async fn gateway_handles_users_me_usage_active_locally_without_proxying_upstream
)
.await;
let client_send_unix_ms = unix_epoch_millis_for_tests();
let response = reqwest::Client::new()
.get(format!("{gateway_url}/api/users/me/usage/active"))
.header("authorization", format!("Bearer {access_token}"))
@@ -5675,9 +5685,17 @@ async fn gateway_handles_users_me_usage_active_locally_without_proxying_upstream
.send()
.await
.expect("request should succeed");
let client_receive_unix_ms = unix_epoch_millis_for_tests();
assert_eq!(response.status(), StatusCode::OK);
let response_headers = response.headers().clone();
assert_usage_server_now_header_between(
&response_headers,
client_send_unix_ms,
client_receive_unix_ms,
);
let payload: serde_json::Value = response.json().await.expect("json body should parse");
assert!(payload.get("server_now_unix_ms").is_none());
let requests = payload["requests"].as_array().expect("requests array");
assert_eq!(requests.len(), 2);
assert_eq!(requests[0]["status"], "streaming");
apps/aether-gateway/src/tests/mod.rs
+35
View File
@@ -1,6 +1,9 @@
use std::time::{SystemTime, UNIX_EPOCH};
pub(super) use std::convert::Infallible;
pub(super) use std::sync::{Arc, Mutex};
use aether_contracts::USAGE_SERVER_NOW_UNIX_MS_HEADER;
pub(super) use axum::body::{to_bytes, Body, Bytes};
pub(super) use axum::response::Response;
pub(super) use axum::routing::any;
@@ -29,6 +32,38 @@ pub(super) use super::router::{attach_static_frontend, build_router, build_route
pub(super) use super::state::{AppState, FrontdoorCorsConfig};
pub(super) use super::usage::UsageRuntimeConfig;
const SERVER_NOW_HEADER_TEST_TOLERANCE_MS: u64 = 1_000;
pub(super) fn unix_epoch_millis_for_tests() -> u64 {
SystemTime::now()
.duration_since(UNIX_EPOCH)
.expect("current time should be after epoch")
.as_millis()
.try_into()
.expect("current epoch millis should fit in u64")
}
pub(super) fn assert_usage_server_now_header_between(
headers: &reqwest::header::HeaderMap,
lower_bound_unix_ms: u64,
upper_bound_unix_ms: u64,
) {
let server_now_unix_ms = headers
.get(USAGE_SERVER_NOW_UNIX_MS_HEADER)
.expect("usage response should include server timing header")
.to_str()
.expect("server timing header should be valid ASCII")
.parse::<u64>()
.expect("server timing header should be epoch millis");
assert!(
server_now_unix_ms >= lower_bound_unix_ms.saturating_sub(SERVER_NOW_HEADER_TEST_TOLERANCE_MS)
&& server_now_unix_ms
<= upper_bound_unix_ms.saturating_add(SERVER_NOW_HEADER_TEST_TOLERANCE_MS),
"server timing header {server_now_unix_ms} should be near request window {lower_bound_unix_ms}..={upper_bound_unix_ms}"
);
}
pub(super) async fn start_server(app: Router) -> (String, tokio::task::JoinHandle<()>) {
let listener = crate::test_support::bind_loopback_listener()
.await