[model-gateway] minor code clean up (#15578)

sgl-model-gateway/tests/common/mod.rs

@@ -26,6 +26,8 @@ use sgl_model_gateway::{
     middleware::TokenBucket,
     policies::PolicyRegistry,
     protocols::common::{Function, Tool},
+    reasoning_parser::ParserFactory as ReasoningParserFactory,
+    tool_parser::ParserFactory as ToolParserFactory,
 };
 
 /// Helper function to create AppContext for tests
@@ -159,6 +161,141 @@ pub async fn create_test_context(config: RouterConfig) -> Arc<AppContext> {
     app_context
 }
 
+/// Helper function to create AppContext for tests with parser factories initialized
+pub async fn create_test_context_with_parsers(config: RouterConfig) -> Arc<AppContext> {
+    let client = reqwest::Client::new();
+
+    // Initialize rate limiter
+    let rate_limiter = match config.max_concurrent_requests {
+        n if n <= 0 => None,
+        n => {
+            let rate_limit_tokens = config
+                .rate_limit_tokens_per_second
+                .filter(|&t| t > 0)
+                .unwrap_or(n);
+            Some(Arc::new(TokenBucket::new(
+                n as usize,
+                rate_limit_tokens as usize,
+            )))
+        }
+    };
+
+    // Initialize registries
+    let worker_registry = Arc::new(WorkerRegistry::new());
+    let policy_registry = Arc::new(PolicyRegistry::new(config.policy.clone()));
+
+    // Initialize storage backends (Memory for tests)
+    let response_storage = Arc::new(MemoryResponseStorage::new());
+    let conversation_storage = Arc::new(MemoryConversationStorage::new());
+    let conversation_item_storage = Arc::new(MemoryConversationItemStorage::new());
+
+    // Initialize load monitor
+    let load_monitor = Some(Arc::new(LoadMonitor::new(
+        worker_registry.clone(),
+        policy_registry.clone(),
+        client.clone(),
+        config.worker_startup_check_interval_secs,
+    )));
+
+    // Create empty OnceLock for worker job queue, workflow engine, and mcp manager
+    let worker_job_queue = Arc::new(OnceLock::new());
+    let workflow_engine = Arc::new(OnceLock::new());
+    let mcp_manager_lock = Arc::new(OnceLock::new());
+
+    // Initialize parser factories
+    let reasoning_parser_factory = Some(ReasoningParserFactory::new());
+    let tool_parser_factory = Some(ToolParserFactory::new());
+
+    let app_context = Arc::new(
+        AppContext::builder()
+            .router_config(config.clone())
+            .client(client)
+            .rate_limiter(rate_limiter)
+            .tokenizer(None) // tokenizer
+            .reasoning_parser_factory(reasoning_parser_factory)
+            .tool_parser_factory(tool_parser_factory)
+            .worker_registry(worker_registry)
+            .policy_registry(policy_registry)
+            .response_storage(response_storage)
+            .conversation_storage(conversation_storage)
+            .conversation_item_storage(conversation_item_storage)
+            .load_monitor(load_monitor)
+            .worker_job_queue(worker_job_queue)
+            .workflow_engine(workflow_engine)
+            .mcp_manager(mcp_manager_lock)
+            .build()
+            .unwrap(),
+    );
+
+    // Initialize JobQueue after AppContext is created
+    let weak_context = Arc::downgrade(&app_context);
+    let job_queue = sgl_model_gateway::core::JobQueue::new(
+        sgl_model_gateway::core::JobQueueConfig::default(),
+        weak_context,
+    );
+    app_context
+        .worker_job_queue
+        .set(job_queue)
+        .expect("JobQueue should only be initialized once");
+
+    // Initialize WorkflowEngine and register workflows
+    use sgl_model_gateway::{
+        core::steps::{create_worker_registration_workflow, create_worker_removal_workflow},
+        workflow::WorkflowEngine,
+    };
+    let engine = Arc::new(WorkflowEngine::new());
+    engine
+        .register_workflow(create_worker_registration_workflow(&config))
+        .expect("worker_registration workflow should be valid");
+    engine
+        .register_workflow(create_worker_removal_workflow())
+        .expect("worker_removal workflow should be valid");
+    app_context
+        .workflow_engine
+        .set(engine)
+        .expect("WorkflowEngine should only be initialized once");
+
+    // Register external workers for OpenAI mode
+    if let RoutingMode::OpenAI { worker_urls, .. } = &config.mode {
+        for url in worker_urls {
+            // Create a worker that supports common test models
+            let models = vec![
+                ModelCard::new("mock-model"),
+                ModelCard::new("gpt-4"),
+                ModelCard::new("gpt-3.5-turbo"),
+            ];
+            let worker: Arc<dyn Worker> = Arc::new(
+                BasicWorkerBuilder::new(url)
+                    .worker_type(WorkerType::Regular)
+                    .runtime_type(RuntimeType::External)
+                    .models(models)
+                    .build(),
+            );
+            app_context.worker_registry.register(worker);
+        }
+    }
+
+    // Initialize MCP manager with empty config
+    use sgl_model_gateway::mcp::{McpConfig, McpManager};
+    let empty_config = McpConfig {
+        servers: vec![],
+        pool: Default::default(),
+        proxy: None,
+        warmup: vec![],
+        inventory: Default::default(),
+    };
+    let mcp_manager = McpManager::with_defaults(empty_config)
+        .await
+        .expect("Failed to create MCP manager");
+    app_context
+        .mcp_manager
+        .set(Arc::new(mcp_manager))
+        .ok()
+        .expect("McpManager should only be initialized once");
+
+    app_context
+}
+
 /// Helper function to create AppContext for tests with MCP config from file
 pub async fn create_test_context_with_mcp_config(
     config: RouterConfig,