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