diff --git a/backend/crates/gateway-admin/tests/use_case/accounts.rs b/backend/crates/gateway-admin/tests/use_case/accounts.rs index d34a19ed8..f03a2e648 100644 --- a/backend/crates/gateway-admin/tests/use_case/accounts.rs +++ b/backend/crates/gateway-admin/tests/use_case/accounts.rs @@ -1,6 +1,7 @@ -use std::sync::{Arc, Mutex}; - -mod capacity; +use std::{ + collections::BTreeMap, + sync::{Arc, Mutex}, +}; use async_trait::async_trait; use chrono::{TimeDelta, Utc}; @@ -57,7 +58,8 @@ use gateway_admin::{ ProviderAdmin, ProviderAdminError, ProviderAdminErrorKind, ProviderAdminRegistry, }, store::{ - AccountStore, AdminStoreError, AdminStoreErrorKind, AdminStoreResult, SettingsStore, + AccountRuntimeStore, AccountStore, AdminStoreError, AdminStoreErrorKind, + AdminStoreResult, SettingsStore, }, }, }; @@ -3444,3 +3446,105 @@ async fn plugin_quota_refresh_uses_native_provider_without_changing_configuratio ); assert!(!recorded(&events).contains(&"snapshot.publish_committed")); } + +struct CapacityRuntime { + counts: Option>, + fail: bool, + reads: Mutex>>, +} + +#[async_trait] +impl AccountRuntimeStore for CapacityRuntime { + async fn active_rate_limits(&self) -> AdminStoreResult { + Ok(AccountRuntimeSnapshot::default()) + } + + async fn account_runtime(&self, ids: &[String]) -> AdminStoreResult { + self.reads.lock().unwrap().push(ids.to_vec()); + if self.fail { + return Err(store_unavailable()); + } + Ok(AccountRuntimeSnapshot { + in_flight: self.counts.clone(), + ..Default::default() + }) + } + + async fn active_freezes( + &self, + ) -> AdminStoreResult> { + Ok(BTreeMap::new()) + } + + async fn capacity_peaks(&self, _: &[String]) -> AdminStoreResult> { + Ok(BTreeMap::new()) + } + + async fn finish_freeze( + &self, + _: &str, + _: &gateway_admin::model::accounts::AccountFreeze, + _: Option>, + ) -> AdminStoreResult { + Ok(false) + } +} + +#[tokio::test] +async fn account_capacity_should_batch_page_ids_and_distinguish_idle_from_unavailable() { + for (counts, fail, expected) in [ + ( + Some(BTreeMap::from([("acct_test".to_owned(), 3)])), + false, + Some(3), + ), + (Some(BTreeMap::new()), false, Some(0)), + (None, false, None), + (None, true, None), + ] { + for empty in [false, true] { + let runtime = Arc::new(CapacityRuntime { + counts: counts.clone(), + fail, + reads: Mutex::default(), + }); + let mut account = account_record("openai"); + account.concurrency_limit = gateway_core::account::AccountConcurrencyLimit::new(5); + let store = FakeAccountStore::with_account(account, events()); + if empty { + store.set_accounts(Vec::new()); + } + let services = super::AdminHarness::new() + .accounts(store) + .account_runtime(runtime.clone()) + .settings(Arc::new(StaticSettingsStore)) + .provider(FakeProviderAdmin::new("openai", events())) + .build() + .await; + let page = services + .accounts() + .list(AccountListQuery { + page: 1, + page_size: gateway_admin::model::PageSize::new(20).unwrap(), + provider_kind: None, + group_filter: None, + search: None, + status: None, + sort: None, + }) + .await + .unwrap(); + if empty { + assert!(page.items.is_empty()); + assert!(runtime.reads.lock().unwrap().is_empty()); + } else { + assert_eq!( + *runtime.reads.lock().unwrap(), + [vec!["acct_test".to_owned()]] + ); + assert_eq!(page.items[0].capacity.used_slots, expected); + assert_eq!(page.items[0].capacity.total_slots, Some(5)); + } + } + } +} diff --git a/backend/crates/gateway-admin/tests/use_case/accounts/capacity.rs b/backend/crates/gateway-admin/tests/use_case/accounts/capacity.rs deleted file mode 100644 index bafb43568..000000000 --- a/backend/crates/gateway-admin/tests/use_case/accounts/capacity.rs +++ /dev/null @@ -1,107 +0,0 @@ -use std::collections::BTreeMap; - -use gateway_admin::ports::store::AccountRuntimeStore; - -use super::*; - -struct CapacityRuntime { - counts: Option>, - fail: bool, - reads: Mutex>>, -} - -#[async_trait] -impl AccountRuntimeStore for CapacityRuntime { - async fn active_rate_limits(&self) -> AdminStoreResult { - Ok(AccountRuntimeSnapshot::default()) - } - - async fn account_runtime(&self, ids: &[String]) -> AdminStoreResult { - self.reads.lock().unwrap().push(ids.to_vec()); - if self.fail { - return Err(store_unavailable()); - } - Ok(AccountRuntimeSnapshot { - in_flight: self.counts.clone(), - ..Default::default() - }) - } - - async fn active_freezes( - &self, - ) -> AdminStoreResult> { - Ok(BTreeMap::new()) - } - - async fn capacity_peaks(&self, _: &[String]) -> AdminStoreResult> { - Ok(BTreeMap::new()) - } - - async fn finish_freeze( - &self, - _: &str, - _: &gateway_admin::model::accounts::AccountFreeze, - _: Option>, - ) -> AdminStoreResult { - Ok(false) - } -} - -#[tokio::test] -async fn account_capacity_should_batch_page_ids_and_distinguish_idle_from_unavailable() { - for (counts, fail, expected) in [ - ( - Some(BTreeMap::from([("acct_test".to_owned(), 3)])), - false, - Some(3), - ), - (Some(BTreeMap::new()), false, Some(0)), - (None, false, None), - (None, true, None), - ] { - for empty in [false, true] { - let runtime = Arc::new(CapacityRuntime { - counts: counts.clone(), - fail, - reads: Mutex::default(), - }); - let mut account = account_record("openai"); - account.concurrency_limit = gateway_core::account::AccountConcurrencyLimit::new(5); - let store = FakeAccountStore::with_account(account, events()); - if empty { - store.set_accounts(Vec::new()); - } - let services = super::super::AdminHarness::new() - .accounts(store) - .account_runtime(runtime.clone()) - .settings(Arc::new(StaticSettingsStore)) - .provider(FakeProviderAdmin::new("openai", events())) - .build() - .await; - let page = services - .accounts() - .list(AccountListQuery { - page: 1, - page_size: gateway_admin::model::PageSize::new(20).unwrap(), - provider_kind: None, - group_filter: None, - search: None, - status: None, - sort: None, - }) - .await - .unwrap(); - if empty { - assert!(page.items.is_empty()); - assert!(runtime.reads.lock().unwrap().is_empty()); - } else { - assert_eq!( - *runtime.reads.lock().unwrap(), - [vec!["acct_test".to_owned()]] - ); - assert_eq!(page.items[0].capacity.used_slots, expected); - assert_eq!(page.items[0].capacity.total_slots, Some(5)); - } - } - } -}