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