rust_expect/multi/
group.rs1use std::collections::HashMap;
4use std::time::Duration;
5
6pub type GroupId = String;
8
9pub type SessionId = usize;
11
12#[derive(Debug, Clone)]
14pub struct GroupResult {
15 pub session_id: SessionId,
17 pub success: bool,
19 pub output: String,
21}
22
23#[derive(Debug, Default)]
25pub struct SessionGroup {
26 name: String,
28 sessions: HashMap<SessionId, SessionInfo>,
30 next_id: SessionId,
32 timeout: Duration,
34}
35
36#[derive(Debug, Clone)]
38struct SessionInfo {
39 label: String,
41 active: bool,
43 output: String,
45}
46
47impl SessionGroup {
48 #[must_use]
50 pub fn new(name: impl Into<String>) -> Self {
51 Self {
52 name: name.into(),
53 timeout: Duration::from_secs(30),
54 ..Default::default()
55 }
56 }
57
58 #[must_use]
60 pub const fn with_timeout(mut self, timeout: Duration) -> Self {
61 self.timeout = timeout;
62 self
63 }
64
65 #[must_use]
67 pub fn name(&self) -> &str {
68 &self.name
69 }
70
71 pub fn add(&mut self, label: impl Into<String>) -> SessionId {
73 let id = self.next_id;
74 self.next_id += 1;
75 self.sessions.insert(
76 id,
77 SessionInfo {
78 label: label.into(),
79 active: true,
80 output: String::new(),
81 },
82 );
83 id
84 }
85
86 pub fn remove(&mut self, id: SessionId) {
88 self.sessions.remove(&id);
89 }
90
91 #[must_use]
93 pub fn len(&self) -> usize {
94 self.sessions.len()
95 }
96
97 #[must_use]
99 pub fn is_empty(&self) -> bool {
100 self.sessions.is_empty()
101 }
102
103 #[must_use]
105 pub fn active_count(&self) -> usize {
106 self.sessions.values().filter(|s| s.active).count()
107 }
108
109 #[must_use]
111 pub fn label(&self, id: SessionId) -> Option<&str> {
112 self.sessions.get(&id).map(|s| s.label.as_str())
113 }
114
115 pub fn set_active(&mut self, id: SessionId, active: bool) {
117 if let Some(session) = self.sessions.get_mut(&id) {
118 session.active = active;
119 }
120 }
121
122 pub fn append_output(&mut self, id: SessionId, output: &str) {
124 if let Some(session) = self.sessions.get_mut(&id) {
125 session.output.push_str(output);
126 }
127 }
128
129 #[must_use]
131 pub fn output(&self, id: SessionId) -> Option<&str> {
132 self.sessions.get(&id).map(|s| s.output.as_str())
133 }
134
135 pub fn clear_output(&mut self) {
137 for session in self.sessions.values_mut() {
138 session.output.clear();
139 }
140 }
141
142 #[must_use]
144 pub fn session_ids(&self) -> Vec<SessionId> {
145 self.sessions.keys().copied().collect()
146 }
147
148 #[must_use]
150 pub fn active_ids(&self) -> Vec<SessionId> {
151 self.sessions
152 .iter()
153 .filter(|(_, s)| s.active)
154 .map(|(id, _)| *id)
155 .collect()
156 }
157
158 pub fn for_each<F>(&self, mut f: F)
160 where
161 F: FnMut(SessionId, &str),
162 {
163 for (id, session) in &self.sessions {
164 if session.active {
165 f(*id, &session.label);
166 }
167 }
168 }
169}
170
171#[derive(Debug, Default)]
173pub struct GroupBuilder {
174 name: String,
175 timeout: Duration,
176 labels: Vec<String>,
177}
178
179impl GroupBuilder {
180 #[must_use]
182 pub fn new(name: impl Into<String>) -> Self {
183 Self {
184 name: name.into(),
185 timeout: Duration::from_secs(30),
186 labels: Vec::new(),
187 }
188 }
189
190 #[must_use]
192 pub const fn timeout(mut self, timeout: Duration) -> Self {
193 self.timeout = timeout;
194 self
195 }
196
197 #[must_use]
199 #[allow(clippy::should_implement_trait)]
200 pub fn add(mut self, label: impl Into<String>) -> Self {
201 self.labels.push(label.into());
202 self
203 }
204
205 #[must_use]
207 pub fn build(self) -> SessionGroup {
208 let mut group = SessionGroup::new(self.name).with_timeout(self.timeout);
209 for label in self.labels {
210 group.add(label);
211 }
212 group
213 }
214}
215
216#[derive(Debug, Default)]
218pub struct GroupManager {
219 groups: HashMap<GroupId, SessionGroup>,
221}
222
223impl GroupManager {
224 #[must_use]
226 pub fn new() -> Self {
227 Self::default()
228 }
229
230 pub fn create(&mut self, name: impl Into<String>) -> &mut SessionGroup {
232 let name = name.into();
233 self.groups
234 .entry(name.clone())
235 .or_insert_with(|| SessionGroup::new(name))
236 }
237
238 #[must_use]
240 pub fn get(&self, name: &str) -> Option<&SessionGroup> {
241 self.groups.get(name)
242 }
243
244 pub fn get_mut(&mut self, name: &str) -> Option<&mut SessionGroup> {
246 self.groups.get_mut(name)
247 }
248
249 pub fn remove(&mut self, name: &str) -> Option<SessionGroup> {
251 self.groups.remove(name)
252 }
253
254 #[must_use]
256 pub fn names(&self) -> Vec<&str> {
257 self.groups
258 .keys()
259 .map(std::string::String::as_str)
260 .collect()
261 }
262
263 #[must_use]
265 pub fn total_sessions(&self) -> usize {
266 self.groups.values().map(SessionGroup::len).sum()
267 }
268}
269
270#[cfg(test)]
271mod tests {
272 use super::*;
273
274 #[test]
275 fn group_basic() {
276 let mut group = SessionGroup::new("test");
277 let id1 = group.add("server1");
278 let id2 = group.add("server2");
279
280 assert_eq!(group.len(), 2);
281 assert_eq!(group.label(id1), Some("server1"));
282 assert_eq!(group.label(id2), Some("server2"));
283 }
284
285 #[test]
286 fn group_builder() {
287 let group = GroupBuilder::new("servers")
288 .add("server1")
289 .add("server2")
290 .add("server3")
291 .build();
292
293 assert_eq!(group.name(), "servers");
294 assert_eq!(group.len(), 3);
295 }
296
297 #[test]
298 fn group_manager() {
299 let mut manager = GroupManager::new();
300 manager.create("web").add("web1");
301 manager.create("db").add("db1");
302
303 assert_eq!(manager.names().len(), 2);
304 assert_eq!(manager.total_sessions(), 2);
305 }
306
307 #[test]
308 fn group_active() {
309 let mut group = SessionGroup::new("test");
310 let id = group.add("server");
311
312 assert_eq!(group.active_count(), 1);
313
314 group.set_active(id, false);
315 assert_eq!(group.active_count(), 0);
316 }
317}