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

//! Test helpers for the agent pool — in-memory pool construction and assertions.
use crate::service::status::buffer::{BufferedItem, StatusEventBuffer};
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::AgentPool;
use super::types::{StoryAgent, composite_key};

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,
                termination_reason: None,
                status_buffer: None,
            },
        );
        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,
                termination_reason: None,
                status_buffer: None,
            },
        );
        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,
                termination_reason: None,
                status_buffer: None,
            },
        );
        tx
    }

    /// Test helper: inject an agent with a specific log session ID.
    /// Used by watchdog tests to simulate per-session counting.
    pub fn inject_test_agent_with_session(
        &self,
        story_id: &str,
        agent_name: &str,
        status: AgentStatus,
        log_session_id: &str,
    ) -> 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: Some(log_session_id.to_string()),
                merge_failure_reported: false,
                throttled: false,
                termination_reason: None,
                status_buffer: None,
            },
        );
        tx
    }

    /// Test helper: inject an agent whose `status_buffer` is subscribed to the
    /// pool's [`StatusBroadcaster`].  Use [`drain_agent_status_buffer`] to read
    /// accumulated events after publishing to `pool.status_broadcaster()`.
    pub fn inject_test_agent_with_live_buffer(
        &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,
                termination_reason: None,
                status_buffer: Some(StatusEventBuffer::new(&self.status_broadcaster)),
            },
        );
        tx
    }

    /// Test helper: drain all buffered status events from the specified agent's
    /// [`StatusEventBuffer`].  Returns `None` if the agent does not exist or has
    /// no buffer attached.
    pub fn drain_agent_status_buffer(
        &self,
        story_id: &str,
        agent_name: &str,
    ) -> Option<Vec<BufferedItem>> {
        let agents = self.agents.lock().unwrap();
        let key = composite_key(story_id, agent_name);
        agents
            .get(&key)
            .and_then(|a| a.status_buffer.as_ref().map(|b| b.drain()))
    }

    /// Test helper: inject an agent with a project root AND a worktree path.
    ///
    /// Use this when the full server-owned completion path needs both a
    /// `project_root` (so `run_pipeline_advance` can load config and advance
    /// the story) and a `worktree_info` (so gate checks can inspect the branch).
    pub fn inject_test_agent_with_root_and_path(
        &self,
        story_id: &str,
        agent_name: &str,
        status: AgentStatus,
        project_root: PathBuf,
        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: Some(project_root),
                log_session_id: None,
                merge_failure_reported: false,
                throttled: false,
                termination_reason: None,
                status_buffer: None,
            },
        );
        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,
                termination_reason: None,
                status_buffer: None,
            },
        );
        tx
    }
}