1use chrono::NaiveDate;
8use datasynth_core::models::{
9 CustomerPool, Employee, EmployeePool, EntityId, EntityRegistry, EntityType, FixedAsset,
10 FixedAssetPool, Material, MaterialPool, Vendor, VendorPool,
11};
12
13use super::{
14 AssetGenerator, AssetGeneratorConfig, CustomerGenerator, CustomerGeneratorConfig,
15 EmployeeGenerator, EmployeeGeneratorConfig, MaterialGenerator, MaterialGeneratorConfig,
16 VendorGenerator, VendorGeneratorConfig,
17};
18
19#[derive(Debug, Clone, Default)]
21pub struct EntityRegistryManagerConfig {
22 pub vendor_config: VendorGeneratorConfig,
24 pub customer_config: CustomerGeneratorConfig,
26 pub material_config: MaterialGeneratorConfig,
28 pub asset_config: AssetGeneratorConfig,
30 pub employee_config: EmployeeGeneratorConfig,
32}
33
34#[derive(Debug, Clone)]
36pub struct MasterDataCounts {
37 pub vendors: usize,
39 pub customers: usize,
41 pub materials: usize,
43 pub assets: usize,
45 pub employees: Option<usize>,
47}
48
49impl Default for MasterDataCounts {
50 fn default() -> Self {
51 Self {
52 vendors: 100,
53 customers: 200,
54 materials: 500,
55 assets: 150,
56 employees: None,
57 }
58 }
59}
60
61#[derive(Debug)]
63pub struct GeneratedMasterData {
64 pub vendors: VendorPool,
66 pub customers: CustomerPool,
68 pub materials: MaterialPool,
70 pub assets: FixedAssetPool,
72 pub employees: EmployeePool,
74 pub registry: EntityRegistry,
76}
77
78pub struct EntityRegistryManager {
80 #[allow(dead_code)]
82 seed: u64,
83 #[allow(dead_code)]
84 config: EntityRegistryManagerConfig,
85 vendor_generator: VendorGenerator,
86 customer_generator: CustomerGenerator,
87 material_generator: MaterialGenerator,
88 asset_generator: AssetGenerator,
89 employee_generator: EmployeeGenerator,
90 registry: EntityRegistry,
91}
92
93impl EntityRegistryManager {
94 pub fn new(seed: u64) -> Self {
96 Self::with_config(seed, EntityRegistryManagerConfig::default())
97 }
98
99 pub fn with_config(seed: u64, config: EntityRegistryManagerConfig) -> Self {
101 Self {
102 seed,
103 vendor_generator: VendorGenerator::with_config(seed, config.vendor_config.clone()),
104 customer_generator: CustomerGenerator::with_config(
105 seed + 1,
106 config.customer_config.clone(),
107 ),
108 material_generator: MaterialGenerator::with_config(
109 seed + 2,
110 config.material_config.clone(),
111 ),
112 asset_generator: AssetGenerator::with_config(seed + 3, config.asset_config.clone()),
113 employee_generator: EmployeeGenerator::with_config(
114 seed + 4,
115 config.employee_config.clone(),
116 ),
117 registry: EntityRegistry::new(),
118 config,
119 }
120 }
121
122 pub fn generate_master_data(
124 &mut self,
125 company_code: &str,
126 counts: &MasterDataCounts,
127 effective_date: NaiveDate,
128 date_range: (NaiveDate, NaiveDate),
129 ) -> GeneratedMasterData {
130 let vendors = self.generate_vendors(company_code, counts.vendors, effective_date);
132
133 let customers = self.generate_customers(company_code, counts.customers, effective_date);
135
136 let materials = self.generate_materials(company_code, counts.materials, effective_date);
138
139 let assets = self.generate_assets(company_code, counts.assets, date_range);
141
142 let employees = self.generate_employees(company_code, date_range);
144
145 GeneratedMasterData {
146 vendors,
147 customers,
148 materials,
149 assets,
150 employees,
151 registry: self.registry.clone(),
152 }
153 }
154
155 pub fn generate_multi_company_master_data(
157 &mut self,
158 company_codes: &[String],
159 counts: &MasterDataCounts,
160 effective_date: NaiveDate,
161 date_range: (NaiveDate, NaiveDate),
162 ) -> Vec<GeneratedMasterData> {
163 let mut results = Vec::new();
164
165 for company_code in company_codes {
166 let data = self.generate_master_data(company_code, counts, effective_date, date_range);
167 results.push(data);
168 }
169
170 results
171 }
172
173 pub fn generate_master_data_with_ic(
175 &mut self,
176 company_codes: &[String],
177 counts: &MasterDataCounts,
178 effective_date: NaiveDate,
179 date_range: (NaiveDate, NaiveDate),
180 ) -> Vec<GeneratedMasterData> {
181 let mut results = Vec::new();
182
183 for company_code in company_codes {
184 let partners: Vec<String> = company_codes
186 .iter()
187 .filter(|c| *c != company_code)
188 .cloned()
189 .collect();
190
191 let vendors = self.vendor_generator.generate_vendor_pool_with_ic(
193 counts.vendors,
194 company_code,
195 &partners,
196 effective_date,
197 );
198
199 for vendor in &vendors.vendors {
201 self.register_vendor(vendor, company_code, effective_date);
202 }
203
204 let customers = self.customer_generator.generate_customer_pool_with_ic(
206 counts.customers,
207 company_code,
208 &partners,
209 effective_date,
210 );
211
212 for customer in &customers.customers {
214 self.register_customer(customer, company_code, effective_date);
215 }
216
217 let materials = self.generate_materials(company_code, counts.materials, effective_date);
219
220 let assets = self.generate_assets(company_code, counts.assets, date_range);
222
223 let employees = self.generate_employees(company_code, date_range);
225
226 results.push(GeneratedMasterData {
227 vendors,
228 customers,
229 materials,
230 assets,
231 employees,
232 registry: self.registry.clone(),
233 });
234 }
235
236 results
237 }
238
239 fn generate_vendors(
241 &mut self,
242 company_code: &str,
243 count: usize,
244 effective_date: NaiveDate,
245 ) -> VendorPool {
246 let pool = self
247 .vendor_generator
248 .generate_vendor_pool(count, company_code, effective_date);
249
250 for vendor in &pool.vendors {
252 self.register_vendor(vendor, company_code, effective_date);
253 }
254
255 pool
256 }
257
258 fn generate_customers(
260 &mut self,
261 company_code: &str,
262 count: usize,
263 effective_date: NaiveDate,
264 ) -> CustomerPool {
265 let pool =
266 self.customer_generator
267 .generate_customer_pool(count, company_code, effective_date);
268
269 for customer in &pool.customers {
271 self.register_customer(customer, company_code, effective_date);
272 }
273
274 pool
275 }
276
277 fn generate_materials(
279 &mut self,
280 company_code: &str,
281 count: usize,
282 effective_date: NaiveDate,
283 ) -> MaterialPool {
284 let pool = self.material_generator.generate_material_pool_with_bom(
285 count,
286 0.25, company_code,
288 effective_date,
289 );
290
291 for material in &pool.materials {
293 self.register_material(material, company_code, effective_date);
294 }
295
296 pool
297 }
298
299 fn generate_assets(
301 &mut self,
302 company_code: &str,
303 count: usize,
304 date_range: (NaiveDate, NaiveDate),
305 ) -> FixedAssetPool {
306 let pool = self
307 .asset_generator
308 .generate_diverse_pool(count, company_code, date_range);
309
310 for asset in &pool.assets {
312 self.register_asset(asset, company_code, asset.acquisition_date);
313 }
314
315 pool
316 }
317
318 fn generate_employees(
320 &mut self,
321 company_code: &str,
322 date_range: (NaiveDate, NaiveDate),
323 ) -> EmployeePool {
324 let pool = self
325 .employee_generator
326 .generate_company_pool(company_code, date_range);
327
328 for employee in &pool.employees {
330 if let Some(hire_date) = employee.hire_date {
331 self.register_employee(employee, company_code, hire_date);
332 }
333 }
334
335 pool
336 }
337
338 fn register_vendor(&mut self, vendor: &Vendor, company_code: &str, effective_date: NaiveDate) {
340 let entity_id = EntityId::vendor(&vendor.vendor_id);
341 let record =
342 datasynth_core::models::EntityRecord::new(entity_id, &vendor.name, effective_date)
343 .with_company_code(company_code);
344 self.registry.register(record);
345 }
346
347 fn register_customer(
349 &mut self,
350 customer: &datasynth_core::models::Customer,
351 company_code: &str,
352 effective_date: NaiveDate,
353 ) {
354 let entity_id = EntityId::customer(&customer.customer_id);
355 let record =
356 datasynth_core::models::EntityRecord::new(entity_id, &customer.name, effective_date)
357 .with_company_code(company_code);
358 self.registry.register(record);
359 }
360
361 fn register_material(
363 &mut self,
364 material: &Material,
365 company_code: &str,
366 effective_date: NaiveDate,
367 ) {
368 let entity_id = EntityId::material(&material.material_id);
369 let record = datasynth_core::models::EntityRecord::new(
370 entity_id,
371 &material.description,
372 effective_date,
373 )
374 .with_company_code(company_code);
375 self.registry.register(record);
376 }
377
378 fn register_asset(
380 &mut self,
381 asset: &FixedAsset,
382 company_code: &str,
383 effective_date: NaiveDate,
384 ) {
385 let entity_id = EntityId::fixed_asset(&asset.asset_id);
386 let record = datasynth_core::models::EntityRecord::new(
387 entity_id,
388 &asset.description,
389 effective_date,
390 )
391 .with_company_code(company_code);
392 self.registry.register(record);
393 }
394
395 fn register_employee(
397 &mut self,
398 employee: &Employee,
399 company_code: &str,
400 effective_date: NaiveDate,
401 ) {
402 let entity_id = EntityId::employee(&employee.employee_id);
403 let record = datasynth_core::models::EntityRecord::new(
404 entity_id,
405 &employee.display_name,
406 effective_date,
407 )
408 .with_company_code(company_code);
409 self.registry.register(record);
410 }
411
412 pub fn registry(&self) -> &EntityRegistry {
414 &self.registry
415 }
416
417 pub fn validate_entity(
419 &self,
420 entity_type: EntityType,
421 entity_id: &str,
422 transaction_date: NaiveDate,
423 ) -> bool {
424 let id = EntityId {
425 entity_type,
426 id: entity_id.to_string(),
427 };
428 self.registry.is_valid_on(&id, transaction_date)
429 }
430
431 pub fn get_active_entities(&self, entity_type: EntityType, date: NaiveDate) -> Vec<EntityId> {
433 self.registry
434 .get_ids_by_type(entity_type)
435 .into_iter()
436 .filter(|id| self.registry.is_valid_on(id, date))
437 .cloned()
438 .collect()
439 }
440
441 pub fn get_random_vendor(
443 &self,
444 company_code: &str,
445 date: NaiveDate,
446 rng: &mut impl rand::Rng,
447 ) -> Option<String> {
448 let vendors = self
449 .registry
450 .get_by_company(company_code)
451 .into_iter()
452 .filter(|rec| rec.entity_id.entity_type == EntityType::Vendor)
453 .filter(|rec| self.registry.is_valid(&rec.entity_id, date))
454 .collect::<Vec<_>>();
455
456 if vendors.is_empty() {
457 None
458 } else {
459 use rand::seq::SliceRandom;
460 vendors.choose(rng).map(|rec| rec.entity_id.id.clone())
461 }
462 }
463
464 pub fn get_random_customer(
466 &self,
467 company_code: &str,
468 date: NaiveDate,
469 rng: &mut impl rand::Rng,
470 ) -> Option<String> {
471 let customers = self
472 .registry
473 .get_by_company(company_code)
474 .into_iter()
475 .filter(|rec| rec.entity_id.entity_type == EntityType::Customer)
476 .filter(|rec| self.registry.is_valid(&rec.entity_id, date))
477 .collect::<Vec<_>>();
478
479 if customers.is_empty() {
480 None
481 } else {
482 use rand::seq::SliceRandom;
483 customers.choose(rng).map(|rec| rec.entity_id.id.clone())
484 }
485 }
486
487 pub fn reset(&mut self) {
489 self.vendor_generator.reset();
490 self.customer_generator.reset();
491 self.material_generator.reset();
492 self.asset_generator.reset();
493 self.employee_generator.reset();
494 self.registry = EntityRegistry::new();
495 }
496}
497
498#[cfg(test)]
499#[allow(clippy::unwrap_used)]
500mod tests {
501 use super::*;
502
503 #[test]
504 fn test_manager_creation() {
505 let manager = EntityRegistryManager::new(42);
506 assert_eq!(manager.registry().total_count(), 0);
507 }
508
509 #[test]
510 fn test_master_data_generation() {
511 let mut manager = EntityRegistryManager::new(42);
512 let counts = MasterDataCounts {
513 vendors: 10,
514 customers: 20,
515 materials: 50,
516 assets: 15,
517 employees: None,
518 };
519
520 let data = manager.generate_master_data(
521 "1000",
522 &counts,
523 NaiveDate::from_ymd_opt(2024, 1, 1).unwrap(),
524 (
525 NaiveDate::from_ymd_opt(2020, 1, 1).unwrap(),
526 NaiveDate::from_ymd_opt(2024, 1, 1).unwrap(),
527 ),
528 );
529
530 assert_eq!(data.vendors.vendors.len(), 10);
531 assert_eq!(data.customers.customers.len(), 20);
532 assert_eq!(data.materials.materials.len(), 50);
533 assert_eq!(data.assets.assets.len(), 15);
534 assert!(!data.employees.employees.is_empty());
535
536 assert!(data.registry.total_count() > 0);
538 }
539
540 #[test]
541 fn test_entity_validation() {
542 let mut manager = EntityRegistryManager::new(42);
543 let counts = MasterDataCounts {
544 vendors: 5,
545 ..Default::default()
546 };
547
548 let data = manager.generate_master_data(
549 "1000",
550 &counts,
551 NaiveDate::from_ymd_opt(2024, 1, 1).unwrap(),
552 (
553 NaiveDate::from_ymd_opt(2020, 1, 1).unwrap(),
554 NaiveDate::from_ymd_opt(2024, 1, 1).unwrap(),
555 ),
556 );
557
558 let vendor_id = &data.vendors.vendors[0].vendor_id;
560 assert!(manager.validate_entity(
561 EntityType::Vendor,
562 vendor_id,
563 NaiveDate::from_ymd_opt(2024, 1, 1).unwrap(),
564 ));
565
566 assert!(!manager.validate_entity(
568 EntityType::Vendor,
569 "V-NONEXISTENT",
570 NaiveDate::from_ymd_opt(2024, 1, 1).unwrap(),
571 ));
572 }
573
574 #[test]
575 fn test_multi_company_generation() {
576 let mut manager = EntityRegistryManager::new(42);
577 let counts = MasterDataCounts {
578 vendors: 5,
579 customers: 10,
580 materials: 20,
581 assets: 5,
582 employees: None,
583 };
584
585 let results = manager.generate_multi_company_master_data(
586 &["1000".to_string(), "2000".to_string()],
587 &counts,
588 NaiveDate::from_ymd_opt(2024, 1, 1).unwrap(),
589 (
590 NaiveDate::from_ymd_opt(2020, 1, 1).unwrap(),
591 NaiveDate::from_ymd_opt(2024, 1, 1).unwrap(),
592 ),
593 );
594
595 assert_eq!(results.len(), 2);
596 }
597
598 #[test]
599 fn test_intercompany_generation() {
600 let mut manager = EntityRegistryManager::new(42);
601 let counts = MasterDataCounts {
602 vendors: 10,
603 customers: 15,
604 materials: 20,
605 assets: 5,
606 employees: None,
607 };
608
609 let results = manager.generate_master_data_with_ic(
610 &["1000".to_string(), "2000".to_string()],
611 &counts,
612 NaiveDate::from_ymd_opt(2024, 1, 1).unwrap(),
613 (
614 NaiveDate::from_ymd_opt(2020, 1, 1).unwrap(),
615 NaiveDate::from_ymd_opt(2024, 1, 1).unwrap(),
616 ),
617 );
618
619 let ic_vendors: Vec<_> = results[0]
621 .vendors
622 .vendors
623 .iter()
624 .filter(|v| v.is_intercompany)
625 .collect();
626 assert!(!ic_vendors.is_empty());
627 }
628
629 #[test]
630 fn test_get_random_vendor() {
631 use datasynth_core::utils::seeded_rng;
632
633 let mut manager = EntityRegistryManager::new(42);
634 let counts = MasterDataCounts {
635 vendors: 10,
636 ..Default::default()
637 };
638
639 let _data = manager.generate_master_data(
640 "1000",
641 &counts,
642 NaiveDate::from_ymd_opt(2024, 1, 1).unwrap(),
643 (
644 NaiveDate::from_ymd_opt(2020, 1, 1).unwrap(),
645 NaiveDate::from_ymd_opt(2024, 1, 1).unwrap(),
646 ),
647 );
648
649 let mut rng = seeded_rng(42, 0);
650 let vendor = manager.get_random_vendor(
651 "1000",
652 NaiveDate::from_ymd_opt(2024, 1, 1).unwrap(),
653 &mut rng,
654 );
655
656 assert!(vendor.is_some());
657 }
658}