use crate::models::ExecutionPlan;
use serde::{Deserialize, Serialize};
use std::sync::{Arc, Mutex};
use std::time::{Duration, Instant};
use tracing::{debug, info};
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ProgressUpdate {
pub current_step: usize,
pub total_steps: usize,
pub progress_percentage: f32,
pub estimated_time_remaining: Duration,
pub timestamp: chrono::DateTime<chrono::Utc>,
}
pub type ProgressCallback = Box<dyn Fn(ProgressUpdate) + Send + Sync>;
pub struct ProgressTracker {
total_steps: usize,
current_step: usize,
completed_steps: usize,
start_time: Instant,
step_durations: Vec<Duration>,
callbacks: Arc<Mutex<Vec<ProgressCallback>>>,
}
impl ProgressTracker {
pub fn new(plan: &ExecutionPlan) -> Self {
let total_steps = plan.steps.len();
info!(total_steps = total_steps, "Creating progress tracker");
Self {
total_steps,
current_step: 0,
completed_steps: 0,
start_time: Instant::now(),
step_durations: Vec::new(),
callbacks: Arc::new(Mutex::new(Vec::new())),
}
}
pub fn on_progress<F>(&self, callback: F)
where
F: Fn(ProgressUpdate) + Send + Sync + 'static,
{
let mut callbacks = self.callbacks.lock().unwrap();
callbacks.push(Box::new(callback));
debug!(
callback_count = callbacks.len(),
"Progress callback registered"
);
}
pub fn step_completed(&mut self, step_duration: Duration) {
self.step_durations.push(step_duration);
self.completed_steps += 1;
self.current_step += 1;
debug!(
current_step = self.current_step,
completed_steps = self.completed_steps,
step_duration_ms = step_duration.as_millis(),
"Step completed"
);
self.notify_progress();
}
pub fn step_skipped(&mut self) {
self.step_durations.push(Duration::from_secs(0));
self.current_step += 1;
debug!(current_step = self.current_step, "Step skipped");
self.notify_progress();
}
pub fn get_progress(&self) -> ProgressUpdate {
let progress_percentage = if self.total_steps > 0 {
(self.completed_steps as f32 / self.total_steps as f32) * 100.0
} else {
0.0
};
let estimated_time_remaining = self.estimated_time_remaining();
ProgressUpdate {
current_step: self.current_step,
total_steps: self.total_steps,
progress_percentage,
estimated_time_remaining,
timestamp: chrono::Utc::now(),
}
}
pub fn current_step(&self) -> usize {
self.current_step
}
pub fn total_steps(&self) -> usize {
self.total_steps
}
pub fn completed_steps(&self) -> usize {
self.completed_steps
}
pub fn progress_percentage(&self) -> f32 {
if self.total_steps > 0 {
(self.completed_steps as f32 / self.total_steps as f32) * 100.0
} else {
0.0
}
}
pub fn estimated_time_remaining(&self) -> Duration {
if self.step_durations.is_empty() || self.completed_steps == 0 {
return Duration::from_secs(0);
}
let total_duration: Duration = self.step_durations.iter().sum();
let average_duration = total_duration / self.step_durations.len() as u32;
let remaining_steps = self.total_steps.saturating_sub(self.completed_steps);
average_duration * remaining_steps as u32
}
pub fn elapsed_time(&self) -> Duration {
self.start_time.elapsed()
}
pub fn average_step_duration(&self) -> Duration {
if self.step_durations.is_empty() {
return Duration::from_secs(0);
}
let total_duration: Duration = self.step_durations.iter().sum();
total_duration / self.step_durations.len() as u32
}
fn notify_progress(&self) {
let progress = self.get_progress();
let callbacks = self.callbacks.lock().unwrap();
for callback in callbacks.iter() {
callback(progress.clone());
}
}
pub fn reset(&mut self) {
self.current_step = 0;
self.completed_steps = 0;
self.start_time = Instant::now();
self.step_durations.clear();
debug!("Progress tracker reset");
}
}
impl Default for ProgressTracker {
fn default() -> Self {
Self {
total_steps: 0,
current_step: 0,
completed_steps: 0,
start_time: Instant::now(),
step_durations: Vec::new(),
callbacks: Arc::new(Mutex::new(Vec::new())),
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::models::{ExecutionPlan, ExecutionStep, RiskScore, StepAction, StepStatus};
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::Arc as StdArc;
fn create_test_plan(step_count: usize) -> ExecutionPlan {
let steps = (0..step_count)
.map(|i| ExecutionStep {
id: format!("step-{}", i),
description: format!("Step {}", i),
action: StepAction::RunCommand {
command: "echo".to_string(),
args: vec![format!("step {}", i)],
},
risk_score: RiskScore::default(),
dependencies: Vec::new(),
rollback_action: None,
status: StepStatus::Pending,
})
.collect();
ExecutionPlan::new("Test Plan".to_string(), steps)
}
#[test]
fn test_create_tracker() {
let plan = create_test_plan(5);
let tracker = ProgressTracker::new(&plan);
assert_eq!(tracker.total_steps(), 5);
assert_eq!(tracker.current_step(), 0);
assert_eq!(tracker.completed_steps(), 0);
assert_eq!(tracker.progress_percentage(), 0.0);
}
#[test]
fn test_step_completed() {
let plan = create_test_plan(5);
let mut tracker = ProgressTracker::new(&plan);
tracker.step_completed(Duration::from_secs(1));
assert_eq!(tracker.completed_steps(), 1);
assert_eq!(tracker.current_step(), 1);
assert_eq!(tracker.progress_percentage(), 20.0);
}
#[test]
fn test_multiple_steps_completed() {
let plan = create_test_plan(5);
let mut tracker = ProgressTracker::new(&plan);
tracker.step_completed(Duration::from_secs(1));
tracker.step_completed(Duration::from_secs(2));
tracker.step_completed(Duration::from_secs(1));
assert_eq!(tracker.completed_steps(), 3);
assert!((tracker.progress_percentage() - 60.0).abs() < 0.01);
}
#[test]
fn test_step_skipped() {
let plan = create_test_plan(5);
let mut tracker = ProgressTracker::new(&plan);
tracker.step_completed(Duration::from_secs(1));
tracker.step_skipped();
assert_eq!(tracker.completed_steps(), 1);
assert_eq!(tracker.current_step(), 2);
}
#[test]
fn test_progress_percentage() {
let plan = create_test_plan(10);
let mut tracker = ProgressTracker::new(&plan);
for _ in 0..5 {
tracker.step_completed(Duration::from_secs(1));
}
assert_eq!(tracker.progress_percentage(), 50.0);
}
#[test]
fn test_estimated_time_remaining() {
let plan = create_test_plan(10);
let mut tracker = ProgressTracker::new(&plan);
tracker.step_completed(Duration::from_secs(1));
tracker.step_completed(Duration::from_secs(1));
let estimated = tracker.estimated_time_remaining();
assert!(estimated.as_secs() >= 7 && estimated.as_secs() <= 9);
}
#[test]
fn test_average_step_duration() {
let plan = create_test_plan(5);
let mut tracker = ProgressTracker::new(&plan);
tracker.step_completed(Duration::from_secs(2));
tracker.step_completed(Duration::from_secs(4));
let average = tracker.average_step_duration();
assert_eq!(average, Duration::from_secs(3));
}
#[test]
fn test_elapsed_time() {
let plan = create_test_plan(5);
let tracker = ProgressTracker::new(&plan);
let elapsed = tracker.elapsed_time();
let _ = elapsed;
}
#[test]
fn test_progress_callback() {
let plan = create_test_plan(5);
let mut tracker = ProgressTracker::new(&plan);
let callback_count = StdArc::new(AtomicUsize::new(0));
let callback_count_clone = callback_count.clone();
tracker.on_progress(move |_progress| {
callback_count_clone.fetch_add(1, Ordering::SeqCst);
});
tracker.step_completed(Duration::from_secs(1));
tracker.step_completed(Duration::from_secs(1));
assert_eq!(callback_count.load(Ordering::SeqCst), 2);
}
#[test]
fn test_get_progress() {
let plan = create_test_plan(5);
let mut tracker = ProgressTracker::new(&plan);
tracker.step_completed(Duration::from_secs(1));
let progress = tracker.get_progress();
assert_eq!(progress.current_step, 1);
assert_eq!(progress.total_steps, 5);
assert_eq!(progress.progress_percentage, 20.0);
}
#[test]
fn test_reset() {
let plan = create_test_plan(5);
let mut tracker = ProgressTracker::new(&plan);
tracker.step_completed(Duration::from_secs(1));
tracker.step_completed(Duration::from_secs(1));
tracker.reset();
assert_eq!(tracker.completed_steps(), 0);
assert_eq!(tracker.current_step(), 0);
assert_eq!(tracker.progress_percentage(), 0.0);
}
#[test]
fn test_empty_plan() {
let plan = create_test_plan(0);
let tracker = ProgressTracker::new(&plan);
assert_eq!(tracker.total_steps(), 0);
assert_eq!(tracker.progress_percentage(), 0.0);
}
#[test]
fn test_progress_update_serialization() {
let update = ProgressUpdate {
current_step: 1,
total_steps: 5,
progress_percentage: 20.0,
estimated_time_remaining: Duration::from_secs(4),
timestamp: chrono::Utc::now(),
};
let json = serde_json::to_string(&update).unwrap();
let deserialized: ProgressUpdate = serde_json::from_str(&json).unwrap();
assert_eq!(deserialized.current_step, 1);
assert_eq!(deserialized.total_steps, 5);
assert_eq!(deserialized.progress_percentage, 20.0);
}
}