huskies/server/src/agents/pool/test_helpers.rs

use crate::worktree::WorktreeInfo;
use std::path::PathBuf;
use std::sync::{Arc, Mutex};
use tokio::sync::broadcast;

use super::super::{AgentEvent, AgentStatus, CompletionReport};
use super::types::{StoryAgent, composite_key};
use super::AgentPool;

impl AgentPool {
    /// Test helper: inject a pre-built agent entry so unit tests can exercise
    /// wait/subscribe logic without spawning a real process.
    pub fn inject_test_agent(
        &self,
        story_id: &str,
        agent_name: &str,
        status: AgentStatus,
    ) -> broadcast::Sender<AgentEvent> {
        let (tx, _) = broadcast::channel::<AgentEvent>(64);
        let key = composite_key(story_id, agent_name);
        let mut agents = self.agents.lock().unwrap();
        agents.insert(
            key,
            StoryAgent {
                agent_name: agent_name.to_string(),
                status,
                worktree_info: None,
                session_id: None,
                tx: tx.clone(),
                task_handle: None,
                event_log: Arc::new(Mutex::new(Vec::new())),
                completion: None,
                project_root: None,
                log_session_id: None,
                merge_failure_reported: false,
                throttled: false,
            },
        );
        tx
    }

    /// Test helper: inject an agent with a specific worktree path for testing
    /// gate-related logic.
    pub fn inject_test_agent_with_path(
        &self,
        story_id: &str,
        agent_name: &str,
        status: AgentStatus,
        worktree_path: PathBuf,
    ) -> broadcast::Sender<AgentEvent> {
        let (tx, _) = broadcast::channel::<AgentEvent>(64);
        let key = composite_key(story_id, agent_name);
        let mut agents = self.agents.lock().unwrap();
        agents.insert(
            key,
            StoryAgent {
                agent_name: agent_name.to_string(),
                status,
                worktree_info: Some(WorktreeInfo {
                    path: worktree_path,
                    branch: format!("feature/story-{story_id}"),
                    base_branch: "master".to_string(),
                }),
                session_id: None,
                tx: tx.clone(),
                task_handle: None,
                event_log: Arc::new(Mutex::new(Vec::new())),
                completion: None,
                project_root: None,
                log_session_id: None,
                merge_failure_reported: false,
                throttled: false,
            },
        );
        tx
    }

    /// Test helper: inject an agent with a completion report and project_root
    /// for testing pipeline advance logic without spawning real agents.
    pub fn inject_test_agent_with_completion(
        &self,
        story_id: &str,
        agent_name: &str,
        status: AgentStatus,
        project_root: PathBuf,
        completion: CompletionReport,
    ) -> broadcast::Sender<AgentEvent> {
        let (tx, _) = broadcast::channel::<AgentEvent>(64);
        let key = composite_key(story_id, agent_name);
        let mut agents = self.agents.lock().unwrap();
        agents.insert(
            key,
            StoryAgent {
                agent_name: agent_name.to_string(),
                status,
                worktree_info: None,
                session_id: None,
                tx: tx.clone(),
                task_handle: None,
                event_log: Arc::new(Mutex::new(Vec::new())),
                completion: Some(completion),
                project_root: Some(project_root),
                log_session_id: None,
                merge_failure_reported: false,
                throttled: false,
            },
        );
        tx
    }

    /// Inject a Running agent with a pre-built (possibly finished) task handle.
    /// Used by watchdog tests to simulate an orphaned agent.
    pub fn inject_test_agent_with_handle(
        &self,
        story_id: &str,
        agent_name: &str,
        status: AgentStatus,
        task_handle: tokio::task::JoinHandle<()>,
    ) -> broadcast::Sender<AgentEvent> {
        let (tx, _) = broadcast::channel::<AgentEvent>(64);
        let key = composite_key(story_id, agent_name);
        let mut agents = self.agents.lock().unwrap();
        agents.insert(
            key,
            StoryAgent {
                agent_name: agent_name.to_string(),
                status,
                worktree_info: None,
                session_id: None,
                tx: tx.clone(),
                task_handle: Some(task_handle),
                event_log: Arc::new(Mutex::new(Vec::new())),
                completion: None,
                project_root: None,
                log_session_id: None,
                merge_failure_reported: false,
                throttled: false,
            },
        );
        tx
    }
}