mockforge_vbr/
snapshots.rs

1//! State snapshot and reset functionality
2//!
3//! This module provides functionality to create, restore, and manage database snapshots
4//! for point-in-time recovery and environment state management.
5
6use crate::database::VirtualDatabase;
7use crate::entities::EntityRegistry;
8use crate::{Error, Result};
9use serde::{Deserialize, Serialize};
10use serde_json::Value;
11use std::collections::HashMap;
12use std::path::{Path, PathBuf};
13use tokio::fs;
14
15/// Snapshot metadata
16#[derive(Debug, Clone, Serialize, Deserialize)]
17pub struct SnapshotMetadata {
18    /// Snapshot name
19    pub name: String,
20    /// Timestamp when snapshot was created
21    pub created_at: chrono::DateTime<chrono::Utc>,
22    /// Optional description
23    pub description: Option<String>,
24    /// Entity counts in the snapshot
25    pub entity_counts: HashMap<String, usize>,
26    /// Database size in bytes (if available)
27    pub database_size: Option<u64>,
28    /// Storage backend type
29    pub storage_backend: String,
30    /// Time travel state (if included in snapshot)
31    #[serde(default)]
32    pub time_travel_state: Option<TimeTravelSnapshotState>,
33}
34
35/// Time travel state included in snapshots
36#[derive(Debug, Clone, Serialize, Deserialize)]
37pub struct TimeTravelSnapshotState {
38    /// Time travel enabled status
39    pub enabled: bool,
40    /// Current virtual time (if enabled)
41    pub current_time: Option<chrono::DateTime<chrono::Utc>>,
42    /// Time scale factor
43    pub scale_factor: f64,
44    /// Cron jobs (serialized)
45    #[serde(default)]
46    pub cron_jobs: Vec<serde_json::Value>,
47    /// Mutation rules (serialized)
48    #[serde(default)]
49    pub mutation_rules: Vec<serde_json::Value>,
50}
51
52/// Snapshot manager
53pub struct SnapshotManager {
54    /// Base directory for storing snapshots
55    snapshots_dir: PathBuf,
56}
57
58impl SnapshotManager {
59    /// Create a new snapshot manager
60    ///
61    /// # Arguments
62    /// * `snapshots_dir` - Base directory for storing snapshots
63    pub fn new<P: AsRef<Path>>(snapshots_dir: P) -> Self {
64        Self {
65            snapshots_dir: snapshots_dir.as_ref().to_path_buf(),
66        }
67    }
68
69    /// Get the path for a snapshot directory
70    fn snapshot_path(&self, name: &str) -> PathBuf {
71        self.snapshots_dir.join(name)
72    }
73
74    /// Get the path for snapshot metadata file
75    fn metadata_path(&self, name: &str) -> PathBuf {
76        self.snapshot_path(name).join("metadata.json")
77    }
78
79    /// Create a snapshot of the current database state
80    ///
81    /// # Arguments
82    /// * `name` - Name for the snapshot
83    /// * `description` - Optional description
84    /// * `database` - The virtual database instance
85    /// * `registry` - The entity registry
86    /// * `include_time_travel` - Whether to include time travel state (cron jobs, mutation rules)
87    /// * `time_travel_state` - Optional time travel state to include
88    pub async fn create_snapshot(
89        &self,
90        name: &str,
91        description: Option<String>,
92        database: &dyn crate::database::VirtualDatabase,
93        registry: &EntityRegistry,
94    ) -> Result<SnapshotMetadata> {
95        self.create_snapshot_with_time_travel(name, description, database, registry, false, None)
96            .await
97    }
98
99    /// Create a snapshot with optional time travel state
100    ///
101    /// # Arguments
102    /// * `name` - Name for the snapshot
103    /// * `description` - Optional description
104    /// * `database` - The virtual database instance
105    /// * `registry` - The entity registry
106    /// * `include_time_travel` - Whether to include time travel state
107    /// * `time_travel_state` - Optional time travel state to include
108    pub async fn create_snapshot_with_time_travel(
109        &self,
110        name: &str,
111        description: Option<String>,
112        database: &dyn crate::database::VirtualDatabase,
113        registry: &EntityRegistry,
114        include_time_travel: bool,
115        time_travel_state: Option<TimeTravelSnapshotState>,
116    ) -> Result<SnapshotMetadata> {
117        // Create snapshot directory
118        let snapshot_dir = self.snapshot_path(name);
119        fs::create_dir_all(&snapshot_dir)
120            .await
121            .map_err(|e| Error::generic(format!("Failed to create snapshot directory: {}", e)))?;
122
123        // Get entity counts
124        let mut entity_counts = HashMap::new();
125        for entity_name in registry.list() {
126            if let Some(entity) = registry.get(&entity_name) {
127                let table_name = entity.table_name();
128                let count_query = format!("SELECT COUNT(*) as count FROM {}", table_name);
129                let results = database.query(&count_query, &[]).await?;
130                let count = results
131                    .first()
132                    .and_then(|r| r.get("count"))
133                    .and_then(|v| v.as_u64())
134                    .map(|v| v as usize)
135                    .unwrap_or(0);
136                entity_counts.insert(entity_name.clone(), count);
137            }
138        }
139
140        // Create snapshot based on storage backend
141        let storage_backend = database.connection_info();
142        let database_size = self.create_snapshot_data(name, database, registry).await?;
143
144        // Create metadata
145        let metadata = SnapshotMetadata {
146            name: name.to_string(),
147            created_at: chrono::Utc::now(),
148            description,
149            time_travel_state: if include_time_travel {
150                time_travel_state
151            } else {
152                None
153            },
154            entity_counts,
155            database_size,
156            storage_backend,
157        };
158
159        // Save metadata
160        let metadata_json = serde_json::to_string_pretty(&metadata)
161            .map_err(|e| Error::generic(format!("Failed to serialize metadata: {}", e)))?;
162        fs::write(self.metadata_path(name), metadata_json)
163            .await
164            .map_err(|e| Error::generic(format!("Failed to write metadata: {}", e)))?;
165
166        Ok(metadata)
167    }
168
169    /// Create snapshot data based on storage backend
170    async fn create_snapshot_data(
171        &self,
172        name: &str,
173        database: &dyn crate::database::VirtualDatabase,
174        registry: &EntityRegistry,
175    ) -> Result<Option<u64>> {
176        let snapshot_dir = self.snapshot_path(name);
177        let storage_backend = database.connection_info().to_lowercase();
178
179        if storage_backend.contains("sqlite") {
180            // For SQLite, we'll export all data to JSON
181            // In a production system, you might use SQLite backup API
182            self.export_sqlite_to_json(&snapshot_dir, database, registry).await?;
183            Ok(None) // Size calculation would require file system access
184        } else if storage_backend.contains("json") {
185            // For JSON backend, copy the JSON file
186            // This would require access to the JSON database implementation
187            Ok(None)
188        } else {
189            // For in-memory, export to JSON
190            self.export_memory_to_json(&snapshot_dir, database, registry).await?;
191            Ok(None)
192        }
193    }
194
195    /// Export SQLite database to JSON files
196    async fn export_sqlite_to_json(
197        &self,
198        snapshot_dir: &Path,
199        database: &dyn crate::database::VirtualDatabase,
200        registry: &EntityRegistry,
201    ) -> Result<()> {
202        let data_dir = snapshot_dir.join("data");
203        fs::create_dir_all(&data_dir)
204            .await
205            .map_err(|e| Error::generic(format!("Failed to create data directory: {}", e)))?;
206
207        // Export each entity table to JSON
208        for entity_name in registry.list() {
209            if let Some(entity) = registry.get(&entity_name) {
210                let table_name = entity.table_name();
211                let query = format!("SELECT * FROM {}", table_name);
212                let records = database.query(&query, &[]).await?;
213
214                let json_file = data_dir.join(format!("{}.json", entity_name.to_lowercase()));
215                let json_content = serde_json::to_string_pretty(&records)
216                    .map_err(|e| Error::generic(format!("Failed to serialize data: {}", e)))?;
217                fs::write(&json_file, json_content)
218                    .await
219                    .map_err(|e| Error::generic(format!("Failed to write data file: {}", e)))?;
220            }
221        }
222
223        Ok(())
224    }
225
226    /// Export in-memory database to JSON files
227    async fn export_memory_to_json(
228        &self,
229        snapshot_dir: &Path,
230        database: &dyn crate::database::VirtualDatabase,
231        registry: &EntityRegistry,
232    ) -> Result<()> {
233        self.export_sqlite_to_json(snapshot_dir, database, registry).await
234    }
235
236    /// Restore a snapshot
237    ///
238    /// # Arguments
239    /// * `name` - Name of the snapshot to restore
240    /// * `database` - The virtual database instance (will be reset)
241    /// * `registry` - The entity registry
242    pub async fn restore_snapshot(
243        &self,
244        name: &str,
245        database: &dyn crate::database::VirtualDatabase,
246        registry: &EntityRegistry,
247    ) -> Result<()> {
248        self.restore_snapshot_with_time_travel(
249            name,
250            database,
251            registry,
252            false,
253            None::<fn(TimeTravelSnapshotState) -> Result<()>>,
254        )
255        .await
256    }
257
258    /// Restore a snapshot with optional time travel state restoration
259    ///
260    /// # Arguments
261    /// * `name` - Name of the snapshot to restore
262    /// * `database` - The virtual database instance (will be reset)
263    /// * `registry` - The entity registry
264    /// * `restore_time_travel` - Whether to restore time travel state
265    /// * `time_travel_restore_callback` - Optional callback to restore time travel state
266    pub async fn restore_snapshot_with_time_travel<F>(
267        &self,
268        name: &str,
269        database: &dyn crate::database::VirtualDatabase,
270        registry: &EntityRegistry,
271        restore_time_travel: bool,
272        time_travel_restore_callback: Option<F>,
273    ) -> Result<()>
274    where
275        F: FnOnce(TimeTravelSnapshotState) -> Result<()>,
276    {
277        // Load metadata to verify snapshot exists
278        let metadata = self.get_snapshot_metadata(name).await?;
279
280        // Clear existing data
281        self.reset_database(database, registry).await?;
282
283        // Restore data based on storage backend
284        let snapshot_dir = self.snapshot_path(name);
285        let storage_backend_lower = metadata.storage_backend.to_lowercase();
286        if storage_backend_lower.contains("sqlite") || storage_backend_lower.contains("memory") {
287            self.import_json_to_database(&snapshot_dir, database, registry).await?;
288        } else if storage_backend_lower.contains("json") {
289            // For JSON backend, would need to copy the JSON file
290            // This requires access to the JSON database implementation
291            return Err(Error::generic(
292                "JSON backend snapshot restore not yet implemented".to_string(),
293            ));
294        }
295
296        // Restore time travel state if requested and available
297        if restore_time_travel {
298            if let Some(ref time_travel_state) = metadata.time_travel_state {
299                if let Some(callback) = time_travel_restore_callback {
300                    callback(time_travel_state.clone())?;
301                }
302            }
303        }
304
305        Ok(())
306    }
307
308    /// Import JSON data into database
309    async fn import_json_to_database(
310        &self,
311        snapshot_dir: &Path,
312        database: &dyn crate::database::VirtualDatabase,
313        registry: &EntityRegistry,
314    ) -> Result<()> {
315        let data_dir = snapshot_dir.join("data");
316
317        if !data_dir.exists() {
318            return Err(Error::generic("Snapshot data directory not found".to_string()));
319        }
320
321        // Import each entity
322        for entity_name in registry.list() {
323            let json_file = data_dir.join(format!("{}.json", entity_name.to_lowercase()));
324            if !json_file.exists() {
325                // Skip if file doesn't exist (entity had no data in snapshot)
326                continue;
327            }
328
329            {
330                let content = fs::read_to_string(&json_file)
331                    .await
332                    .map_err(|e| Error::generic(format!("Failed to read data file: {}", e)))?;
333
334                let records: Vec<HashMap<String, Value>> = serde_json::from_str(&content)
335                    .map_err(|e| Error::generic(format!("Failed to parse data file: {}", e)))?;
336
337                if let Some(entity) = registry.get(&entity_name) {
338                    let table_name = entity.table_name();
339
340                    // Ensure table exists before inserting
341                    // For Memory database, this is handled by execute, but we need to make sure
342                    // the table structure is preserved after reset
343                    if !database.table_exists(table_name).await.unwrap_or(false) {
344                        // Table was removed during reset, we need to recreate it
345                        // For Memory database, this happens automatically on first INSERT
346                        // But we should ensure the table entry exists
347                    }
348
349                    for record in records {
350                        let fields: Vec<String> = record.keys().cloned().collect();
351                        let placeholders: Vec<String> =
352                            (0..fields.len()).map(|_| "?".to_string()).collect();
353                        let values: Vec<Value> = fields
354                            .iter()
355                            .map(|f| record.get(f).cloned().unwrap_or(Value::Null))
356                            .collect();
357
358                        let query = format!(
359                            "INSERT INTO {} ({}) VALUES ({})",
360                            table_name,
361                            fields.join(", "),
362                            placeholders.join(", ")
363                        );
364
365                        database.execute(&query, &values).await?;
366                    }
367                }
368            }
369        }
370
371        Ok(())
372    }
373
374    /// Reset database to empty state
375    async fn reset_database(
376        &self,
377        database: &dyn crate::database::VirtualDatabase,
378        registry: &EntityRegistry,
379    ) -> Result<()> {
380        // Delete all data from all tables
381        for entity_name in registry.list() {
382            if let Some(entity) = registry.get(&entity_name) {
383                let table_name = entity.table_name();
384                let query = format!("DELETE FROM {}", table_name);
385                let _ = database.execute(&query, &[]).await;
386            }
387        }
388
389        // Reset counters
390        let counter_table = "_vbr_counters";
391        if database.table_exists(counter_table).await.unwrap_or(false) {
392            let query = format!("DELETE FROM {}", counter_table);
393            let _ = database.execute(&query, &[]).await;
394        }
395
396        Ok(())
397    }
398
399    /// List all snapshots
400    pub async fn list_snapshots(&self) -> Result<Vec<SnapshotMetadata>> {
401        if !self.snapshots_dir.exists() {
402            return Ok(Vec::new());
403        }
404
405        let mut snapshots = Vec::new();
406
407        let mut entries = fs::read_dir(&self.snapshots_dir)
408            .await
409            .map_err(|e| Error::generic(format!("Failed to read snapshots directory: {}", e)))?;
410
411        while let Some(entry) = entries
412            .next_entry()
413            .await
414            .map_err(|e| Error::generic(format!("Failed to read directory entry: {}", e)))?
415        {
416            let path = entry.path();
417            if path.is_dir() {
418                if let Some(name) = path.file_name().and_then(|n| n.to_str()) {
419                    if let Ok(metadata) = self.get_snapshot_metadata(name).await {
420                        snapshots.push(metadata);
421                    }
422                }
423            }
424        }
425
426        // Sort by creation time (newest first)
427        snapshots.sort_by(|a, b| b.created_at.cmp(&a.created_at));
428
429        Ok(snapshots)
430    }
431
432    /// Get snapshot metadata
433    pub async fn get_snapshot_metadata(&self, name: &str) -> Result<SnapshotMetadata> {
434        let metadata_path = self.metadata_path(name);
435        let content = fs::read_to_string(&metadata_path)
436            .await
437            .map_err(|e| Error::generic(format!("Failed to read snapshot metadata: {}", e)))?;
438
439        let metadata: SnapshotMetadata = serde_json::from_str(&content)
440            .map_err(|e| Error::generic(format!("Failed to parse snapshot metadata: {}", e)))?;
441
442        Ok(metadata)
443    }
444
445    /// Delete a snapshot
446    pub async fn delete_snapshot(&self, name: &str) -> Result<()> {
447        let snapshot_dir = self.snapshot_path(name);
448
449        if !snapshot_dir.exists() {
450            return Err(Error::generic(format!("Snapshot '{}' not found", name)));
451        }
452
453        fs::remove_dir_all(&snapshot_dir)
454            .await
455            .map_err(|e| Error::generic(format!("Failed to delete snapshot: {}", e)))?;
456
457        Ok(())
458    }
459}
460
461/// Reset database to empty state (public API)
462pub async fn reset_database(
463    database: &dyn crate::database::VirtualDatabase,
464    registry: &EntityRegistry,
465) -> Result<()> {
466    // This is a simplified reset - in production, you might want to
467    // drop and recreate tables, but for now we'll just delete all data
468    for entity_name in registry.list() {
469        if let Some(entity) = registry.get(&entity_name) {
470            let table_name = entity.table_name();
471            let query = format!("DELETE FROM {}", table_name);
472            let _ = database.execute(&query, &[]).await;
473        }
474    }
475
476    // Reset counters
477    let counter_table = "_vbr_counters";
478    if database.table_exists(counter_table).await.unwrap_or(false) {
479        let query = format!("DELETE FROM {}", counter_table);
480        let _ = database.execute(&query, &[]).await;
481    }
482
483    Ok(())
484}
485
486#[cfg(test)]
487mod tests {
488    use super::*;
489    use crate::database::InMemoryDatabase;
490    use crate::entities::{Entity, EntityRegistry};
491    use crate::migration::MigrationManager;
492    use crate::schema::VbrSchemaDefinition;
493    use mockforge_data::{FieldDefinition, SchemaDefinition};
494    use std::sync::Arc;
495
496    async fn setup_test_env(
497    ) -> (Arc<dyn VirtualDatabase + Send + Sync>, EntityRegistry, tempfile::TempDir) {
498        let mut db = InMemoryDatabase::new().await.unwrap();
499        db.initialize().await.unwrap();
500        let mut registry = EntityRegistry::new();
501
502        // Create a test entity
503        let base_schema = SchemaDefinition::new("User".to_string())
504            .with_field(FieldDefinition::new("id".to_string(), "string".to_string()))
505            .with_field(FieldDefinition::new("name".to_string(), "string".to_string()));
506
507        let vbr_schema = VbrSchemaDefinition::new(base_schema);
508        let entity = Entity::new("User".to_string(), vbr_schema);
509
510        let manager = MigrationManager::new();
511        let create_sql = manager.generate_create_table(&entity).unwrap();
512        db.create_table(&create_sql).await.unwrap();
513
514        registry.register(entity).unwrap();
515
516        let temp_dir = tempfile::tempdir().unwrap();
517        (Arc::new(db), registry, temp_dir)
518    }
519
520    // TimeTravelSnapshotState tests
521    #[test]
522    fn test_time_travel_snapshot_state_serialize() {
523        let state = TimeTravelSnapshotState {
524            enabled: true,
525            current_time: Some(chrono::Utc::now()),
526            scale_factor: 1.0,
527            cron_jobs: vec![],
528            mutation_rules: vec![],
529        };
530
531        let json = serde_json::to_string(&state).unwrap();
532        assert!(json.contains("enabled"));
533        assert!(json.contains("scale_factor"));
534    }
535
536    #[test]
537    fn test_time_travel_snapshot_state_deserialize() {
538        let json = r#"{
539            "enabled": true,
540            "current_time": "2024-01-01T00:00:00Z",
541            "scale_factor": 2.0,
542            "cron_jobs": [],
543            "mutation_rules": []
544        }"#;
545
546        let state: TimeTravelSnapshotState = serde_json::from_str(json).unwrap();
547        assert!(state.enabled);
548        assert_eq!(state.scale_factor, 2.0);
549    }
550
551    #[test]
552    fn test_time_travel_snapshot_state_clone() {
553        let state = TimeTravelSnapshotState {
554            enabled: false,
555            current_time: None,
556            scale_factor: 1.5,
557            cron_jobs: vec![],
558            mutation_rules: vec![],
559        };
560
561        let cloned = state.clone();
562        assert_eq!(state.enabled, cloned.enabled);
563        assert_eq!(state.scale_factor, cloned.scale_factor);
564    }
565
566    // SnapshotMetadata tests
567    #[test]
568    fn test_snapshot_metadata_serialize() {
569        let metadata = SnapshotMetadata {
570            name: "test-snapshot".to_string(),
571            created_at: chrono::Utc::now(),
572            description: Some("Test description".to_string()),
573            entity_counts: HashMap::new(),
574            database_size: Some(1024),
575            storage_backend: "In-Memory".to_string(),
576            time_travel_state: None,
577        };
578
579        let json = serde_json::to_string(&metadata).unwrap();
580        assert!(json.contains("test-snapshot"));
581        assert!(json.contains("In-Memory"));
582    }
583
584    #[test]
585    fn test_snapshot_metadata_deserialize() {
586        let json = r#"{
587            "name": "test",
588            "created_at": "2024-01-01T00:00:00Z",
589            "description": null,
590            "entity_counts": {},
591            "database_size": null,
592            "storage_backend": "SQLite",
593            "time_travel_state": null
594        }"#;
595
596        let metadata: SnapshotMetadata = serde_json::from_str(json).unwrap();
597        assert_eq!(metadata.name, "test");
598        assert_eq!(metadata.storage_backend, "SQLite");
599    }
600
601    #[test]
602    fn test_snapshot_metadata_clone() {
603        let metadata = SnapshotMetadata {
604            name: "snap1".to_string(),
605            created_at: chrono::Utc::now(),
606            description: None,
607            entity_counts: HashMap::new(),
608            database_size: None,
609            storage_backend: "Memory".to_string(),
610            time_travel_state: None,
611        };
612
613        let cloned = metadata.clone();
614        assert_eq!(metadata.name, cloned.name);
615        assert_eq!(metadata.storage_backend, cloned.storage_backend);
616    }
617
618    // SnapshotManager tests
619    #[test]
620    fn test_snapshot_manager_new() {
621        let temp_dir = tempfile::tempdir().unwrap();
622        let manager = SnapshotManager::new(temp_dir.path());
623        assert_eq!(manager.snapshots_dir, temp_dir.path());
624    }
625
626    #[test]
627    fn test_snapshot_path() {
628        let temp_dir = tempfile::tempdir().unwrap();
629        let manager = SnapshotManager::new(temp_dir.path());
630        let path = manager.snapshot_path("test-snapshot");
631        assert!(path.ends_with("test-snapshot"));
632    }
633
634    #[test]
635    fn test_metadata_path() {
636        let temp_dir = tempfile::tempdir().unwrap();
637        let manager = SnapshotManager::new(temp_dir.path());
638        let path = manager.metadata_path("test-snapshot");
639        assert!(path.ends_with("metadata.json"));
640    }
641
642    #[tokio::test]
643    async fn test_create_snapshot() {
644        let (database, registry, temp_dir) = setup_test_env().await;
645        let manager = SnapshotManager::new(temp_dir.path());
646
647        // Insert test data
648        database
649            .execute(
650                "INSERT INTO users (id, name) VALUES (?, ?)",
651                &[
652                    serde_json::Value::String("1".to_string()),
653                    serde_json::Value::String("Test User".to_string()),
654                ],
655            )
656            .await
657            .unwrap();
658
659        let result = manager
660            .create_snapshot(
661                "test-snapshot",
662                Some("Test description".to_string()),
663                database.as_ref(),
664                &registry,
665            )
666            .await;
667
668        assert!(result.is_ok());
669        let metadata = result.unwrap();
670        assert_eq!(metadata.name, "test-snapshot");
671        assert_eq!(metadata.description, Some("Test description".to_string()));
672        assert!(metadata.entity_counts.contains_key("User"));
673    }
674
675    #[tokio::test]
676    async fn test_create_snapshot_with_time_travel() {
677        let (database, registry, temp_dir) = setup_test_env().await;
678        let manager = SnapshotManager::new(temp_dir.path());
679
680        let time_travel_state = TimeTravelSnapshotState {
681            enabled: true,
682            current_time: Some(chrono::Utc::now()),
683            scale_factor: 2.0,
684            cron_jobs: vec![],
685            mutation_rules: vec![],
686        };
687
688        let result = manager
689            .create_snapshot_with_time_travel(
690                "tt-snapshot",
691                None,
692                database.as_ref(),
693                &registry,
694                true,
695                Some(time_travel_state),
696            )
697            .await;
698
699        assert!(result.is_ok());
700        let metadata = result.unwrap();
701        assert!(metadata.time_travel_state.is_some());
702        assert_eq!(metadata.time_travel_state.unwrap().scale_factor, 2.0);
703    }
704
705    #[tokio::test]
706    async fn test_list_snapshots_empty() {
707        let temp_dir = tempfile::tempdir().unwrap();
708        let manager = SnapshotManager::new(temp_dir.path());
709
710        let result = manager.list_snapshots().await;
711        assert!(result.is_ok());
712        assert_eq!(result.unwrap().len(), 0);
713    }
714
715    #[tokio::test]
716    async fn test_list_snapshots() {
717        let (database, registry, temp_dir) = setup_test_env().await;
718        let manager = SnapshotManager::new(temp_dir.path());
719
720        // Create multiple snapshots
721        manager
722            .create_snapshot("snap1", None, database.as_ref(), &registry)
723            .await
724            .unwrap();
725        manager
726            .create_snapshot("snap2", None, database.as_ref(), &registry)
727            .await
728            .unwrap();
729
730        let result = manager.list_snapshots().await;
731        assert!(result.is_ok());
732        let snapshots = result.unwrap();
733        assert_eq!(snapshots.len(), 2);
734    }
735
736    #[tokio::test]
737    async fn test_get_snapshot_metadata() {
738        let (database, registry, temp_dir) = setup_test_env().await;
739        let manager = SnapshotManager::new(temp_dir.path());
740
741        manager
742            .create_snapshot("test", None, database.as_ref(), &registry)
743            .await
744            .unwrap();
745
746        let result = manager.get_snapshot_metadata("test").await;
747        assert!(result.is_ok());
748        let metadata = result.unwrap();
749        assert_eq!(metadata.name, "test");
750    }
751
752    #[tokio::test]
753    async fn test_get_snapshot_metadata_not_found() {
754        let temp_dir = tempfile::tempdir().unwrap();
755        let manager = SnapshotManager::new(temp_dir.path());
756
757        let result = manager.get_snapshot_metadata("nonexistent").await;
758        assert!(result.is_err());
759    }
760
761    #[tokio::test]
762    async fn test_delete_snapshot() {
763        let (database, registry, temp_dir) = setup_test_env().await;
764        let manager = SnapshotManager::new(temp_dir.path());
765
766        manager
767            .create_snapshot("to-delete", None, database.as_ref(), &registry)
768            .await
769            .unwrap();
770
771        let result = manager.delete_snapshot("to-delete").await;
772        assert!(result.is_ok());
773
774        // Verify it's gone
775        let snapshots = manager.list_snapshots().await.unwrap();
776        assert_eq!(snapshots.len(), 0);
777    }
778
779    #[tokio::test]
780    async fn test_delete_snapshot_not_found() {
781        let temp_dir = tempfile::tempdir().unwrap();
782        let manager = SnapshotManager::new(temp_dir.path());
783
784        let result = manager.delete_snapshot("nonexistent").await;
785        assert!(result.is_err());
786    }
787
788    #[tokio::test]
789    async fn test_restore_snapshot() {
790        let (database, registry, temp_dir) = setup_test_env().await;
791        let manager = SnapshotManager::new(temp_dir.path());
792
793        // Insert test data and create snapshot
794        database
795            .execute(
796                "INSERT INTO users (id, name) VALUES (?, ?)",
797                &[
798                    serde_json::Value::String("1".to_string()),
799                    serde_json::Value::String("Original".to_string()),
800                ],
801            )
802            .await
803            .unwrap();
804
805        manager
806            .create_snapshot("backup", None, database.as_ref(), &registry)
807            .await
808            .unwrap();
809
810        // Modify data
811        database
812            .execute(
813                "UPDATE users SET name = ? WHERE id = ?",
814                &[
815                    serde_json::Value::String("Modified".to_string()),
816                    serde_json::Value::String("1".to_string()),
817                ],
818            )
819            .await
820            .unwrap();
821
822        // Restore snapshot
823        let result = manager.restore_snapshot("backup", database.as_ref(), &registry).await;
824        assert!(result.is_ok());
825
826        // Verify restoration
827        let rows = database
828            .query(
829                "SELECT * FROM users WHERE id = ?",
830                &[serde_json::Value::String("1".to_string())],
831            )
832            .await
833            .unwrap();
834        assert_eq!(rows.len(), 1);
835        assert_eq!(rows[0].get("name").unwrap().as_str().unwrap(), "Original");
836    }
837
838    #[tokio::test]
839    async fn test_restore_snapshot_not_found() {
840        let (database, registry, temp_dir) = setup_test_env().await;
841        let manager = SnapshotManager::new(temp_dir.path());
842
843        let result = manager.restore_snapshot("nonexistent", database.as_ref(), &registry).await;
844        assert!(result.is_err());
845    }
846
847    #[tokio::test]
848    async fn test_reset_database() {
849        let (database, registry, _temp_dir) = setup_test_env().await;
850
851        // Insert test data
852        database
853            .execute(
854                "INSERT INTO users (id, name) VALUES (?, ?)",
855                &[
856                    serde_json::Value::String("1".to_string()),
857                    serde_json::Value::String("Test".to_string()),
858                ],
859            )
860            .await
861            .unwrap();
862
863        // Verify data exists
864        let rows = database.query("SELECT * FROM users", &[]).await.unwrap();
865        assert_eq!(rows.len(), 1);
866
867        // Reset database
868        let result = reset_database(database.as_ref(), &registry).await;
869        assert!(result.is_ok());
870
871        // Verify data is cleared
872        let rows = database.query("SELECT * FROM users", &[]).await.unwrap();
873        assert_eq!(rows.len(), 0);
874    }
875
876    #[tokio::test]
877    async fn test_snapshot_ordering() {
878        let (database, registry, temp_dir) = setup_test_env().await;
879        let manager = SnapshotManager::new(temp_dir.path());
880
881        // Create snapshots with slight delay to ensure different timestamps
882        manager
883            .create_snapshot("first", None, database.as_ref(), &registry)
884            .await
885            .unwrap();
886        tokio::time::sleep(tokio::time::Duration::from_millis(10)).await;
887        manager
888            .create_snapshot("second", None, database.as_ref(), &registry)
889            .await
890            .unwrap();
891
892        // List should be sorted by creation time (newest first)
893        let snapshots = manager.list_snapshots().await.unwrap();
894        assert_eq!(snapshots.len(), 2);
895        assert_eq!(snapshots[0].name, "second");
896        assert_eq!(snapshots[1].name, "first");
897    }
898
899    #[tokio::test]
900    async fn test_snapshot_entity_counts() {
901        let (database, mut registry, temp_dir) = setup_test_env().await;
902
903        // Add another entity
904        let base_schema = SchemaDefinition::new("Product".to_string())
905            .with_field(FieldDefinition::new("id".to_string(), "string".to_string()));
906        let vbr_schema = VbrSchemaDefinition::new(base_schema);
907        let entity = Entity::new("Product".to_string(), vbr_schema);
908
909        let manager_m = MigrationManager::new();
910        let create_sql = manager_m.generate_create_table(&entity).unwrap();
911        database.create_table(&create_sql).await.unwrap();
912        registry.register(entity).unwrap();
913
914        // Insert data
915        database
916            .execute(
917                "INSERT INTO users (id, name) VALUES (?, ?)",
918                &[
919                    serde_json::Value::String("1".to_string()),
920                    serde_json::Value::String("User1".to_string()),
921                ],
922            )
923            .await
924            .unwrap();
925
926        database
927            .execute(
928                "INSERT INTO products (id) VALUES (?)",
929                &[serde_json::Value::String("1".to_string())],
930            )
931            .await
932            .unwrap();
933        database
934            .execute(
935                "INSERT INTO products (id) VALUES (?)",
936                &[serde_json::Value::String("2".to_string())],
937            )
938            .await
939            .unwrap();
940
941        let manager = SnapshotManager::new(temp_dir.path());
942        let metadata = manager
943            .create_snapshot("multi-entity", None, database.as_ref(), &registry)
944            .await
945            .unwrap();
946
947        assert_eq!(metadata.entity_counts.get("User").unwrap(), &1);
948        assert_eq!(metadata.entity_counts.get("Product").unwrap(), &2);
949    }
950
951    #[tokio::test]
952    async fn test_snapshot_with_empty_tables() {
953        let (database, registry, temp_dir) = setup_test_env().await;
954        let manager = SnapshotManager::new(temp_dir.path());
955
956        let result = manager.create_snapshot("empty", None, database.as_ref(), &registry).await;
957        assert!(result.is_ok());
958
959        let metadata = result.unwrap();
960        assert_eq!(metadata.entity_counts.get("User").unwrap(), &0);
961    }
962}
mockforge_vbr/snapshots.rs

mockforge_vbr/
snapshots.rs
