fastapi_output/components/
shutdown_progress.rs1use crate::mode::OutputMode;
7use crate::themes::FastApiTheme;
8
9const ANSI_RESET: &str = "\x1b[0m";
10
11#[derive(Debug, Clone, Copy, PartialEq, Eq)]
13pub enum ShutdownPhase {
14 GracePeriod,
16 ForceClose,
18 Complete,
20}
21
22impl ShutdownPhase {
23 #[must_use]
25 pub const fn label(self) -> &'static str {
26 match self {
27 Self::GracePeriod => "Grace Period",
28 Self::ForceClose => "Force Close",
29 Self::Complete => "Complete",
30 }
31 }
32
33 fn color(self, theme: &FastApiTheme) -> crate::themes::Color {
34 match self {
35 Self::GracePeriod => theme.info,
36 Self::ForceClose => theme.error,
37 Self::Complete => theme.success,
38 }
39 }
40}
41
42#[derive(Debug, Clone)]
44pub struct ShutdownProgress {
45 pub phase: ShutdownPhase,
47 pub total_connections: usize,
49 pub drained_connections: usize,
51 pub in_flight_requests: usize,
53 pub background_tasks: usize,
55 pub cleanup_done: usize,
57 pub cleanup_total: usize,
59 pub notes: Vec<String>,
61}
62
63impl ShutdownProgress {
64 #[must_use]
66 pub fn new(phase: ShutdownPhase) -> Self {
67 Self {
68 phase,
69 total_connections: 0,
70 drained_connections: 0,
71 in_flight_requests: 0,
72 background_tasks: 0,
73 cleanup_done: 0,
74 cleanup_total: 0,
75 notes: Vec::new(),
76 }
77 }
78
79 #[must_use]
81 pub fn connections(mut self, drained: usize, total: usize) -> Self {
82 self.drained_connections = drained;
83 self.total_connections = total;
84 self
85 }
86
87 #[must_use]
89 pub fn in_flight(mut self, in_flight: usize) -> Self {
90 self.in_flight_requests = in_flight;
91 self
92 }
93
94 #[must_use]
96 pub fn background_tasks(mut self, tasks: usize) -> Self {
97 self.background_tasks = tasks;
98 self
99 }
100
101 #[must_use]
103 pub fn cleanup(mut self, done: usize, total: usize) -> Self {
104 self.cleanup_done = done;
105 self.cleanup_total = total;
106 self
107 }
108
109 #[must_use]
111 pub fn note(mut self, note: impl Into<String>) -> Self {
112 self.notes.push(note.into());
113 self
114 }
115}
116
117#[derive(Debug, Clone)]
119pub struct ShutdownProgressDisplay {
120 mode: OutputMode,
121 theme: FastApiTheme,
122 progress_width: usize,
123 title: Option<String>,
124}
125
126impl ShutdownProgressDisplay {
127 #[must_use]
129 pub fn new(mode: OutputMode) -> Self {
130 Self {
131 mode,
132 theme: FastApiTheme::default(),
133 progress_width: 24,
134 title: Some("Shutdown Progress".to_string()),
135 }
136 }
137
138 #[must_use]
140 pub fn theme(mut self, theme: FastApiTheme) -> Self {
141 self.theme = theme;
142 self
143 }
144
145 #[must_use]
147 pub fn progress_width(mut self, width: usize) -> Self {
148 self.progress_width = width.max(8);
149 self
150 }
151
152 #[must_use]
154 pub fn title(mut self, title: Option<String>) -> Self {
155 self.title = title;
156 self
157 }
158
159 #[must_use]
161 pub fn render(&self, progress: &ShutdownProgress) -> String {
162 let mut lines = Vec::new();
163
164 if let Some(title) = &self.title {
165 lines.push(title.clone());
166 lines.push("-".repeat(title.len()));
167 }
168
169 lines.push(self.render_phase(progress.phase));
170
171 if progress.total_connections > 0 {
172 lines.push(self.render_connections(progress));
173 } else {
174 lines.push("Connections: none".to_string());
175 }
176
177 if progress.in_flight_requests > 0 {
178 lines.push(format!(
179 "In-flight requests: {}",
180 progress.in_flight_requests
181 ));
182 }
183
184 if progress.background_tasks > 0 {
185 lines.push(format!("Background tasks: {}", progress.background_tasks));
186 }
187
188 if progress.cleanup_total > 0 {
189 lines.push(format!(
190 "Cleanup: {}/{} steps",
191 progress.cleanup_done, progress.cleanup_total
192 ));
193 }
194
195 for note in &progress.notes {
196 lines.push(format!("Note: {note}"));
197 }
198
199 if progress.phase == ShutdownPhase::Complete {
200 lines.push(self.render_complete());
201 }
202
203 lines.join("\n")
204 }
205
206 fn render_phase(&self, phase: ShutdownPhase) -> String {
207 if self.mode.uses_ansi() {
208 let mut line = format!(
209 "{}Phase:{} {}{}",
210 self.theme.muted.to_ansi_fg(),
211 ANSI_RESET,
212 phase.color(&self.theme).to_ansi_fg(),
213 phase.label()
214 );
215 line.push_str(ANSI_RESET);
216 line
217 } else {
218 format!("Phase: {}", phase.label())
219 }
220 }
221
222 fn render_connections(&self, progress: &ShutdownProgress) -> String {
223 let bar = shutdown_bar(
224 progress.drained_connections,
225 progress.total_connections,
226 self.progress_width,
227 self.mode,
228 &self.theme,
229 );
230 format!(
231 "Connections: {}/{} drained {bar}",
232 progress.drained_connections, progress.total_connections
233 )
234 }
235
236 fn render_complete(&self) -> String {
237 if self.mode.uses_ansi() {
238 format!(
239 "{}Shutdown complete{}",
240 self.theme.success.to_ansi_fg(),
241 ANSI_RESET
242 )
243 } else {
244 "Shutdown complete".to_string()
245 }
246 }
247}
248
249fn shutdown_bar(
250 drained: usize,
251 total: usize,
252 width: usize,
253 mode: OutputMode,
254 theme: &FastApiTheme,
255) -> String {
256 if total == 0 {
257 return String::new();
258 }
259
260 let width = width.max(8);
261 let filled = drained.saturating_mul(width) / total;
262 let filled = filled.min(width);
263 let remaining = width.saturating_sub(filled);
264
265 let mut bar = String::new();
266 bar.push('[');
267
268 if mode.uses_ansi() {
269 if filled > 0 {
270 bar.push_str(&theme.success.to_ansi_fg());
271 bar.push_str(&"#".repeat(filled));
272 bar.push_str(ANSI_RESET);
273 }
274 if remaining > 0 {
275 bar.push_str(&theme.muted.to_ansi_fg());
276 bar.push_str(&"-".repeat(remaining));
277 bar.push_str(ANSI_RESET);
278 }
279 } else {
280 bar.push_str(&"#".repeat(filled));
281 bar.push_str(&"-".repeat(remaining));
282 }
283
284 bar.push(']');
285 bar
286}
287
288#[cfg(test)]
289mod tests {
290 use super::*;
291 use crate::testing::{assert_contains, assert_has_ansi, assert_no_ansi};
292
293 #[test]
298 fn test_shutdown_phase_labels() {
299 assert_eq!(ShutdownPhase::GracePeriod.label(), "Grace Period");
300 assert_eq!(ShutdownPhase::ForceClose.label(), "Force Close");
301 assert_eq!(ShutdownPhase::Complete.label(), "Complete");
302 }
303
304 #[test]
305 fn test_shutdown_phase_equality() {
306 assert_eq!(ShutdownPhase::GracePeriod, ShutdownPhase::GracePeriod);
307 assert_ne!(ShutdownPhase::GracePeriod, ShutdownPhase::ForceClose);
308 assert_ne!(ShutdownPhase::ForceClose, ShutdownPhase::Complete);
309 }
310
311 #[test]
316 fn test_shutdown_progress_new() {
317 let progress = ShutdownProgress::new(ShutdownPhase::GracePeriod);
318 assert_eq!(progress.phase, ShutdownPhase::GracePeriod);
319 assert_eq!(progress.total_connections, 0);
320 assert_eq!(progress.drained_connections, 0);
321 assert_eq!(progress.in_flight_requests, 0);
322 assert_eq!(progress.background_tasks, 0);
323 assert_eq!(progress.cleanup_done, 0);
324 assert_eq!(progress.cleanup_total, 0);
325 assert!(progress.notes.is_empty());
326 }
327
328 #[test]
329 fn test_shutdown_progress_connections() {
330 let progress = ShutdownProgress::new(ShutdownPhase::GracePeriod).connections(5, 10);
331 assert_eq!(progress.drained_connections, 5);
332 assert_eq!(progress.total_connections, 10);
333 }
334
335 #[test]
336 fn test_shutdown_progress_in_flight() {
337 let progress = ShutdownProgress::new(ShutdownPhase::GracePeriod).in_flight(3);
338 assert_eq!(progress.in_flight_requests, 3);
339 }
340
341 #[test]
342 fn test_shutdown_progress_background_tasks() {
343 let progress = ShutdownProgress::new(ShutdownPhase::GracePeriod).background_tasks(2);
344 assert_eq!(progress.background_tasks, 2);
345 }
346
347 #[test]
348 fn test_shutdown_progress_cleanup() {
349 let progress = ShutdownProgress::new(ShutdownPhase::GracePeriod).cleanup(1, 5);
350 assert_eq!(progress.cleanup_done, 1);
351 assert_eq!(progress.cleanup_total, 5);
352 }
353
354 #[test]
355 fn test_shutdown_progress_note() {
356 let progress = ShutdownProgress::new(ShutdownPhase::GracePeriod)
357 .note("First note")
358 .note("Second note");
359 assert_eq!(progress.notes.len(), 2);
360 assert_eq!(progress.notes[0], "First note");
361 assert_eq!(progress.notes[1], "Second note");
362 }
363
364 #[test]
365 fn test_shutdown_progress_full_builder() {
366 let progress = ShutdownProgress::new(ShutdownPhase::ForceClose)
367 .connections(8, 10)
368 .in_flight(1)
369 .background_tasks(2)
370 .cleanup(3, 4)
371 .note("Forcing connections");
372
373 assert_eq!(progress.phase, ShutdownPhase::ForceClose);
374 assert_eq!(progress.drained_connections, 8);
375 assert_eq!(progress.total_connections, 10);
376 assert_eq!(progress.in_flight_requests, 1);
377 assert_eq!(progress.background_tasks, 2);
378 assert_eq!(progress.cleanup_done, 3);
379 assert_eq!(progress.cleanup_total, 4);
380 assert_eq!(progress.notes.len(), 1);
381 }
382
383 #[test]
388 fn test_display_custom_title() {
389 let progress = ShutdownProgress::new(ShutdownPhase::GracePeriod);
390 let display = ShutdownProgressDisplay::new(OutputMode::Plain)
391 .title(Some("Server Shutdown".to_string()));
392 let output = display.render(&progress);
393
394 assert_contains(&output, "Server Shutdown");
395 assert!(!output.contains("Shutdown Progress"));
396 }
397
398 #[test]
399 fn test_display_no_title() {
400 let progress = ShutdownProgress::new(ShutdownPhase::GracePeriod);
401 let display = ShutdownProgressDisplay::new(OutputMode::Plain).title(None);
402 let output = display.render(&progress);
403
404 assert!(!output.contains("Shutdown Progress"));
405 assert_contains(&output, "Phase:");
406 }
407
408 #[test]
409 fn test_display_progress_width() {
410 let progress = ShutdownProgress::new(ShutdownPhase::GracePeriod).connections(5, 10);
411 let display = ShutdownProgressDisplay::new(OutputMode::Plain).progress_width(10);
412 let output = display.render(&progress);
413
414 assert_contains(&output, "[#####-----]");
416 }
417
418 #[test]
419 fn test_display_progress_width_minimum() {
420 let progress = ShutdownProgress::new(ShutdownPhase::GracePeriod).connections(4, 8);
421 let display = ShutdownProgressDisplay::new(OutputMode::Plain).progress_width(2); let output = display.render(&progress);
423
424 assert!(output.contains("[####----]"));
426 }
427
428 #[test]
433 fn renders_plain_shutdown_progress() {
434 let progress = ShutdownProgress::new(ShutdownPhase::GracePeriod)
435 .connections(3, 10)
436 .in_flight(2)
437 .background_tasks(1)
438 .cleanup(1, 3)
439 .note("Waiting for DB pool");
440 let display = ShutdownProgressDisplay::new(OutputMode::Plain);
441 let output = display.render(&progress);
442
443 assert_contains(&output, "Shutdown Progress");
444 assert_contains(&output, "Phase: Grace Period");
445 assert_contains(&output, "Connections: 3/10 drained");
446 assert_contains(&output, "In-flight requests: 2");
447 assert_contains(&output, "Background tasks: 1");
448 assert_contains(&output, "Cleanup: 1/3 steps");
449 assert_contains(&output, "Note: Waiting for DB pool");
450 assert_no_ansi(&output);
451 }
452
453 #[test]
454 fn renders_rich_shutdown_progress() {
455 let progress = ShutdownProgress::new(ShutdownPhase::GracePeriod).connections(5, 10);
456 let display = ShutdownProgressDisplay::new(OutputMode::Rich);
457 let output = display.render(&progress);
458
459 assert_has_ansi(&output);
460 assert_contains(&output, "Grace Period");
461 assert_contains(&output, "5/10 drained");
462 }
463
464 #[test]
465 fn renders_complete_phase() {
466 let progress = ShutdownProgress::new(ShutdownPhase::Complete);
467 let display = ShutdownProgressDisplay::new(OutputMode::Plain);
468 let output = display.render(&progress);
469 assert_contains(&output, "Shutdown complete");
470 }
471
472 #[test]
473 fn renders_rich_complete_phase_with_ansi() {
474 let progress = ShutdownProgress::new(ShutdownPhase::Complete);
475 let display = ShutdownProgressDisplay::new(OutputMode::Rich);
476 let output = display.render(&progress);
477
478 assert_has_ansi(&output);
479 assert_contains(&output, "Shutdown complete");
480 }
481
482 #[test]
487 fn test_grace_period_phase() {
488 let progress = ShutdownProgress::new(ShutdownPhase::GracePeriod);
489 let display = ShutdownProgressDisplay::new(OutputMode::Plain);
490 let output = display.render(&progress);
491
492 assert_contains(&output, "Phase: Grace Period");
493 }
494
495 #[test]
496 fn test_force_close_phase() {
497 let progress = ShutdownProgress::new(ShutdownPhase::ForceClose);
498 let display = ShutdownProgressDisplay::new(OutputMode::Plain);
499 let output = display.render(&progress);
500
501 assert_contains(&output, "Phase: Force Close");
502 }
503
504 #[test]
509 fn test_no_connections() {
510 let progress = ShutdownProgress::new(ShutdownPhase::GracePeriod);
511 let display = ShutdownProgressDisplay::new(OutputMode::Plain);
512 let output = display.render(&progress);
513
514 assert_contains(&output, "Connections: none");
515 }
516
517 #[test]
518 fn test_zero_total_connections() {
519 let progress = ShutdownProgress::new(ShutdownPhase::GracePeriod).connections(0, 0);
520 let display = ShutdownProgressDisplay::new(OutputMode::Plain);
521 let output = display.render(&progress);
522
523 assert_contains(&output, "Connections: none");
525 }
526
527 #[test]
528 fn test_all_connections_drained() {
529 let progress = ShutdownProgress::new(ShutdownPhase::GracePeriod).connections(10, 10);
530 let display = ShutdownProgressDisplay::new(OutputMode::Plain).progress_width(10);
531 let output = display.render(&progress);
532
533 assert_contains(&output, "10/10 drained");
534 assert_contains(&output, "[##########]");
535 }
536
537 #[test]
538 fn test_no_in_flight_requests_omitted() {
539 let progress = ShutdownProgress::new(ShutdownPhase::GracePeriod).in_flight(0);
540 let display = ShutdownProgressDisplay::new(OutputMode::Plain);
541 let output = display.render(&progress);
542
543 assert!(!output.contains("In-flight"));
544 }
545
546 #[test]
547 fn test_no_background_tasks_omitted() {
548 let progress = ShutdownProgress::new(ShutdownPhase::GracePeriod).background_tasks(0);
549 let display = ShutdownProgressDisplay::new(OutputMode::Plain);
550 let output = display.render(&progress);
551
552 assert!(!output.contains("Background tasks"));
553 }
554
555 #[test]
556 fn test_no_cleanup_omitted() {
557 let progress = ShutdownProgress::new(ShutdownPhase::GracePeriod).cleanup(0, 0);
558 let display = ShutdownProgressDisplay::new(OutputMode::Plain);
559 let output = display.render(&progress);
560
561 assert!(!output.contains("Cleanup:"));
562 }
563
564 #[test]
565 fn test_multiple_notes() {
566 let progress = ShutdownProgress::new(ShutdownPhase::GracePeriod)
567 .note("Note 1")
568 .note("Note 2")
569 .note("Note 3");
570 let display = ShutdownProgressDisplay::new(OutputMode::Plain);
571 let output = display.render(&progress);
572
573 assert_contains(&output, "Note: Note 1");
574 assert_contains(&output, "Note: Note 2");
575 assert_contains(&output, "Note: Note 3");
576 }
577
578 #[test]
583 fn test_progress_bar_empty() {
584 let progress = ShutdownProgress::new(ShutdownPhase::GracePeriod).connections(0, 10);
585 let display = ShutdownProgressDisplay::new(OutputMode::Plain).progress_width(10);
586 let output = display.render(&progress);
587
588 assert_contains(&output, "[----------]");
589 }
590
591 #[test]
592 fn test_progress_bar_half() {
593 let progress = ShutdownProgress::new(ShutdownPhase::GracePeriod).connections(5, 10);
594 let display = ShutdownProgressDisplay::new(OutputMode::Plain).progress_width(10);
595 let output = display.render(&progress);
596
597 assert_contains(&output, "[#####-----]");
598 }
599
600 #[test]
601 fn test_progress_bar_full() {
602 let progress = ShutdownProgress::new(ShutdownPhase::GracePeriod).connections(10, 10);
603 let display = ShutdownProgressDisplay::new(OutputMode::Plain).progress_width(10);
604 let output = display.render(&progress);
605
606 assert_contains(&output, "[##########]");
607 }
608
609 #[test]
610 fn test_progress_bar_one_of_many() {
611 let progress = ShutdownProgress::new(ShutdownPhase::GracePeriod).connections(1, 100);
612 let display = ShutdownProgressDisplay::new(OutputMode::Plain).progress_width(20);
613 let output = display.render(&progress);
614
615 assert!(output.contains("[--------------------]"));
617 }
618}