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                use_sdk_mode: false,
55                auto_respawn_on_crash: false,
56                shim_health_check_interval_secs: 60,
57                shim_health_timeout_secs: 120,
58                shim_shutdown_timeout_secs: 30,
59                shim_working_state_timeout_secs: 1800,
60                pending_queue_max_age_secs: 600,
61                event_log_max_bytes: super::DEFAULT_EVENT_LOG_MAX_BYTES,
62                retro_min_duration_secs: 60,
63                roles: Vec::new(),
64            },
65            session: "test".to_string(),
66            members: Vec::new(),
67            pane_map: HashMap::new(),
68            availability: HashMap::new(),
69        }
70    }
71
72    pub fn project_root(&self) -> &Path {
73        &self.project_root
74    }
75
76    pub fn board_tasks_dir(&self) -> PathBuf {
77        super::team_config_dir(&self.project_root)
78            .join("board")
79            .join("tasks")
80    }
81
82    pub fn inbox_root(&self) -> PathBuf {
83        inbox::inboxes_root(&self.project_root)
84    }
85
86    pub fn with_roles(mut self, roles: Vec<RoleDef>) -> Self {
87        self.team_config.roles = roles;
88        self
89    }
90
91    pub fn with_member(mut self, member: MemberInstance) -> Self {
92        self.members.push(member);
93        self
94    }
95
96    pub fn with_members(mut self, members: Vec<MemberInstance>) -> Self {
97        self.members.extend(members);
98        self
99    }
100
101    pub fn with_availability(mut self, availability: HashMap<String, MemberState>) -> Self {
102        self.availability = availability;
103        self
104    }
105
106    pub fn with_member_state(mut self, member: &str, state: MemberState) -> Self {
107        self.availability.insert(member.to_string(), state);
108        self
109    }
110
111    pub fn with_pane(mut self, member: &str, pane_id: &str) -> Self {
112        self.pane_map
113            .insert(member.to_string(), pane_id.to_string());
114        self
115    }
116
117    pub fn with_board_task(
118        self,
119        id: u32,
120        title: &str,
121        status: &str,
122        claimed_by: Option<&str>,
123    ) -> Self {
124        let tasks_dir = self.board_tasks_dir();
125        std::fs::create_dir_all(&tasks_dir).unwrap();
126
127        let mut content = format!(
128            "---\nid: {id}\ntitle: {title}\nstatus: {status}\npriority: high\nclass: standard\n"
129        );
130        if let Some(claimed_by) = claimed_by {
131            content.push_str(&format!("claimed_by: {claimed_by}\n"));
132        }
133        content.push_str("---\n\nTask description.\n");
134
135        std::fs::write(tasks_dir.join(format!("{id:03}-{title}.md")), content).unwrap();
136        self
137    }
138
139    pub fn with_inbox_message(self, member: &str, message: InboxMessage, delivered: bool) -> Self {
140        let inbox_root = self.inbox_root();
141        inbox::init_inbox(&inbox_root, member).unwrap();
142        let id = inbox::deliver_to_inbox(&inbox_root, &message).unwrap();
143        if delivered {
144            inbox::mark_delivered(&inbox_root, member, &id).unwrap();
145        }
146        self
147    }
148
149    pub fn availability_map(&self) -> HashMap<String, MemberState> {
150        self.availability.clone()
151    }
152
153    pub fn pending_inbox_messages(&self, member: &str) -> Result<Vec<InboxMessage>> {
154        inbox::pending_messages(&self.inbox_root(), member)
155    }
156
157    pub fn build_daemon(&self) -> Result<TeamDaemon> {
158        let mut daemon = TeamDaemon::new(DaemonConfig {
159            project_root: self.project_root.clone(),
160            team_config: self.team_config.clone(),
161            session: self.session.clone(),
162            members: self.members.clone(),
163            pane_map: self.pane_map.clone(),
164        })?;
165        daemon.states = self.availability.clone();
166        // Test panes are assumed to be already running — pre-confirm readiness so
167        // that delivery goes through the normal inject-then-inbox path rather than
168        // the pending-delivery-queue path (which is only for freshly-spawned agents).
169        for watcher in daemon.watchers.values_mut() {
170            watcher.confirm_ready();
171        }
172        Ok(daemon)
173    }
174
175    pub fn daemon_member_count(&self, daemon: &TeamDaemon) -> usize {
176        daemon.config.members.len()
177    }
178
179    pub fn daemon_state(&self, daemon: &TeamDaemon, member: &str) -> Option<MemberState> {
180        daemon.states.get(member).copied()
181    }
182}
183
184impl Drop for TestHarness {
185    fn drop(&mut self) {
186        let _ = std::fs::remove_dir_all(&self.project_root);
187    }
188}
189
190pub fn architect_member(name: &str) -> MemberInstance {
191    MemberInstance {
192        name: name.to_string(),
193        role_name: "architect".to_string(),
194        role_type: RoleType::Architect,
195        agent: Some("claude".to_string()),
196        prompt: None,
197        reports_to: None,
198        use_worktrees: false,
199    }
200}
201
202pub fn manager_member(name: &str, reports_to: Option<&str>) -> MemberInstance {
203    MemberInstance {
204        name: name.to_string(),
205        role_name: "manager".to_string(),
206        role_type: RoleType::Manager,
207        agent: Some("claude".to_string()),
208        prompt: None,
209        reports_to: reports_to.map(str::to_string),
210        use_worktrees: false,
211    }
212}
213
214pub fn engineer_member(
215    name: &str,
216    reports_to: Option<&str>,
217    use_worktrees: bool,
218) -> MemberInstance {
219    MemberInstance {
220        name: name.to_string(),
221        role_name: "eng".to_string(),
222        role_type: RoleType::Engineer,
223        agent: Some("codex".to_string()),
224        prompt: None,
225        reports_to: reports_to.map(str::to_string),
226        use_worktrees,
227    }
228}