Skip to main content

batty_cli/team/
harness.rs

1use std::collections::HashMap;
2use std::path::{Path, PathBuf};
3
4use anyhow::Result;
5use uuid::Uuid;
6
7use super::config::{
8    AutomationConfig, BoardConfig, OrchestratorPosition, RoleDef, RoleType, StandupConfig,
9    TeamConfig, WorkflowMode, WorkflowPolicy,
10};
11use super::daemon::{DaemonConfig, TeamDaemon};
12use super::hierarchy::MemberInstance;
13use super::inbox::{self, InboxMessage};
14use super::standup::MemberState;
15
16pub struct TestHarness {
17    project_root: PathBuf,
18    team_config: TeamConfig,
19    session: String,
20    members: Vec<MemberInstance>,
21    pane_map: HashMap<String, String>,
22    availability: HashMap<String, MemberState>,
23}
24
25impl Default for TestHarness {
26    fn default() -> Self {
27        Self::new()
28    }
29}
30
31impl TestHarness {
32    pub fn new() -> Self {
33        let project_root =
34            std::env::temp_dir().join(format!("batty-test-harness-{}", Uuid::new_v4()));
35        std::fs::create_dir_all(super::team_config_dir(&project_root)).unwrap();
36        Self {
37            project_root,
38            team_config: TeamConfig {
39                name: "test".to_string(),
40                agent: None,
41                workflow_mode: WorkflowMode::Legacy,
42                workflow_policy: WorkflowPolicy::default(),
43                board: BoardConfig::default(),
44                standup: StandupConfig::default(),
45                automation: AutomationConfig::default(),
46                automation_sender: None,
47                external_senders: Vec::new(),
48                orchestrator_pane: false,
49                orchestrator_position: OrchestratorPosition::Bottom,
50                layout: None,
51                cost: Default::default(),
52                grafana: Default::default(),
53                use_shim: false,
54                auto_respawn_on_crash: false,
55                shim_health_check_interval_secs: 60,
56                shim_health_timeout_secs: 120,
57                shim_shutdown_timeout_secs: 30,
58                event_log_max_bytes: super::DEFAULT_EVENT_LOG_MAX_BYTES,
59                retro_min_duration_secs: 60,
60                roles: Vec::new(),
61            },
62            session: "test".to_string(),
63            members: Vec::new(),
64            pane_map: HashMap::new(),
65            availability: HashMap::new(),
66        }
67    }
68
69    pub fn project_root(&self) -> &Path {
70        &self.project_root
71    }
72
73    pub fn board_tasks_dir(&self) -> PathBuf {
74        super::team_config_dir(&self.project_root)
75            .join("board")
76            .join("tasks")
77    }
78
79    pub fn inbox_root(&self) -> PathBuf {
80        inbox::inboxes_root(&self.project_root)
81    }
82
83    pub fn with_roles(mut self, roles: Vec<RoleDef>) -> Self {
84        self.team_config.roles = roles;
85        self
86    }
87
88    pub fn with_member(mut self, member: MemberInstance) -> Self {
89        self.members.push(member);
90        self
91    }
92
93    pub fn with_members(mut self, members: Vec<MemberInstance>) -> Self {
94        self.members.extend(members);
95        self
96    }
97
98    pub fn with_availability(mut self, availability: HashMap<String, MemberState>) -> Self {
99        self.availability = availability;
100        self
101    }
102
103    pub fn with_member_state(mut self, member: &str, state: MemberState) -> Self {
104        self.availability.insert(member.to_string(), state);
105        self
106    }
107
108    pub fn with_pane(mut self, member: &str, pane_id: &str) -> Self {
109        self.pane_map
110            .insert(member.to_string(), pane_id.to_string());
111        self
112    }
113
114    pub fn with_board_task(
115        self,
116        id: u32,
117        title: &str,
118        status: &str,
119        claimed_by: Option<&str>,
120    ) -> Self {
121        let tasks_dir = self.board_tasks_dir();
122        std::fs::create_dir_all(&tasks_dir).unwrap();
123
124        let mut content = format!(
125            "---\nid: {id}\ntitle: {title}\nstatus: {status}\npriority: high\nclass: standard\n"
126        );
127        if let Some(claimed_by) = claimed_by {
128            content.push_str(&format!("claimed_by: {claimed_by}\n"));
129        }
130        content.push_str("---\n\nTask description.\n");
131
132        std::fs::write(tasks_dir.join(format!("{id:03}-{title}.md")), content).unwrap();
133        self
134    }
135
136    pub fn with_inbox_message(self, member: &str, message: InboxMessage, delivered: bool) -> Self {
137        let inbox_root = self.inbox_root();
138        inbox::init_inbox(&inbox_root, member).unwrap();
139        let id = inbox::deliver_to_inbox(&inbox_root, &message).unwrap();
140        if delivered {
141            inbox::mark_delivered(&inbox_root, member, &id).unwrap();
142        }
143        self
144    }
145
146    pub fn availability_map(&self) -> HashMap<String, MemberState> {
147        self.availability.clone()
148    }
149
150    pub fn pending_inbox_messages(&self, member: &str) -> Result<Vec<InboxMessage>> {
151        inbox::pending_messages(&self.inbox_root(), member)
152    }
153
154    pub fn build_daemon(&self) -> Result<TeamDaemon> {
155        let mut daemon = TeamDaemon::new(DaemonConfig {
156            project_root: self.project_root.clone(),
157            team_config: self.team_config.clone(),
158            session: self.session.clone(),
159            members: self.members.clone(),
160            pane_map: self.pane_map.clone(),
161        })?;
162        daemon.states = self.availability.clone();
163        // Test panes are assumed to be already running — pre-confirm readiness so
164        // that delivery goes through the normal inject-then-inbox path rather than
165        // the pending-delivery-queue path (which is only for freshly-spawned agents).
166        for watcher in daemon.watchers.values_mut() {
167            watcher.confirm_ready();
168        }
169        Ok(daemon)
170    }
171
172    pub fn daemon_member_count(&self, daemon: &TeamDaemon) -> usize {
173        daemon.config.members.len()
174    }
175
176    pub fn daemon_state(&self, daemon: &TeamDaemon, member: &str) -> Option<MemberState> {
177        daemon.states.get(member).copied()
178    }
179}
180
181impl Drop for TestHarness {
182    fn drop(&mut self) {
183        let _ = std::fs::remove_dir_all(&self.project_root);
184    }
185}
186
187pub fn architect_member(name: &str) -> MemberInstance {
188    MemberInstance {
189        name: name.to_string(),
190        role_name: "architect".to_string(),
191        role_type: RoleType::Architect,
192        agent: Some("claude".to_string()),
193        prompt: None,
194        reports_to: None,
195        use_worktrees: false,
196    }
197}
198
199pub fn manager_member(name: &str, reports_to: Option<&str>) -> MemberInstance {
200    MemberInstance {
201        name: name.to_string(),
202        role_name: "manager".to_string(),
203        role_type: RoleType::Manager,
204        agent: Some("claude".to_string()),
205        prompt: None,
206        reports_to: reports_to.map(str::to_string),
207        use_worktrees: false,
208    }
209}
210
211pub fn engineer_member(
212    name: &str,
213    reports_to: Option<&str>,
214    use_worktrees: bool,
215) -> MemberInstance {
216    MemberInstance {
217        name: name.to_string(),
218        role_name: "eng".to_string(),
219        role_type: RoleType::Engineer,
220        agent: Some("codex".to_string()),
221        prompt: None,
222        reports_to: reports_to.map(str::to_string),
223        use_worktrees,
224    }
225}