use crate::database::VirtualDatabase;
use crate::schema::VbrSchemaDefinition;
use crate::{Error, Result};
use mockforge_core::intelligent_behavior::rules::StateMachine;
use std::collections::HashMap;
use tracing::warn;
#[derive(Debug, Clone)]
pub struct Entity {
pub name: String,
pub schema: VbrSchemaDefinition,
pub table_name: String,
pub state_machine: Option<StateMachine>,
}
impl Entity {
pub fn new(name: String, schema: VbrSchemaDefinition) -> Self {
let table_name = name.to_lowercase() + "s"; Self {
name,
schema,
table_name,
state_machine: None,
}
}
pub fn with_state_machine(
name: String,
schema: VbrSchemaDefinition,
state_machine: StateMachine,
) -> Self {
let table_name = name.to_lowercase() + "s";
Self {
name,
schema,
table_name,
state_machine: Some(state_machine),
}
}
pub fn set_state_machine(&mut self, state_machine: StateMachine) {
self.state_machine = Some(state_machine);
}
pub fn state_machine(&self) -> Option<&StateMachine> {
self.state_machine.as_ref()
}
pub fn has_state_machine(&self) -> bool {
self.state_machine.is_some()
}
pub async fn apply_state_transition(
&self,
database: &dyn VirtualDatabase,
record_id: &str,
new_state: &str,
state_field_name: Option<&str>,
) -> Result<()> {
let derived_field;
let field_name = if let Some(name) = state_field_name {
name
} else if let Some(ref sm) = self.state_machine {
let candidate = format!("{}_status", sm.resource_type.to_ascii_lowercase());
if self.schema.base.fields.iter().any(|f| f.name == candidate) {
derived_field = candidate;
&derived_field
} else {
"status"
}
} else {
"status"
};
let field_exists = self.schema.base.fields.iter().any(|f| f.name == field_name);
if !field_exists {
warn!(
"State field '{}' not found in entity schema, attempting update anyway",
field_name
);
}
let query = format!("UPDATE {} SET {} = ? WHERE id = ?", self.table_name, field_name);
database
.execute(
&query,
&[
serde_json::Value::String(new_state.to_string()),
serde_json::Value::String(record_id.to_string()),
],
)
.await
.map_err(|e| Error::generic(format!("Failed to update entity state: {}", e)))?;
Ok(())
}
pub async fn get_current_state(
&self,
database: &dyn VirtualDatabase,
record_id: &str,
state_field_name: Option<&str>,
) -> Result<Option<String>> {
let field_name = state_field_name.unwrap_or("status");
let query = format!("SELECT {} FROM {} WHERE id = ?", field_name, self.table_name);
let results = database
.query(&query, &[serde_json::Value::String(record_id.to_string())])
.await
.map_err(|e| Error::generic(format!("Failed to query entity state: {}", e)))?;
if let Some(row) = results.first() {
if let Some(value) = row.get(field_name) {
if let Some(state) = value.as_str() {
return Ok(Some(state.to_string()));
}
}
}
Ok(None)
}
pub async fn can_transition(
&self,
database: &dyn VirtualDatabase,
record_id: &str,
to_state: &str,
state_field_name: Option<&str>,
) -> Result<bool> {
let state_machine = self
.state_machine
.as_ref()
.ok_or_else(|| Error::generic("Entity does not have a state machine configured"))?;
let current_state = self
.get_current_state(database, record_id, state_field_name)
.await?
.ok_or_else(|| {
Error::generic(format!("Record '{}' not found or has no state", record_id))
})?;
Ok(state_machine.can_transition(¤t_state, to_state))
}
pub fn name(&self) -> &str {
&self.name
}
pub fn table_name(&self) -> &str {
&self.table_name
}
}
#[derive(Clone)]
pub struct EntityRegistry {
entities: HashMap<String, Entity>,
}
impl EntityRegistry {
pub fn new() -> Self {
Self {
entities: HashMap::new(),
}
}
pub fn register(&mut self, entity: Entity) -> Result<()> {
let name = entity.name.clone();
if self.entities.contains_key(&name) {
return Err(Error::generic(format!("Entity '{}' already registered", name)));
}
self.entities.insert(name, entity);
Ok(())
}
pub fn get(&self, name: &str) -> Option<&Entity> {
self.entities.get(name)
}
pub fn list(&self) -> Vec<String> {
self.entities.keys().cloned().collect()
}
pub fn exists(&self, name: &str) -> bool {
self.entities.contains_key(name)
}
pub fn remove(&mut self, name: &str) -> Result<()> {
self.entities
.remove(name)
.ok_or_else(|| Error::generic(format!("Entity '{}' not found", name)))?;
Ok(())
}
}
impl Default for EntityRegistry {
fn default() -> Self {
Self::new()
}
}
#[cfg(test)]
mod tests {
use super::*;
use mockforge_data::SchemaDefinition;
fn create_test_schema(name: &str) -> VbrSchemaDefinition {
let base_schema = SchemaDefinition::new(name.to_string());
VbrSchemaDefinition::new(base_schema)
}
#[test]
fn test_entity_creation() {
let base_schema = SchemaDefinition::new("User".to_string());
let vbr_schema = VbrSchemaDefinition::new(base_schema);
let entity = Entity::new("User".to_string(), vbr_schema);
assert_eq!(entity.name(), "User");
assert_eq!(entity.table_name(), "users");
}
#[test]
fn test_entity_table_name_pluralization() {
let entity = Entity::new("Order".to_string(), create_test_schema("Order"));
assert_eq!(entity.table_name(), "orders");
let entity2 = Entity::new("Product".to_string(), create_test_schema("Product"));
assert_eq!(entity2.table_name(), "products");
}
#[test]
fn test_entity_table_name_lowercase() {
let entity = Entity::new("UserProfile".to_string(), create_test_schema("UserProfile"));
assert_eq!(entity.table_name(), "userprofiles");
}
#[test]
fn test_entity_no_state_machine() {
let entity = Entity::new("Item".to_string(), create_test_schema("Item"));
assert!(!entity.has_state_machine());
assert!(entity.state_machine().is_none());
}
#[test]
fn test_entity_clone() {
let entity = Entity::new("Test".to_string(), create_test_schema("Test"));
let cloned = entity.clone();
assert_eq!(entity.name(), cloned.name());
assert_eq!(entity.table_name(), cloned.table_name());
}
#[test]
fn test_entity_debug() {
let entity = Entity::new("DebugEntity".to_string(), create_test_schema("DebugEntity"));
let debug = format!("{:?}", entity);
assert!(debug.contains("Entity"));
assert!(debug.contains("DebugEntity"));
}
#[test]
fn test_entity_name_getter() {
let entity = Entity::new("Customer".to_string(), create_test_schema("Customer"));
assert_eq!(entity.name(), "Customer");
assert_eq!(entity.name, "Customer");
}
#[test]
fn test_entity_registry() {
let mut registry = EntityRegistry::new();
let base_schema = SchemaDefinition::new("User".to_string());
let vbr_schema = VbrSchemaDefinition::new(base_schema);
let entity = Entity::new("User".to_string(), vbr_schema);
assert!(registry.register(entity).is_ok());
assert!(registry.exists("User"));
assert!(registry.get("User").is_some());
}
#[test]
fn test_entity_registry_duplicate() {
let mut registry = EntityRegistry::new();
let base_schema1 = SchemaDefinition::new("User".to_string());
let vbr_schema1 = VbrSchemaDefinition::new(base_schema1);
let entity1 = Entity::new("User".to_string(), vbr_schema1);
let base_schema2 = SchemaDefinition::new("User".to_string());
let vbr_schema2 = VbrSchemaDefinition::new(base_schema2);
let entity2 = Entity::new("User".to_string(), vbr_schema2);
assert!(registry.register(entity1).is_ok());
assert!(registry.register(entity2).is_err());
}
#[test]
fn test_entity_registry_default() {
let registry = EntityRegistry::default();
assert!(registry.list().is_empty());
}
#[test]
fn test_entity_registry_new() {
let registry = EntityRegistry::new();
assert!(registry.list().is_empty());
assert!(!registry.exists("Anything"));
}
#[test]
fn test_entity_registry_get_nonexistent() {
let registry = EntityRegistry::new();
assert!(registry.get("NonExistent").is_none());
}
#[test]
fn test_entity_registry_list() {
let mut registry = EntityRegistry::new();
registry
.register(Entity::new("User".to_string(), create_test_schema("User")))
.unwrap();
registry
.register(Entity::new("Order".to_string(), create_test_schema("Order")))
.unwrap();
registry
.register(Entity::new("Product".to_string(), create_test_schema("Product")))
.unwrap();
let list = registry.list();
assert_eq!(list.len(), 3);
assert!(list.contains(&"User".to_string()));
assert!(list.contains(&"Order".to_string()));
assert!(list.contains(&"Product".to_string()));
}
#[test]
fn test_entity_registry_exists() {
let mut registry = EntityRegistry::new();
registry
.register(Entity::new("User".to_string(), create_test_schema("User")))
.unwrap();
assert!(registry.exists("User"));
assert!(!registry.exists("Order"));
assert!(!registry.exists("user")); }
#[test]
fn test_entity_registry_remove() {
let mut registry = EntityRegistry::new();
registry
.register(Entity::new("User".to_string(), create_test_schema("User")))
.unwrap();
assert!(registry.exists("User"));
assert!(registry.remove("User").is_ok());
assert!(!registry.exists("User"));
}
#[test]
fn test_entity_registry_remove_nonexistent() {
let mut registry = EntityRegistry::new();
let result = registry.remove("NonExistent");
assert!(result.is_err());
}
#[test]
fn test_entity_registry_clone() {
let mut registry = EntityRegistry::new();
registry
.register(Entity::new("User".to_string(), create_test_schema("User")))
.unwrap();
let cloned = registry.clone();
assert!(cloned.exists("User"));
assert_eq!(cloned.list().len(), 1);
}
#[test]
fn test_entity_registry_multiple_operations() {
let mut registry = EntityRegistry::new();
registry
.register(Entity::new("A".to_string(), create_test_schema("A")))
.unwrap();
registry
.register(Entity::new("B".to_string(), create_test_schema("B")))
.unwrap();
assert_eq!(registry.list().len(), 2);
assert!(registry.exists("A"));
assert!(registry.exists("B"));
registry.remove("A").unwrap();
assert_eq!(registry.list().len(), 1);
assert!(!registry.exists("A"));
assert!(registry.exists("B"));
registry
.register(Entity::new("C".to_string(), create_test_schema("C")))
.unwrap();
assert_eq!(registry.list().len(), 2);
}
#[test]
fn test_entity_registry_get_returns_reference() {
let mut registry = EntityRegistry::new();
registry
.register(Entity::new("User".to_string(), create_test_schema("User")))
.unwrap();
let entity = registry.get("User").unwrap();
assert_eq!(entity.name(), "User");
assert_eq!(entity.table_name(), "users");
}
}