use std::collections::BTreeMap; use aether_data_contracts::repository::provider_catalog::StoredProviderCatalogProvider; use aether_data_contracts::repository::quota::StoredProviderQuotaSnapshot; use aether_wallet::{ProviderBillingType, ProviderQuotaSnapshot}; pub fn should_skip_provider_quota(quota: &StoredProviderQuotaSnapshot, now_unix_secs: u64) -> bool { let snapshot = ProviderQuotaSnapshot { provider_id: quota.provider_id.clone(), billing_type: ProviderBillingType::parse("a.billing_type), monthly_quota_usd: quota.monthly_quota_usd, monthly_used_usd: quota.monthly_used_usd, quota_reset_day: quota.quota_reset_day, quota_last_reset_at_unix_secs: quota.quota_last_reset_at_unix_secs, quota_expires_at_unix_secs: quota.quota_expires_at_unix_secs, is_active: quota.is_active, }; if !snapshot.is_active || snapshot.is_expired(now_unix_secs) { return true; } match snapshot.billing_type { ProviderBillingType::MonthlyQuota | ProviderBillingType::FreeTier => snapshot .remaining_quota_usd() .is_some_and(|remaining| remaining <= 0.0), ProviderBillingType::PayAsYouGo | ProviderBillingType::Unknown => false, } } pub fn build_provider_concurrent_limit_map( providers: Vec, ) -> BTreeMap { providers .into_iter() .filter_map(|provider| { provider .concurrent_limit .and_then(|limit| usize::try_from(limit).ok()) .filter(|limit| *limit > 0) .map(|limit| (provider.id, limit)) }) .collect() } #[cfg(test)] mod tests { use super::{build_provider_concurrent_limit_map, should_skip_provider_quota}; use aether_data_contracts::repository::provider_catalog::StoredProviderCatalogProvider; use aether_data_contracts::repository::quota::StoredProviderQuotaSnapshot; fn sample_provider(id: &str, concurrent_limit: Option) -> StoredProviderCatalogProvider { StoredProviderCatalogProvider::new( id.to_string(), format!("provider-{id}"), Some("https://example.com".to_string()), "custom".to_string(), ) .expect("provider should build") .with_transport_fields( true, false, false, concurrent_limit, None, None, None, None, None, ) } #[test] fn skips_inactive_or_exhausted_monthly_quota_provider() { let inactive = StoredProviderQuotaSnapshot::new( "provider-1".to_string(), "monthly_quota".to_string(), Some(10.0), 1.0, Some(30), Some(1_000), None, false, ) .expect("quota should build"); assert!(should_skip_provider_quota(&inactive, 2_000)); let exhausted = StoredProviderQuotaSnapshot::new( "provider-1".to_string(), "monthly_quota".to_string(), Some(10.0), 10.0, Some(30), Some(1_000), None, true, ) .expect("quota should build"); assert!(should_skip_provider_quota(&exhausted, 2_000)); let payg = StoredProviderQuotaSnapshot::new( "provider-1".to_string(), "pay_as_you_go".to_string(), None, 10.0, None, None, None, true, ) .expect("quota should build"); assert!(!should_skip_provider_quota(&payg, 2_000)); } #[test] fn builds_provider_concurrent_limit_map_for_positive_limits_only() { let limits = build_provider_concurrent_limit_map(vec![ sample_provider("provider-a", Some(10)), sample_provider("provider-b", Some(0)), sample_provider("provider-c", None), ]); assert_eq!(limits.get("provider-a"), Some(&10)); assert!(!limits.contains_key("provider-b")); assert!(!limits.contains_key("provider-c")); } }