pub use mockforge_core::{Error, Result};
use std::collections::HashMap;
use std::sync::Arc;
pub mod config;
pub mod database;
pub mod entities;
pub mod schema;
pub mod constraints;
pub mod migration;
pub mod api_generator;
pub mod handlers;
pub mod auth;
pub mod session;
pub mod aging;
pub mod mutation_rules;
pub mod scheduler;
pub mod integration;
pub mod openapi;
pub mod seeding;
pub mod id_generation;
pub mod snapshots;
pub use config::{StorageBackend, VbrConfig};
pub use database::VirtualDatabase;
pub use entities::{Entity, EntityRegistry};
pub use mutation_rules::{
ComparisonOperator, MutationOperation, MutationRule, MutationRuleManager, MutationTrigger,
};
pub use schema::{ManyToManyDefinition, VbrSchemaDefinition};
pub use snapshots::{SnapshotMetadata, TimeTravelSnapshotState};
pub struct VbrEngine {
config: VbrConfig,
database: Arc<dyn VirtualDatabase + Send + Sync>,
registry: EntityRegistry,
}
impl VbrEngine {
pub async fn new(config: VbrConfig) -> Result<Self> {
let database = database::create_database(&config.storage).await?;
let registry = EntityRegistry::new();
Ok(Self {
config,
database,
registry,
})
}
pub fn config(&self) -> &VbrConfig {
&self.config
}
pub fn database_arc(&self) -> Arc<dyn VirtualDatabase + Send + Sync> {
Arc::clone(&self.database)
}
pub fn database(&self) -> &dyn VirtualDatabase {
self.database.as_ref()
}
pub fn registry(&self) -> &EntityRegistry {
&self.registry
}
pub fn registry_mut(&mut self) -> &mut EntityRegistry {
&mut self.registry
}
pub async fn from_openapi(
config: VbrConfig,
openapi_content: &str,
) -> Result<(Self, openapi::OpenApiConversionResult)> {
let json_value: serde_json::Value = if openapi_content.trim_start().starts_with('{') {
serde_json::from_str(openapi_content)
.map_err(|e| Error::generic(format!("Failed to parse OpenAPI JSON: {}", e)))?
} else {
serde_yaml::from_str(openapi_content)
.map_err(|e| Error::generic(format!("Failed to parse OpenAPI YAML: {}", e)))?
};
let spec = mockforge_core::openapi::OpenApiSpec::from_json(json_value)
.map_err(|e| Error::generic(format!("Failed to load OpenAPI spec: {}", e)))?;
spec.validate()
.map_err(|e| Error::generic(format!("Invalid OpenAPI specification: {}", e)))?;
let conversion_result = openapi::convert_openapi_to_vbr(&spec)?;
let mut engine = Self::new(config).await?;
for (entity_name, vbr_schema) in &conversion_result.entities {
let entity = Entity::new(entity_name.clone(), vbr_schema.clone());
engine.registry_mut().register(entity.clone())?;
migration::create_table_for_entity(engine.database.as_ref(), &entity).await?;
}
Ok((engine, conversion_result))
}
pub async fn from_openapi_file<P: AsRef<std::path::Path>>(
config: VbrConfig,
file_path: P,
) -> Result<(Self, openapi::OpenApiConversionResult)> {
let content = tokio::fs::read_to_string(file_path.as_ref())
.await
.map_err(|e| Error::generic(format!("Failed to read OpenAPI file: {}", e)))?;
Self::from_openapi(config, &content).await
}
pub async fn seed_entity(
&self,
entity_name: &str,
records: &[HashMap<String, serde_json::Value>],
) -> Result<usize> {
seeding::seed_entity(self.database.as_ref(), &self.registry, entity_name, records).await
}
pub async fn seed_all(&self, seed_data: &seeding::SeedData) -> Result<HashMap<String, usize>> {
seeding::seed_all(self.database.as_ref(), &self.registry, seed_data).await
}
pub async fn seed_from_file<P: AsRef<std::path::Path>>(
&self,
file_path: P,
) -> Result<HashMap<String, usize>> {
let seed_data = seeding::load_seed_file(file_path).await?;
self.seed_all(&seed_data).await
}
pub async fn clear_entity(&self, entity_name: &str) -> Result<()> {
seeding::clear_entity(self.database.as_ref(), &self.registry, entity_name).await
}
pub async fn clear_all(&self) -> Result<()> {
seeding::clear_all(self.database.as_ref(), &self.registry).await
}
pub async fn create_snapshot<P: AsRef<std::path::Path>>(
&self,
name: &str,
description: Option<String>,
snapshots_dir: P,
) -> Result<SnapshotMetadata> {
let manager = snapshots::SnapshotManager::new(snapshots_dir);
manager
.create_snapshot(name, description, self.database.as_ref(), &self.registry)
.await
}
pub async fn create_snapshot_with_time_travel<P: AsRef<std::path::Path>>(
&self,
name: &str,
description: Option<String>,
snapshots_dir: P,
include_time_travel: bool,
time_travel_state: Option<TimeTravelSnapshotState>,
) -> Result<SnapshotMetadata> {
let manager = snapshots::SnapshotManager::new(snapshots_dir);
manager
.create_snapshot_with_time_travel(
name,
description,
self.database.as_ref(),
&self.registry,
include_time_travel,
time_travel_state,
)
.await
}
pub async fn restore_snapshot<P: AsRef<std::path::Path>>(
&self,
name: &str,
snapshots_dir: P,
) -> Result<()> {
let manager = snapshots::SnapshotManager::new(snapshots_dir);
manager.restore_snapshot(name, self.database.as_ref(), &self.registry).await
}
pub async fn restore_snapshot_with_time_travel<P, F>(
&self,
name: &str,
snapshots_dir: P,
restore_time_travel: bool,
time_travel_restore_callback: Option<F>,
) -> Result<()>
where
P: AsRef<std::path::Path>,
F: FnOnce(TimeTravelSnapshotState) -> Result<()>,
{
let manager = snapshots::SnapshotManager::new(snapshots_dir);
manager
.restore_snapshot_with_time_travel(
name,
self.database.as_ref(),
&self.registry,
restore_time_travel,
time_travel_restore_callback,
)
.await
}
pub async fn list_snapshots<P: AsRef<std::path::Path>>(
snapshots_dir: P,
) -> Result<Vec<SnapshotMetadata>> {
let manager = snapshots::SnapshotManager::new(snapshots_dir);
manager.list_snapshots().await
}
pub async fn delete_snapshot<P: AsRef<std::path::Path>>(
name: &str,
snapshots_dir: P,
) -> Result<()> {
let manager = snapshots::SnapshotManager::new(snapshots_dir);
manager.delete_snapshot(name).await
}
pub async fn reset(&self) -> Result<()> {
snapshots::reset_database(self.database.as_ref(), &self.registry).await
}
pub async fn export_state(&self) -> Result<serde_json::Value> {
let mut state = serde_json::Map::new();
let mut entities = serde_json::Map::new();
for entity_name in self.registry.list() {
if let Some(entity) = self.registry.get(&entity_name) {
let table_name = entity.table_name();
let query = format!("SELECT * FROM {}", table_name);
let records = self.database.query(&query, &[]).await?;
entities.insert(
entity_name.clone(),
serde_json::Value::Array(
records.into_iter().map(|r| serde_json::json!(r)).collect(),
),
);
}
}
state.insert("entities".to_string(), serde_json::Value::Object(entities));
state.insert(
"storage_backend".to_string(),
serde_json::Value::String(self.database.connection_info()),
);
state.insert(
"exported_at".to_string(),
serde_json::Value::String(chrono::Utc::now().to_rfc3339()),
);
Ok(serde_json::Value::Object(state))
}
pub async fn import_state(&self, state: serde_json::Value) -> Result<()> {
self.reset().await?;
let entities = state
.get("entities")
.and_then(|e| e.as_object())
.ok_or_else(|| Error::generic("Invalid state format: missing entities"))?;
for (entity_name, records_value) in entities {
let records = records_value.as_array().ok_or_else(|| {
Error::generic(format!("Invalid records for entity {}", entity_name))
})?;
if let Some(entity) = self.registry.get(entity_name) {
let table_name = entity.table_name();
for record in records {
let record_obj = record
.as_object()
.ok_or_else(|| Error::generic("Invalid record format"))?;
let fields: Vec<String> = record_obj.keys().cloned().collect();
let placeholders: Vec<String> =
(0..fields.len()).map(|_| "?".to_string()).collect();
let values: Vec<serde_json::Value> = fields
.iter()
.map(|f| record_obj.get(f).cloned().unwrap_or(serde_json::Value::Null))
.collect();
let query = format!(
"INSERT INTO {} ({}) VALUES ({})",
table_name,
fields.join(", "),
placeholders.join(", ")
);
self.database.execute(&query, &values).await?;
}
}
}
Ok(())
}
pub async fn state_summary(&self) -> String {
let mut total_records = 0usize;
let entity_count = self.registry.list().len();
for entity_name in self.registry.list() {
if let Some(entity) = self.registry.get(&entity_name) {
let table_name = entity.table_name();
let count_query = format!("SELECT COUNT(*) as count FROM {}", table_name);
if let Ok(results) = self.database.query(&count_query, &[]).await {
if let Some(count) =
results.first().and_then(|r| r.get("count")).and_then(|v| v.as_u64())
{
total_records += count as usize;
}
}
}
}
format!("{} entities, {} records", entity_count, total_records)
}
}
use async_trait::async_trait;
use mockforge_core::snapshots::ProtocolStateExporter;
#[async_trait]
impl ProtocolStateExporter for VbrEngine {
fn protocol_name(&self) -> &str {
"vbr"
}
async fn export_state(&self) -> Result<serde_json::Value> {
VbrEngine::export_state(self).await
}
async fn import_state(&self, state: serde_json::Value) -> Result<()> {
VbrEngine::import_state(self, state).await
}
async fn state_summary(&self) -> String {
VbrEngine::state_summary(self).await
}
}
#[cfg(test)]
mod tests {
use super::*;
#[tokio::test]
async fn test_vbr_engine_creation() {
let config = VbrConfig::default().with_storage_backend(StorageBackend::Memory);
let engine = VbrEngine::new(config).await;
assert!(engine.is_ok());
}
}