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

fix(stats): scope group usage by providers and add ungrouped view

Browse Source
This commit is contained in:
wangpengxiang
2026-09-21 12:49:48 +08:00
parent ba7c9f8b27
commit a9a7c64e5d
15 changed files with 622 additions and 48 deletions
Download Patch File Download Diff File Expand all files Collapse all files
apps/aether-gateway/src/tests/control/admin/stats.rs
+222 -1
View File
@@ -2031,7 +2031,12 @@ async fn gateway_aggregates_admin_stats_by_current_user_group_membership() {
assert_eq!(response.status(), StatusCode::OK);
let payload: serde_json::Value = response.json().await.expect("json body should parse");
assert_eq!(payload["attribution"], "current_membership");
assert_eq!(payload["total"], 1);
assert_eq!(payload["total"], 2);
let mut payload = payload;
payload["items"]
.as_array_mut()
.unwrap()
.retain(|item| item["id"] == group.id);
assert_eq!(payload["items"][0]["id"], group.id);
assert_eq!(payload["items"][0]["name"], "Engineering");
assert_eq!(payload["items"][0]["requests"], 2);
@@ -2067,6 +2072,222 @@ async fn gateway_aggregates_admin_stats_by_current_user_group_membership() {
upstream_handle.abort();
}
#[tokio::test]
async fn gateway_group_usage_intersects_members_and_providers_with_shared_overlap() {
let rows = vec![
sample_usage_row(
"g",
"g",
Some("user-1"),
None,
None,
"Gemini",
"model",
10,
2,
0.4,
0.4,
DAY_1_UNIX_SECS,
),
sample_usage_row(
"s",
"s",
Some("user-1"),
None,
None,
"Shared",
"model",
10,
2,
0.35,
0.35,
DAY_1_UNIX_SECS,
),
sample_usage_row(
"o",
"o",
Some("user-1"),
None,
None,
"Other",
"model",
10,
2,
1.5,
1.5,
DAY_1_UNIX_SECS,
),
sample_usage_row(
"x",
"x",
Some("user-2"),
None,
None,
"Gemini",
"model",
10,
2,
10.0,
10.0,
DAY_1_UNIX_SECS,
),
];
let users = InMemoryUserReadRepository::seed_auth_users([
sample_auth_user("user-1", "alice", "user", true),
sample_auth_user("user-2", "bob", "user", true),
]);
let export_users = users.list_export_users().await.unwrap();
let users = users.with_export_users(export_users);
let mut group_ids = Vec::new();
for (name, allowed, mode) in [
(
"Gemini group",
vec!["provider-gemini", "Shared", "provider-shared"],
"specific",
),
("Other group", vec!["Other", "shared-type"], "specific"),
("Denied", vec!["Gemini"], "deny_all"),
("Empty", vec![], "specific"),
("Inherited", vec![], "inherit"),
("Unrestricted", vec![], "unrestricted"),
] {
let group = users
.create_user_group(UpsertUserGroupRecord {
name: name.to_string(),
description: None,
priority: 0,
allowed_providers: Some(allowed.into_iter().map(str::to_string).collect()),
allowed_providers_mode: mode.to_string(),
allowed_api_formats: None,
allowed_api_formats_mode: "inherit".to_string(),
allowed_models: None,
allowed_models_mode: "inherit".to_string(),
rate_limit: None,
rate_limit_mode: "inherit".to_string(),
})
.await
.unwrap()
.unwrap();
users
.replace_user_group_members(&group.id, &["user-1".to_string()])
.await
.unwrap();
group_ids.push(group.id);
}
let mut shared = sample_provider("provider-shared", "Shared", 0);
shared.provider_type = "shared-type".to_string();
let providers = InMemoryProviderCatalogReadRepository::seed(
vec![
sample_provider("provider-gemini", "Gemini", 0),
shared,
sample_provider("provider-other", "Other", 0),
],
vec![],
vec![],
);
let data = GatewayDataState::with_usage_reader_for_tests(Arc::new(
InMemoryUsageReadRepository::seed(rows),
))
.with_user_reader(Arc::new(users))
.with_provider_catalog_reader(Arc::new(providers));
let gateway = build_router_with_state(AppState::new().unwrap().with_data_state_for_tests(data));
let (url, handle) = start_server(gateway).await;
let client = reqwest::Client::new();
let range = "start_date=2024-03-21&end_date=2024-03-21&tz_offset_minutes=0";
let paths = [
(
format!("usage/stats?{range}&user_group_id=__ungrouped__"),
Some(10.0),
),
(
format!("stats/time-series?{range}&granularity=day&user_group_id=__ungrouped__"),
Some(10.0),
),
(
format!("stats/leaderboard/users?{range}&metric=cost&user_group_id=__ungrouped__"),
Some(10.0),
),
(
format!("stats/leaderboard/user-groups?{range}&metric=cost"),
None,
),
(
format!("usage/stats?{range}&user_group_id={}", group_ids[0]),
Some(0.75),
),
(
format!(
"stats/time-series?{range}&granularity=day&user_group_id={}",
group_ids[1]
),
Some(1.85),
),
(
format!(
"stats/leaderboard/users?{range}&metric=cost&user_group_id={}",
group_ids[0]
),
Some(0.75),
),
(format!("usage/stats?{range}&user_id=user-1"), Some(2.25)),
(
format!(
"usage/stats?{range}&user_group_id={}&provider=Other",
group_ids[0]
),
Some(0.0),
),
];
for (path, expected) in paths {
let response = admin_request(client.get(format!("{url}/api/admin/{path}")))
.send()
.await
.unwrap();
assert_eq!(response.status(), StatusCode::OK, "{path}");
let body: serde_json::Value = response.json().await.unwrap();
if let Some(expected) = expected {
let value = if path.starts_with("stats/time-series") {
&body[0]["total_cost"]
} else if path.starts_with("stats/leaderboard") {
&body["items"][0]["cost"]
} else {
&body["total_cost"]
};
assert!(
(value.as_f64().unwrap() - expected).abs() < 1e-9,
"{path}: {body}"
);
} else {
let items = body["items"].as_array().unwrap();
assert_eq!(items.len(), 7);
let ungrouped = items
.iter()
.find(|item| item["id"] == "__ungrouped__")
.unwrap();
assert_eq!(ungrouped["cost"], 10.0);
assert_eq!(ungrouped["member_count"], 1);
assert_eq!(ungrouped["active_member_count"], 1);
for (id, cost, requests) in [
(&group_ids[0], 0.75, 2),
(&group_ids[1], 1.85, 2),
(&group_ids[2], 0.0, 0),
(&group_ids[3], 0.0, 0),
(&group_ids[4], 2.25, 3),
(&group_ids[5], 2.25, 3),
] {
let row = items
.iter()
.find(|item| item["id"].as_str() == Some(id.as_str()))
.unwrap();
assert!((row["cost"].as_f64().unwrap() - cost).abs() < 1e-9);
assert_eq!(row["requests"], requests);
}
}
}
handle.abort();
}
#[tokio::test]
async fn gateway_handles_admin_stats_leaderboard_users_locally_without_usage_reader() {
let (upstream_url, upstream_hits, upstream_handle) =