use async_trait::async_trait;
use chrono::{DateTime, Utc};
use serde::{Deserialize, Serialize};
use std::collections::{HashMap, VecDeque};
use std::sync::Arc;
use std::time::Duration;
use tokio::sync::RwLock;
use super::config::RollbackConfig;
use super::error::{ActuatorError, ActuatorResult, DeploymentState};
use super::traits::{ConfigurationManager, HealthMonitor};
use super::types::{
ConfigurationSnapshot, DeploymentMetrics, DeploymentStatus, HealthStatus, RollbackReason,
RollbackRequest, RollbackResult,
};
use crate::decision::{ConfigChange, SystemMetrics};
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum RollbackMode {
Fast,
Gradual,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct RollbackPhase {
pub phase: usize,
pub traffic_percent: f64,
pub start_time: Option<DateTime<Utc>>,
pub end_time: Option<DateTime<Utc>>,
pub completed: bool,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct RollbackStatus {
pub rollback_id: String,
pub deployment_id: String,
pub mode: RollbackMode,
pub current_phase: Option<usize>,
pub start_time: DateTime<Utc>,
pub end_time: Option<DateTime<Utc>>,
pub completed: bool,
pub success: bool,
pub error: Option<String>,
pub phases: Vec<RollbackPhase>,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct RollbackHistoryEntry {
pub rollback_id: String,
pub deployment_id: String,
pub request: RollbackRequest,
pub result: RollbackResult,
pub status: RollbackStatus,
pub trigger: RollbackTrigger,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct RollbackTrigger {
pub trigger_type: RollbackReason,
pub timestamp: DateTime<Utc>,
pub automatic: bool,
pub metrics: Option<DeploymentMetrics>,
pub health_status: Option<HealthStatus>,
pub threshold_exceeded: Option<ThresholdViolation>,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ThresholdViolation {
pub metric: String,
pub actual_value: f64,
pub threshold_value: f64,
pub operator: String,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct RollbackNotification {
pub id: String,
pub deployment_id: String,
pub rollback_id: String,
pub reason: RollbackReason,
pub timestamp: DateTime<Utc>,
pub severity: NotificationSeverity,
pub message: String,
pub details: HashMap<String, serde_json::Value>,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum NotificationSeverity {
Info,
Warning,
Error,
Critical,
}
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
pub struct RollbackStats {
pub total_rollbacks: u64,
pub automatic_rollbacks: u64,
pub manual_rollbacks: u64,
pub successful_rollbacks: u64,
pub failed_rollbacks: u64,
pub fast_rollbacks: u64,
pub gradual_rollbacks: u64,
pub avg_rollback_duration_secs: f64,
pub fastest_rollback_secs: f64,
pub rollbacks_by_reason: HashMap<String, u64>,
}
pub struct RollbackEngine {
config: RollbackConfig,
config_manager: Arc<RwLock<dyn ConfigurationManager>>,
health_monitor: Arc<RwLock<dyn HealthMonitor>>,
history: Arc<RwLock<VecDeque<RollbackHistoryEntry>>>,
active_rollbacks: Arc<RwLock<HashMap<String, RollbackStatus>>>,
snapshots: Arc<RwLock<HashMap<String, ConfigurationSnapshot>>>,
stats: Arc<RwLock<RollbackStats>>,
notification_handlers: Arc<RwLock<Vec<Box<dyn NotificationHandler>>>>,
}
impl RollbackEngine {
pub fn new(
config: RollbackConfig,
config_manager: Arc<RwLock<dyn ConfigurationManager>>,
health_monitor: Arc<RwLock<dyn HealthMonitor>>,
) -> Self {
Self {
config,
config_manager,
health_monitor,
history: Arc::new(RwLock::new(VecDeque::new())),
active_rollbacks: Arc::new(RwLock::new(HashMap::new())),
snapshots: Arc::new(RwLock::new(HashMap::new())),
stats: Arc::new(RwLock::new(RollbackStats::default())),
notification_handlers: Arc::new(RwLock::new(Vec::new())),
}
}
pub async fn should_trigger_rollback(
&self,
deployment_id: &str,
status: &DeploymentStatus,
metrics: &SystemMetrics,
) -> ActuatorResult<Option<RollbackTrigger>> {
if !self.config.auto_rollback {
return Ok(None);
}
if !status.health.healthy {
return Ok(Some(RollbackTrigger {
trigger_type: RollbackReason::HealthCheckFailed,
timestamp: Utc::now(),
automatic: true,
metrics: Some(status.metrics.clone()),
health_status: Some(status.health.clone()),
threshold_exceeded: None,
}));
}
if metrics.error_rate_pct > 5.0 {
return Ok(Some(RollbackTrigger {
trigger_type: RollbackReason::HighErrorRate,
timestamp: Utc::now(),
automatic: true,
metrics: Some(status.metrics.clone()),
health_status: Some(status.health.clone()),
threshold_exceeded: Some(ThresholdViolation {
metric: "error_rate_pct".to_string(),
actual_value: metrics.error_rate_pct,
threshold_value: 5.0,
operator: ">".to_string(),
}),
}));
}
if metrics.p99_latency_ms > 2000.0 {
return Ok(Some(RollbackTrigger {
trigger_type: RollbackReason::HighLatency,
timestamp: Utc::now(),
automatic: true,
metrics: Some(status.metrics.clone()),
health_status: Some(status.health.clone()),
threshold_exceeded: Some(ThresholdViolation {
metric: "p99_latency_ms".to_string(),
actual_value: metrics.p99_latency_ms,
threshold_value: 2000.0,
operator: ">".to_string(),
}),
}));
}
if let Some(avg_rating) = metrics.avg_rating {
if avg_rating < 3.5 {
return Ok(Some(RollbackTrigger {
trigger_type: RollbackReason::QualityDegradation,
timestamp: Utc::now(),
automatic: true,
metrics: Some(status.metrics.clone()),
health_status: Some(status.health.clone()),
threshold_exceeded: Some(ThresholdViolation {
metric: "avg_rating".to_string(),
actual_value: avg_rating,
threshold_value: 3.5,
operator: "<".to_string(),
}),
}));
}
}
let cost_increase_pct = if status.metrics.control_cost_per_request > 0.0 {
((status.metrics.canary_cost_per_request - status.metrics.control_cost_per_request)
/ status.metrics.control_cost_per_request)
* 100.0
} else {
0.0
};
if cost_increase_pct > 20.0 {
return Ok(Some(RollbackTrigger {
trigger_type: RollbackReason::HighCost,
timestamp: Utc::now(),
automatic: true,
metrics: Some(status.metrics.clone()),
health_status: Some(status.health.clone()),
threshold_exceeded: Some(ThresholdViolation {
metric: "cost_increase_pct".to_string(),
actual_value: cost_increase_pct,
threshold_value: 20.0,
operator: ">".to_string(),
}),
}));
}
Ok(None)
}
pub async fn execute_rollback(
&self,
request: RollbackRequest,
deployment_status: &DeploymentStatus,
config_changes: Vec<ConfigChange>,
) -> ActuatorResult<RollbackResult> {
let start_time = Utc::now();
let rollback_id = format!("rollback-{}", uuid::Uuid::new_v4());
let mode = if self.config.fast_rollback || request.reason == RollbackReason::HealthCheckFailed
{
RollbackMode::Fast
} else {
RollbackMode::Gradual
};
let mut status = RollbackStatus {
rollback_id: rollback_id.clone(),
deployment_id: request.deployment_id.clone(),
mode,
current_phase: None,
start_time,
end_time: None,
completed: false,
success: false,
error: None,
phases: Vec::new(),
};
{
let mut active = self.active_rollbacks.write().await;
active.insert(request.deployment_id.clone(), status.clone());
}
if self.config.enable_notifications {
let notification = RollbackNotification {
id: format!("notif-{}", uuid::Uuid::new_v4()),
deployment_id: request.deployment_id.clone(),
rollback_id: rollback_id.clone(),
reason: request.reason.clone(),
timestamp: Utc::now(),
severity: if request.automatic {
NotificationSeverity::Warning
} else {
NotificationSeverity::Info
},
message: format!(
"Rollback initiated for deployment {} ({})",
request.deployment_id,
if request.automatic { "automatic" } else { "manual" }
),
details: request.metadata.clone(),
};
self.send_notification(notification).await?;
}
let result = match mode {
RollbackMode::Fast => {
self.execute_fast_rollback(&request, &config_changes, &mut status)
.await
}
RollbackMode::Gradual => {
self.execute_gradual_rollback(&request, &config_changes, deployment_status, &mut status)
.await
}
};
let end_time = Utc::now();
let duration = end_time.signed_duration_since(start_time);
status.end_time = Some(end_time);
status.completed = true;
let rollback_result = match result {
Ok(_) => {
status.success = true;
let validation_result = self.validate_after_rollback(&request.deployment_id).await;
let success = validation_result.unwrap_or(false);
if !success {
status.success = false;
status.error = Some("Post-rollback validation failed".to_string());
}
RollbackResult {
deployment_id: request.deployment_id.clone(),
success,
error: status.error.clone(),
reverted_changes: config_changes.clone(),
timestamp: end_time,
duration: Duration::from_secs(duration.num_seconds().max(0) as u64),
}
}
Err(e) => {
status.success = false;
status.error = Some(e.to_string());
RollbackResult {
deployment_id: request.deployment_id.clone(),
success: false,
error: Some(e.to_string()),
reverted_changes: Vec::new(),
timestamp: end_time,
duration: Duration::from_secs(duration.num_seconds().max(0) as u64),
}
}
};
self.update_stats(&rollback_result, mode, request.automatic).await;
self.store_history(
rollback_id.clone(),
request.clone(),
rollback_result.clone(),
status.clone(),
)
.await?;
{
let mut active = self.active_rollbacks.write().await;
active.remove(&request.deployment_id);
}
if self.config.enable_notifications {
let notification = RollbackNotification {
id: format!("notif-{}", uuid::Uuid::new_v4()),
deployment_id: request.deployment_id.clone(),
rollback_id,
reason: request.reason,
timestamp: Utc::now(),
severity: if rollback_result.success {
NotificationSeverity::Info
} else {
NotificationSeverity::Critical
},
message: format!(
"Rollback {} for deployment {}",
if rollback_result.success { "completed" } else { "failed" },
request.deployment_id
),
details: HashMap::new(),
};
self.send_notification(notification).await?;
}
Ok(rollback_result)
}
async fn execute_fast_rollback(
&self,
request: &RollbackRequest,
config_changes: &[ConfigChange],
status: &mut RollbackStatus,
) -> ActuatorResult<()> {
let mut config_manager = self.config_manager.write().await;
config_manager.revert_config(config_changes).await?;
Ok(())
}
async fn execute_gradual_rollback(
&self,
request: &RollbackRequest,
config_changes: &[ConfigChange],
deployment_status: &DeploymentStatus,
status: &mut RollbackStatus,
) -> ActuatorResult<()> {
let phases = if let Some(current_phase) = deployment_status.current_phase {
(0..=current_phase)
.rev()
.map(|phase| RollbackPhase {
phase,
traffic_percent: match phase {
0 => 0.0,
1 => 1.0,
2 => 5.0,
3 => 25.0,
4 => 50.0,
_ => 100.0,
},
start_time: None,
end_time: None,
completed: false,
})
.collect()
} else {
vec![RollbackPhase {
phase: 0,
traffic_percent: 0.0,
start_time: None,
end_time: None,
completed: false,
}]
};
status.phases = phases;
for (idx, phase) in status.phases.iter_mut().enumerate() {
status.current_phase = Some(idx);
phase.start_time = Some(Utc::now());
tokio::time::sleep(Duration::from_secs(10)).await;
phase.end_time = Some(Utc::now());
phase.completed = true;
}
let mut config_manager = self.config_manager.write().await;
config_manager.revert_config(config_changes).await?;
Ok(())
}
async fn validate_after_rollback(&self, deployment_id: &str) -> ActuatorResult<bool> {
tokio::time::sleep(Duration::from_secs(5)).await;
let health_monitor = self.health_monitor.read().await;
let health_summary = health_monitor.get_health_summary(deployment_id).await?;
Ok(health_summary.healthy)
}
pub async fn store_snapshot(
&self,
deployment_id: &str,
snapshot: ConfigurationSnapshot,
) -> ActuatorResult<()> {
if !self.config.enable_backup {
return Ok(());
}
let mut snapshots = self.snapshots.write().await;
snapshots.insert(deployment_id.to_string(), snapshot);
Ok(())
}
pub async fn restore_from_snapshot(&self, deployment_id: &str) -> ActuatorResult<()> {
let snapshot = {
let snapshots = self.snapshots.read().await;
snapshots
.get(deployment_id)
.cloned()
.ok_or_else(|| {
ActuatorError::RollbackFailed(format!(
"No snapshot found for deployment {}",
deployment_id
))
})?
};
let mut config_manager = self.config_manager.write().await;
config_manager.restore_snapshot(&snapshot.id).await?;
Ok(())
}
pub async fn get_history(&self, limit: Option<usize>) -> Vec<RollbackHistoryEntry> {
let history = self.history.read().await;
let limit = limit.unwrap_or(100);
history.iter().take(limit).cloned().collect()
}
pub async fn get_active_rollback(
&self,
deployment_id: &str,
) -> Option<RollbackStatus> {
let active = self.active_rollbacks.read().await;
active.get(deployment_id).cloned()
}
pub async fn get_stats(&self) -> RollbackStats {
self.stats.read().await.clone()
}
pub async fn add_notification_handler(
&self,
handler: Box<dyn NotificationHandler>,
) {
let mut handlers = self.notification_handlers.write().await;
handlers.push(handler);
}
async fn send_notification(&self, notification: RollbackNotification) -> ActuatorResult<()> {
let handlers = self.notification_handlers.read().await;
for handler in handlers.iter() {
handler.handle(notification.clone()).await?;
}
Ok(())
}
async fn store_history(
&self,
rollback_id: String,
request: RollbackRequest,
result: RollbackResult,
status: RollbackStatus,
) -> ActuatorResult<()> {
let trigger = RollbackTrigger {
trigger_type: request.reason.clone(),
timestamp: request.timestamp,
automatic: request.automatic,
metrics: None,
health_status: None,
threshold_exceeded: None,
};
let entry = RollbackHistoryEntry {
rollback_id,
deployment_id: request.deployment_id.clone(),
request,
result,
status,
trigger,
};
let mut history = self.history.write().await;
history.push_front(entry);
while history.len() > self.config.max_rollback_history {
history.pop_back();
}
Ok(())
}
async fn update_stats(&self, result: &RollbackResult, mode: RollbackMode, automatic: bool) {
let mut stats = self.stats.write().await;
stats.total_rollbacks += 1;
if automatic {
stats.automatic_rollbacks += 1;
} else {
stats.manual_rollbacks += 1;
}
if result.success {
stats.successful_rollbacks += 1;
} else {
stats.failed_rollbacks += 1;
}
match mode {
RollbackMode::Fast => stats.fast_rollbacks += 1,
RollbackMode::Gradual => stats.gradual_rollbacks += 1,
}
let duration_secs = result.duration.as_secs_f64();
if stats.total_rollbacks == 1 {
stats.avg_rollback_duration_secs = duration_secs;
stats.fastest_rollback_secs = duration_secs;
} else {
stats.avg_rollback_duration_secs = (stats.avg_rollback_duration_secs
* (stats.total_rollbacks - 1) as f64
+ duration_secs)
/ stats.total_rollbacks as f64;
if duration_secs < stats.fastest_rollback_secs {
stats.fastest_rollback_secs = duration_secs;
}
}
}
}
#[async_trait]
pub trait NotificationHandler: Send + Sync {
async fn handle(&self, notification: RollbackNotification) -> ActuatorResult<()>;
}
#[cfg(test)]
pub struct MockConfigManager;
#[cfg(test)]
#[async_trait]
impl ConfigurationManager for MockConfigManager {
async fn apply_config(&mut self, _changes: &[ConfigChange]) -> ActuatorResult<()> {
Ok(())
}
async fn revert_config(&mut self, _changes: &[ConfigChange]) -> ActuatorResult<()> {
Ok(())
}
async fn create_snapshot(&mut self) -> ActuatorResult<ConfigurationSnapshot> {
Ok(ConfigurationSnapshot {
id: "test-snapshot".to_string(),
timestamp: Utc::now(),
config: HashMap::new(),
deployment_id: None,
version: 1,
tags: Vec::new(),
})
}
async fn restore_snapshot(&mut self, _snapshot_id: &str) -> ActuatorResult<()> {
Ok(())
}
async fn get_config(&self) -> ActuatorResult<HashMap<String, serde_json::Value>> {
Ok(HashMap::new())
}
async fn validate_config(&self, _changes: &[ConfigChange]) -> ActuatorResult<()> {
Ok(())
}
}
#[cfg(test)]
pub struct MockHealthMonitor {
pub healthy: bool,
}
#[cfg(test)]
#[async_trait]
impl HealthMonitor for MockHealthMonitor {
async fn check_health(
&self,
_deployment_id: &str,
_metrics: &SystemMetrics,
) -> ActuatorResult<Vec<super::types::HealthCheckResult>> {
Ok(Vec::new())
}
async fn should_rollback(
&self,
_deployment_id: &str,
_status: &DeploymentStatus,
) -> ActuatorResult<bool> {
Ok(!self.healthy)
}
async fn get_health_summary(
&self,
_deployment_id: &str,
) -> ActuatorResult<HealthStatus> {
Ok(HealthStatus {
healthy: self.healthy,
checks_passed: if self.healthy { 5 } else { 0 },
checks_failed: if self.healthy { 0 } else { 5 },
check_results: Vec::new(),
last_check_time: Utc::now(),
})
}
}
#[cfg(test)]
mod tests {
use super::*;
fn create_test_engine() -> RollbackEngine {
let config = RollbackConfig::default();
let config_manager = Arc::new(RwLock::new(MockConfigManager));
let health_monitor = Arc::new(RwLock::new(MockHealthMonitor { healthy: true }));
RollbackEngine::new(config, config_manager, health_monitor)
}
fn create_test_deployment_status() -> DeploymentStatus {
DeploymentStatus {
deployment_id: "test-deployment".to_string(),
state: DeploymentState::RollingOut,
current_phase: Some(2),
traffic_percent: 25.0,
start_time: Utc::now(),
phase_start_time: Some(Utc::now()),
health: HealthStatus {
healthy: true,
checks_passed: 5,
checks_failed: 0,
check_results: Vec::new(),
last_check_time: Utc::now(),
},
metrics: DeploymentMetrics::default(),
errors: Vec::new(),
last_updated: Utc::now(),
}
}
fn create_test_metrics() -> SystemMetrics {
SystemMetrics {
avg_latency_ms: 100.0,
p95_latency_ms: 200.0,
p99_latency_ms: 300.0,
success_rate_pct: 99.5,
error_rate_pct: 0.5,
throughput_rps: 1000.0,
total_cost_usd: 100.0,
daily_cost_usd: 10.0,
avg_rating: Some(4.5),
cache_hit_rate_pct: Some(80.0),
request_count: 10000,
timestamp: Utc::now(),
}
}
#[tokio::test]
async fn test_no_rollback_trigger_when_healthy() {
let engine = create_test_engine();
let status = create_test_deployment_status();
let metrics = create_test_metrics();
let trigger = engine
.should_trigger_rollback("test-deployment", &status, &metrics)
.await
.unwrap();
assert!(trigger.is_none());
}
#[tokio::test]
async fn test_rollback_trigger_on_high_error_rate() {
let engine = create_test_engine();
let status = create_test_deployment_status();
let mut metrics = create_test_metrics();
metrics.error_rate_pct = 10.0;
let trigger = engine
.should_trigger_rollback("test-deployment", &status, &metrics)
.await
.unwrap();
assert!(trigger.is_some());
let trigger = trigger.unwrap();
assert_eq!(trigger.trigger_type, RollbackReason::HighErrorRate);
assert!(trigger.automatic);
}
#[tokio::test]
async fn test_rollback_trigger_on_high_latency() {
let engine = create_test_engine();
let status = create_test_deployment_status();
let mut metrics = create_test_metrics();
metrics.p99_latency_ms = 3000.0;
let trigger = engine
.should_trigger_rollback("test-deployment", &status, &metrics)
.await
.unwrap();
assert!(trigger.is_some());
let trigger = trigger.unwrap();
assert_eq!(trigger.trigger_type, RollbackReason::HighLatency);
}
#[tokio::test]
async fn test_rollback_trigger_on_quality_degradation() {
let engine = create_test_engine();
let status = create_test_deployment_status();
let mut metrics = create_test_metrics();
metrics.avg_rating = Some(2.5);
let trigger = engine
.should_trigger_rollback("test-deployment", &status, &metrics)
.await
.unwrap();
assert!(trigger.is_some());
let trigger = trigger.unwrap();
assert_eq!(trigger.trigger_type, RollbackReason::QualityDegradation);
}
#[tokio::test]
async fn test_rollback_trigger_on_health_check_failure() {
let engine = create_test_engine();
let mut status = create_test_deployment_status();
status.health.healthy = false;
let metrics = create_test_metrics();
let trigger = engine
.should_trigger_rollback("test-deployment", &status, &metrics)
.await
.unwrap();
assert!(trigger.is_some());
let trigger = trigger.unwrap();
assert_eq!(trigger.trigger_type, RollbackReason::HealthCheckFailed);
}
#[tokio::test]
async fn test_fast_rollback_execution() {
let engine = create_test_engine();
let status = create_test_deployment_status();
let request = RollbackRequest {
deployment_id: "test-deployment".to_string(),
reason: RollbackReason::HealthCheckFailed,
timestamp: Utc::now(),
automatic: true,
metadata: HashMap::new(),
};
let config_changes = vec![ConfigChange {
config_type: crate::decision::ConfigType::Model,
path: "model.name".to_string(),
old_value: Some(serde_json::json!("gpt-3.5-turbo")),
new_value: serde_json::json!("gpt-4"),
description: "Model upgrade".to_string(),
}];
let result = engine
.execute_rollback(request, &status, config_changes.clone())
.await
.unwrap();
assert!(result.success);
assert_eq!(result.reverted_changes.len(), config_changes.len());
}
#[tokio::test]
async fn test_gradual_rollback_execution() {
let mut config = RollbackConfig::default();
config.fast_rollback = false;
let config_manager = Arc::new(RwLock::new(MockConfigManager));
let health_monitor = Arc::new(RwLock::new(MockHealthMonitor { healthy: true }));
let engine = RollbackEngine::new(config, config_manager, health_monitor);
let status = create_test_deployment_status();
let request = RollbackRequest {
deployment_id: "test-deployment".to_string(),
reason: RollbackReason::Manual,
timestamp: Utc::now(),
automatic: false,
metadata: HashMap::new(),
};
let config_changes = vec![ConfigChange {
config_type: crate::decision::ConfigType::Model,
path: "model.name".to_string(),
old_value: Some(serde_json::json!("gpt-3.5-turbo")),
new_value: serde_json::json!("gpt-4"),
description: "Model upgrade".to_string(),
}];
let result = engine
.execute_rollback(request, &status, config_changes)
.await
.unwrap();
assert!(result.success);
}
#[tokio::test]
async fn test_snapshot_storage_and_restoration() {
let engine = create_test_engine();
let snapshot = ConfigurationSnapshot {
id: "snapshot-1".to_string(),
timestamp: Utc::now(),
config: HashMap::new(),
deployment_id: Some("test-deployment".to_string()),
version: 1,
tags: vec!["production".to_string()],
};
engine
.store_snapshot("test-deployment", snapshot)
.await
.unwrap();
let result = engine.restore_from_snapshot("test-deployment").await;
assert!(result.is_ok());
}
#[tokio::test]
async fn test_rollback_history_tracking() {
let engine = create_test_engine();
let status = create_test_deployment_status();
let request = RollbackRequest {
deployment_id: "test-deployment".to_string(),
reason: RollbackReason::Manual,
timestamp: Utc::now(),
automatic: false,
metadata: HashMap::new(),
};
let config_changes = vec![ConfigChange {
config_type: crate::decision::ConfigType::Model,
path: "model.name".to_string(),
old_value: Some(serde_json::json!("gpt-3.5-turbo")),
new_value: serde_json::json!("gpt-4"),
description: "Model upgrade".to_string(),
}];
engine
.execute_rollback(request, &status, config_changes)
.await
.unwrap();
let history = engine.get_history(None).await;
assert_eq!(history.len(), 1);
assert_eq!(history[0].deployment_id, "test-deployment");
}
#[tokio::test]
async fn test_rollback_statistics() {
let engine = create_test_engine();
let status = create_test_deployment_status();
for i in 0..3 {
let request = RollbackRequest {
deployment_id: format!("test-deployment-{}", i),
reason: if i % 2 == 0 {
RollbackReason::Manual
} else {
RollbackReason::HighErrorRate
},
timestamp: Utc::now(),
automatic: i % 2 != 0,
metadata: HashMap::new(),
};
let config_changes = vec![ConfigChange {
config_type: crate::decision::ConfigType::Model,
path: "model.name".to_string(),
old_value: Some(serde_json::json!("gpt-3.5-turbo")),
new_value: serde_json::json!("gpt-4"),
description: "Model upgrade".to_string(),
}];
engine
.execute_rollback(request, &status, config_changes)
.await
.unwrap();
}
let stats = engine.get_stats().await;
assert_eq!(stats.total_rollbacks, 3);
assert_eq!(stats.manual_rollbacks, 2);
assert_eq!(stats.automatic_rollbacks, 1);
assert_eq!(stats.successful_rollbacks, 3);
}
#[tokio::test]
async fn test_active_rollback_tracking() {
let engine = create_test_engine();
let status = create_test_deployment_status();
let request = RollbackRequest {
deployment_id: "test-deployment".to_string(),
reason: RollbackReason::Manual,
timestamp: Utc::now(),
automatic: false,
metadata: HashMap::new(),
};
let config_changes = vec![ConfigChange {
config_type: crate::decision::ConfigType::Model,
path: "model.name".to_string(),
old_value: Some(serde_json::json!("gpt-3.5-turbo")),
new_value: serde_json::json!("gpt-4"),
description: "Model upgrade".to_string(),
}];
engine
.execute_rollback(request, &status, config_changes)
.await
.unwrap();
let active = engine.get_active_rollback("test-deployment").await;
assert!(active.is_none());
}
#[tokio::test]
async fn test_post_rollback_validation() {
let config_manager = Arc::new(RwLock::new(MockConfigManager));
let health_monitor = Arc::new(RwLock::new(MockHealthMonitor { healthy: true }));
let engine = RollbackEngine::new(
RollbackConfig::default(),
config_manager,
health_monitor,
);
let result = engine.validate_after_rollback("test-deployment").await;
assert!(result.is_ok());
assert!(result.unwrap());
}
#[tokio::test]
async fn test_post_rollback_validation_failure() {
let config_manager = Arc::new(RwLock::new(MockConfigManager));
let health_monitor = Arc::new(RwLock::new(MockHealthMonitor { healthy: false }));
let engine = RollbackEngine::new(
RollbackConfig::default(),
config_manager,
health_monitor,
);
let result = engine.validate_after_rollback("test-deployment").await;
assert!(result.is_ok());
assert!(!result.unwrap());
}
}