use std::sync::Arc; use aether_data_contracts::repository::candidate_selection::MinimalCandidateSelectionReadRepository; use aether_data_contracts::repository::candidates::{ RequestCandidateReadRepository, RequestCandidateRepository, RequestCandidateWriteRepository, }; use aether_data_contracts::repository::provider_catalog::{ ProviderCatalogReadRepository, ProviderCatalogWriteRepository, }; use aether_data_contracts::repository::usage::{ UsageReadRepository, UsageRepository, UsageWriteRepository, }; use aether_data::repository::auth::AuthApiKeyReadRepository; use aether_data::repository::management_tokens::{ InMemoryManagementTokenRepository, ManagementTokenReadRepository, ManagementTokenWriteRepository, StoredManagementToken, StoredManagementTokenUserSummary, StoredManagementTokenWithUser, }; use aether_data::repository::proxy_nodes::{InMemoryProxyNodeRepository, StoredProxyNode}; use aether_data::repository::proxy_nodes::{ProxyNodeReadRepository, ProxyNodeWriteRepository}; use aether_data::repository::users::{ InMemoryUserReadRepository, StoredUserAuthRecord, UserReadRepository, }; use sha2::{Digest, Sha256}; use super::{GatewayDataConfig, GatewayDataState}; impl GatewayDataState { pub(crate) fn with_tunnel_management_auth_for_testkit( node_id: &str, tunnel_generation: &str, raw_token: &str, encryption_key: impl Into, ) -> Result { const TOKEN_ID: &str = "token-tunnel-harness"; const USER_ID: &str = "user-tunnel-harness"; let node = StoredProxyNode::new( node_id.to_string(), "tunnel harness node".to_string(), "127.0.0.1".to_string(), 0, false, "offline".to_string(), 30, 0, 0, 0, 0, 0, true, false, 0, )? .with_tunnel_generation(tunnel_generation.to_string()); let proxy_repository = Arc::new(InMemoryProxyNodeRepository::seed([node])); let user_summary = StoredManagementTokenUserSummary::new( USER_ID.to_string(), Some("tunnel-harness@example.com".to_string()), "tunnel_harness_admin".to_string(), "admin".to_string(), )?; let token = StoredManagementToken::new( TOKEN_ID.to_string(), USER_ID.to_string(), "tunnel harness token".to_string(), )? .with_permissions(Some(serde_json::json!(["admin:proxy_nodes:admin"]))); let token_hash = format!("{:x}", Sha256::digest(raw_token.as_bytes())); let token_repository = Arc::new(InMemoryManagementTokenRepository::seed_with_hashes( [StoredManagementTokenWithUser::new(token, user_summary)], [(token_hash, TOKEN_ID.to_string())], )); let token_reader: Arc = token_repository.clone(); let token_writer: Arc = token_repository; let user = StoredUserAuthRecord::new( USER_ID.to_string(), Some("tunnel-harness@example.com".to_string()), true, "tunnel_harness_admin".to_string(), None, "admin".to_string(), "local".to_string(), None, None, None, true, false, None, None, )?; let user_reader: Arc = Arc::new(InMemoryUserReadRepository::seed_auth_users([user])); let mut state = Self::with_proxy_node_repository_for_testkit(proxy_repository, encryption_key); state.management_token_reader = Some(token_reader); state.management_token_writer = Some(token_writer); state.user_reader = Some(user_reader); Ok(state) } pub(crate) fn with_proxy_node_repository_for_testkit( repository: Arc, encryption_key: impl Into, ) -> Self where T: ProxyNodeReadRepository + ProxyNodeWriteRepository + 'static, { let proxy_node_reader: Arc = repository.clone(); let proxy_node_writer: Arc = repository; let mut state = Self::disabled(); state.config = GatewayDataConfig::disabled().with_encryption_key(encryption_key); state.proxy_node_reader = Some(proxy_node_reader); state.proxy_node_writer = Some(proxy_node_writer); state } pub(crate) fn with_openai_chat_pressure_repositories_for_testkit( auth_api_key_repository: Arc, candidate_selection_repository: Arc, provider_catalog_repository: Arc, request_candidate_repository: Arc, usage_repository: Arc, encryption_key: impl Into, ) -> Self where T: RequestCandidateRepository + 'static, U: ProviderCatalogReadRepository + ProviderCatalogWriteRepository + 'static, V: UsageRepository + 'static, { let request_candidate_reader: Arc = request_candidate_repository.clone(); let request_candidate_writer: Arc = request_candidate_repository; let provider_catalog_reader: Arc = provider_catalog_repository.clone(); let provider_catalog_writer: Arc = provider_catalog_repository; let usage_reader: Arc = usage_repository.clone(); let usage_writer: Arc = usage_repository; Self { config: GatewayDataConfig::disabled().with_encryption_key(encryption_key), backends: None, auth_api_key_reader: Some(auth_api_key_repository), auth_api_key_writer: None, auth_module_reader: None, auth_module_writer: None, announcement_reader: None, announcement_writer: None, management_token_reader: None, management_token_writer: None, oauth_provider_reader: None, oauth_provider_writer: None, proxy_node_reader: None, proxy_node_writer: None, billing_reader: None, background_task_reader: None, background_task_writer: None, gemini_file_mapping_reader: None, gemini_file_mapping_writer: None, global_model_reader: None, global_model_writer: None, minimal_candidate_selection_reader: Some(candidate_selection_repository), request_candidate_reader: Some(request_candidate_reader), request_candidate_writer: Some(request_candidate_writer), provider_catalog_reader: Some(provider_catalog_reader), provider_catalog_writer: Some(provider_catalog_writer), pool_score_reader: None, pool_score_writer: None, provider_quota_reader: None, provider_quota_writer: None, routing_group_reader: None, routing_group_writer: None, usage_reader: Some(usage_reader), usage_writer: Some(usage_writer), user_reader: None, user_preferences: None, usage_worker_queue: None, video_task_reader: None, video_task_writer: None, wallet_reader: None, wallet_writer: None, settlement_writer: None, system_config_values: None, system_config_value_cache: Default::default(), billing_model_context_cache: Default::default(), } } }