1use crate::database::VirtualDatabase;
7use crate::schema::VbrSchemaDefinition;
8use crate::{Error, Result};
9use mockforge_core::intelligent_behavior::rules::StateMachine;
10use std::collections::HashMap;
11use tracing::warn;
12
13#[derive(Debug, Clone)]
15pub struct Entity {
16 pub name: String,
18
19 pub schema: VbrSchemaDefinition,
21
22 pub table_name: String,
24
25 pub state_machine: Option<StateMachine>,
30}
31
32impl Entity {
33 pub fn new(name: String, schema: VbrSchemaDefinition) -> Self {
35 let table_name = name.to_lowercase() + "s"; Self {
37 name,
38 schema,
39 table_name,
40 state_machine: None,
41 }
42 }
43
44 pub fn with_state_machine(
46 name: String,
47 schema: VbrSchemaDefinition,
48 state_machine: StateMachine,
49 ) -> Self {
50 let table_name = name.to_lowercase() + "s";
51 Self {
52 name,
53 schema,
54 table_name,
55 state_machine: Some(state_machine),
56 }
57 }
58
59 pub fn set_state_machine(&mut self, state_machine: StateMachine) {
61 self.state_machine = Some(state_machine);
62 }
63
64 pub fn state_machine(&self) -> Option<&StateMachine> {
66 self.state_machine.as_ref()
67 }
68
69 pub fn has_state_machine(&self) -> bool {
71 self.state_machine.is_some()
72 }
73
74 pub async fn apply_state_transition(
83 &self,
84 database: &dyn VirtualDatabase,
85 record_id: &str,
86 new_state: &str,
87 state_field_name: Option<&str>,
88 ) -> Result<()> {
89 let field_name = if let Some(name) = state_field_name {
92 name
93 } else if let Some(ref sm) = self.state_machine {
94 "status"
98 } else {
99 "status"
100 };
101
102 let field_exists = self.schema.base.fields.iter().any(|f| f.name == field_name);
104
105 if !field_exists {
106 warn!(
109 "State field '{}' not found in entity schema, attempting update anyway",
110 field_name
111 );
112 }
113
114 let query = format!("UPDATE {} SET {} = ? WHERE id = ?", self.table_name, field_name);
116
117 database
118 .execute(
119 &query,
120 &[
121 serde_json::Value::String(new_state.to_string()),
122 serde_json::Value::String(record_id.to_string()),
123 ],
124 )
125 .await
126 .map_err(|e| Error::generic(format!("Failed to update entity state: {}", e)))?;
127
128 Ok(())
129 }
130
131 pub async fn get_current_state(
135 &self,
136 database: &dyn VirtualDatabase,
137 record_id: &str,
138 state_field_name: Option<&str>,
139 ) -> Result<Option<String>> {
140 let field_name = state_field_name.unwrap_or("status");
141
142 let query = format!("SELECT {} FROM {} WHERE id = ?", field_name, self.table_name);
143
144 let results = database
145 .query(&query, &[serde_json::Value::String(record_id.to_string())])
146 .await
147 .map_err(|e| Error::generic(format!("Failed to query entity state: {}", e)))?;
148
149 if let Some(row) = results.first() {
150 if let Some(value) = row.get(field_name) {
151 if let Some(state) = value.as_str() {
152 return Ok(Some(state.to_string()));
153 }
154 }
155 }
156
157 Ok(None)
158 }
159
160 pub async fn can_transition(
165 &self,
166 database: &dyn VirtualDatabase,
167 record_id: &str,
168 to_state: &str,
169 state_field_name: Option<&str>,
170 ) -> Result<bool> {
171 let state_machine = self
172 .state_machine
173 .as_ref()
174 .ok_or_else(|| Error::generic("Entity does not have a state machine configured"))?;
175
176 let current_state = self
178 .get_current_state(database, record_id, state_field_name)
179 .await?
180 .ok_or_else(|| {
181 Error::generic(format!("Record '{}' not found or has no state", record_id))
182 })?;
183
184 Ok(state_machine.can_transition(¤t_state, to_state))
186 }
187
188 pub fn name(&self) -> &str {
190 &self.name
191 }
192
193 pub fn table_name(&self) -> &str {
195 &self.table_name
196 }
197}
198
199#[derive(Clone)]
201pub struct EntityRegistry {
202 entities: HashMap<String, Entity>,
204}
205
206impl EntityRegistry {
207 pub fn new() -> Self {
209 Self {
210 entities: HashMap::new(),
211 }
212 }
213
214 pub fn register(&mut self, entity: Entity) -> Result<()> {
216 let name = entity.name.clone();
217 if self.entities.contains_key(&name) {
218 return Err(Error::generic(format!("Entity '{}' already registered", name)));
219 }
220 self.entities.insert(name, entity);
221 Ok(())
222 }
223
224 pub fn get(&self, name: &str) -> Option<&Entity> {
226 self.entities.get(name)
227 }
228
229 pub fn list(&self) -> Vec<String> {
231 self.entities.keys().cloned().collect()
232 }
233
234 pub fn exists(&self, name: &str) -> bool {
236 self.entities.contains_key(name)
237 }
238
239 pub fn remove(&mut self, name: &str) -> Result<()> {
241 self.entities
242 .remove(name)
243 .ok_or_else(|| Error::generic(format!("Entity '{}' not found", name)))?;
244 Ok(())
245 }
246}
247
248impl Default for EntityRegistry {
249 fn default() -> Self {
250 Self::new()
251 }
252}
253
254#[cfg(test)]
255mod tests {
256 use super::*;
257 use mockforge_data::SchemaDefinition;
258
259 fn create_test_schema(name: &str) -> VbrSchemaDefinition {
260 let base_schema = SchemaDefinition::new(name.to_string());
261 VbrSchemaDefinition::new(base_schema)
262 }
263
264 #[test]
266 fn test_entity_creation() {
267 let base_schema = SchemaDefinition::new("User".to_string());
268 let vbr_schema = VbrSchemaDefinition::new(base_schema);
269 let entity = Entity::new("User".to_string(), vbr_schema);
270
271 assert_eq!(entity.name(), "User");
272 assert_eq!(entity.table_name(), "users");
273 }
274
275 #[test]
276 fn test_entity_table_name_pluralization() {
277 let entity = Entity::new("Order".to_string(), create_test_schema("Order"));
278 assert_eq!(entity.table_name(), "orders");
279
280 let entity2 = Entity::new("Product".to_string(), create_test_schema("Product"));
281 assert_eq!(entity2.table_name(), "products");
282 }
283
284 #[test]
285 fn test_entity_table_name_lowercase() {
286 let entity = Entity::new("UserProfile".to_string(), create_test_schema("UserProfile"));
287 assert_eq!(entity.table_name(), "userprofiles");
288 }
289
290 #[test]
291 fn test_entity_no_state_machine() {
292 let entity = Entity::new("Item".to_string(), create_test_schema("Item"));
293 assert!(!entity.has_state_machine());
294 assert!(entity.state_machine().is_none());
295 }
296
297 #[test]
298 fn test_entity_clone() {
299 let entity = Entity::new("Test".to_string(), create_test_schema("Test"));
300 let cloned = entity.clone();
301 assert_eq!(entity.name(), cloned.name());
302 assert_eq!(entity.table_name(), cloned.table_name());
303 }
304
305 #[test]
306 fn test_entity_debug() {
307 let entity = Entity::new("DebugEntity".to_string(), create_test_schema("DebugEntity"));
308 let debug = format!("{:?}", entity);
309 assert!(debug.contains("Entity"));
310 assert!(debug.contains("DebugEntity"));
311 }
312
313 #[test]
314 fn test_entity_name_getter() {
315 let entity = Entity::new("Customer".to_string(), create_test_schema("Customer"));
316 assert_eq!(entity.name(), "Customer");
317 assert_eq!(entity.name, "Customer");
318 }
319
320 #[test]
322 fn test_entity_registry() {
323 let mut registry = EntityRegistry::new();
324
325 let base_schema = SchemaDefinition::new("User".to_string());
326 let vbr_schema = VbrSchemaDefinition::new(base_schema);
327 let entity = Entity::new("User".to_string(), vbr_schema);
328
329 assert!(registry.register(entity).is_ok());
330 assert!(registry.exists("User"));
331 assert!(registry.get("User").is_some());
332 }
333
334 #[test]
335 fn test_entity_registry_duplicate() {
336 let mut registry = EntityRegistry::new();
337
338 let base_schema1 = SchemaDefinition::new("User".to_string());
339 let vbr_schema1 = VbrSchemaDefinition::new(base_schema1);
340 let entity1 = Entity::new("User".to_string(), vbr_schema1);
341
342 let base_schema2 = SchemaDefinition::new("User".to_string());
343 let vbr_schema2 = VbrSchemaDefinition::new(base_schema2);
344 let entity2 = Entity::new("User".to_string(), vbr_schema2);
345
346 assert!(registry.register(entity1).is_ok());
347 assert!(registry.register(entity2).is_err());
348 }
349
350 #[test]
351 fn test_entity_registry_default() {
352 let registry = EntityRegistry::default();
353 assert!(registry.list().is_empty());
354 }
355
356 #[test]
357 fn test_entity_registry_new() {
358 let registry = EntityRegistry::new();
359 assert!(registry.list().is_empty());
360 assert!(!registry.exists("Anything"));
361 }
362
363 #[test]
364 fn test_entity_registry_get_nonexistent() {
365 let registry = EntityRegistry::new();
366 assert!(registry.get("NonExistent").is_none());
367 }
368
369 #[test]
370 fn test_entity_registry_list() {
371 let mut registry = EntityRegistry::new();
372
373 registry
374 .register(Entity::new("User".to_string(), create_test_schema("User")))
375 .unwrap();
376 registry
377 .register(Entity::new("Order".to_string(), create_test_schema("Order")))
378 .unwrap();
379 registry
380 .register(Entity::new("Product".to_string(), create_test_schema("Product")))
381 .unwrap();
382
383 let list = registry.list();
384 assert_eq!(list.len(), 3);
385 assert!(list.contains(&"User".to_string()));
386 assert!(list.contains(&"Order".to_string()));
387 assert!(list.contains(&"Product".to_string()));
388 }
389
390 #[test]
391 fn test_entity_registry_exists() {
392 let mut registry = EntityRegistry::new();
393 registry
394 .register(Entity::new("User".to_string(), create_test_schema("User")))
395 .unwrap();
396
397 assert!(registry.exists("User"));
398 assert!(!registry.exists("Order"));
399 assert!(!registry.exists("user")); }
401
402 #[test]
403 fn test_entity_registry_remove() {
404 let mut registry = EntityRegistry::new();
405 registry
406 .register(Entity::new("User".to_string(), create_test_schema("User")))
407 .unwrap();
408
409 assert!(registry.exists("User"));
410 assert!(registry.remove("User").is_ok());
411 assert!(!registry.exists("User"));
412 }
413
414 #[test]
415 fn test_entity_registry_remove_nonexistent() {
416 let mut registry = EntityRegistry::new();
417 let result = registry.remove("NonExistent");
418 assert!(result.is_err());
419 }
420
421 #[test]
422 fn test_entity_registry_clone() {
423 let mut registry = EntityRegistry::new();
424 registry
425 .register(Entity::new("User".to_string(), create_test_schema("User")))
426 .unwrap();
427
428 let cloned = registry.clone();
429 assert!(cloned.exists("User"));
430 assert_eq!(cloned.list().len(), 1);
431 }
432
433 #[test]
434 fn test_entity_registry_multiple_operations() {
435 let mut registry = EntityRegistry::new();
436
437 registry
439 .register(Entity::new("A".to_string(), create_test_schema("A")))
440 .unwrap();
441 registry
442 .register(Entity::new("B".to_string(), create_test_schema("B")))
443 .unwrap();
444
445 assert_eq!(registry.list().len(), 2);
447 assert!(registry.exists("A"));
448 assert!(registry.exists("B"));
449
450 registry.remove("A").unwrap();
452 assert_eq!(registry.list().len(), 1);
453 assert!(!registry.exists("A"));
454 assert!(registry.exists("B"));
455
456 registry
458 .register(Entity::new("C".to_string(), create_test_schema("C")))
459 .unwrap();
460 assert_eq!(registry.list().len(), 2);
461 }
462
463 #[test]
464 fn test_entity_registry_get_returns_reference() {
465 let mut registry = EntityRegistry::new();
466 registry
467 .register(Entity::new("User".to_string(), create_test_schema("User")))
468 .unwrap();
469
470 let entity = registry.get("User").unwrap();
471 assert_eq!(entity.name(), "User");
472 assert_eq!(entity.table_name(), "users");
473 }
474}