1use chrono::NaiveDate;
8use datasynth_core::models::{BenefitEnrollment, BenefitPlanType, BenefitStatus};
9use datasynth_core::utils::seeded_rng;
10use datasynth_core::uuid_factory::{DeterministicUuidFactory, GeneratorType};
11use rand::prelude::*;
12use rand_chacha::ChaCha8Rng;
13use rust_decimal::Decimal;
14use tracing::debug;
15
16const HEALTH_PLANS: &[&str] = &["Blue Cross PPO", "Aetna HMO", "UnitedHealth Choice Plus"];
18const DENTAL_PLANS: &[&str] = &["Delta Dental Basic", "MetLife Dental PPO"];
19const VISION_PLANS: &[&str] = &["VSP Standard", "EyeMed Vision Care"];
20const RETIREMENT_PLANS: &[&str] = &["401(k) Traditional", "401(k) Roth"];
21const LIFE_PLANS: &[&str] = &["Basic Life 1x Salary", "Supplemental Life 2x"];
22
23pub struct BenefitEnrollmentGenerator {
25 rng: ChaCha8Rng,
26 uuid_factory: DeterministicUuidFactory,
27}
28
29impl BenefitEnrollmentGenerator {
30 pub fn new(seed: u64) -> Self {
32 Self {
33 rng: seeded_rng(seed, 0),
34 uuid_factory: DeterministicUuidFactory::new(seed, GeneratorType::BenefitEnrollment),
35 }
36 }
37
38 pub fn generate(
51 &mut self,
52 company_code: &str,
53 employees: &[(String, String)],
54 enrollment_date: NaiveDate,
55 currency: &str,
56 ) -> Vec<BenefitEnrollment> {
57 debug!(
58 company_code,
59 employee_count = employees.len(),
60 %enrollment_date,
61 "Generating benefit enrollments"
62 );
63
64 let mut enrollments = Vec::new();
65 let effective_date = enrollment_date;
66 let period = format!(
67 "{}-{:02}",
68 enrollment_date.format("%Y"),
69 enrollment_date.format("%m")
70 );
71
72 for (employee_id, employee_name) in employees {
73 if self.rng.random_bool(0.90) {
75 let enrollment = self.make_enrollment(
76 company_code,
77 employee_id,
78 employee_name,
79 BenefitPlanType::Health,
80 HEALTH_PLANS,
81 enrollment_date,
82 effective_date,
83 &period,
84 currency,
85 );
86 enrollments.push(enrollment);
87 }
88
89 if self.rng.random_bool(0.70) {
91 let enrollment = self.make_enrollment(
92 company_code,
93 employee_id,
94 employee_name,
95 BenefitPlanType::Dental,
96 DENTAL_PLANS,
97 enrollment_date,
98 effective_date,
99 &period,
100 currency,
101 );
102 enrollments.push(enrollment);
103 }
104
105 if self.rng.random_bool(0.50) {
107 let enrollment = self.make_enrollment(
108 company_code,
109 employee_id,
110 employee_name,
111 BenefitPlanType::Vision,
112 VISION_PLANS,
113 enrollment_date,
114 effective_date,
115 &period,
116 currency,
117 );
118 enrollments.push(enrollment);
119 }
120
121 if self.rng.random_bool(0.60) {
123 let enrollment = self.make_enrollment(
124 company_code,
125 employee_id,
126 employee_name,
127 BenefitPlanType::Retirement401k,
128 RETIREMENT_PLANS,
129 enrollment_date,
130 effective_date,
131 &period,
132 currency,
133 );
134 enrollments.push(enrollment);
135 }
136
137 if self.rng.random_bool(0.40) {
139 let enrollment = self.make_enrollment(
140 company_code,
141 employee_id,
142 employee_name,
143 BenefitPlanType::LifeInsurance,
144 LIFE_PLANS,
145 enrollment_date,
146 effective_date,
147 &period,
148 currency,
149 );
150 enrollments.push(enrollment);
151 }
152 }
153
154 enrollments
155 }
156
157 #[allow(clippy::too_many_arguments)]
159 fn make_enrollment(
160 &mut self,
161 company_code: &str,
162 employee_id: &str,
163 employee_name: &str,
164 plan_type: BenefitPlanType,
165 plan_names: &[&str],
166 enrollment_date: NaiveDate,
167 effective_date: NaiveDate,
168 period: &str,
169 currency: &str,
170 ) -> BenefitEnrollment {
171 let id = self.uuid_factory.next().to_string();
172 let plan_name = plan_names[self.rng.random_range(0..plan_names.len())].to_string();
173
174 let (employee_contrib, employer_contrib) = self.contribution_amounts(plan_type);
175
176 let status_roll: f64 = self.rng.random();
178 let (status, is_active) = if status_roll < 0.95 {
179 (BenefitStatus::Active, true)
180 } else if status_roll < 0.98 {
181 (BenefitStatus::Pending, false)
182 } else {
183 (BenefitStatus::Terminated, false)
184 };
185
186 BenefitEnrollment::new(
187 id,
188 company_code,
189 employee_id,
190 employee_name,
191 plan_type,
192 plan_name,
193 enrollment_date,
194 effective_date,
195 period,
196 employee_contrib,
197 employer_contrib,
198 currency,
199 status,
200 is_active,
201 )
202 }
203
204 fn contribution_amounts(&mut self, plan_type: BenefitPlanType) -> (Decimal, Decimal) {
206 let (emp_min, emp_max, er_min, er_max) = match plan_type {
207 BenefitPlanType::Health => (200.0, 800.0, 400.0, 1200.0),
208 BenefitPlanType::Dental => (25.0, 75.0, 30.0, 90.0),
209 BenefitPlanType::Vision => (10.0, 30.0, 10.0, 30.0),
210 BenefitPlanType::Retirement401k => (200.0, 2000.0, 100.0, 1000.0),
211 BenefitPlanType::LifeInsurance => (15.0, 50.0, 15.0, 50.0),
212 BenefitPlanType::StockPurchase => (100.0, 500.0, 0.0, 0.0),
213 BenefitPlanType::Disability => (20.0, 60.0, 20.0, 60.0),
214 };
215
216 let emp: f64 = self.rng.random_range(emp_min..=emp_max);
217 let er: f64 = self.rng.random_range(er_min..=er_max);
218
219 let employee_contribution = Decimal::from_f64_retain(emp)
220 .unwrap_or(Decimal::from(100))
221 .round_dp(2);
222 let employer_contribution = Decimal::from_f64_retain(er)
223 .unwrap_or(Decimal::from(100))
224 .round_dp(2);
225
226 (employee_contribution, employer_contribution)
227 }
228}
229
230#[cfg(test)]
235#[allow(clippy::unwrap_used)]
236mod tests {
237 use super::*;
238
239 fn test_employees() -> Vec<(String, String)> {
240 vec![
241 ("EMP-001".to_string(), "Alice Smith".to_string()),
242 ("EMP-002".to_string(), "Bob Jones".to_string()),
243 ("EMP-003".to_string(), "Carol White".to_string()),
244 ("EMP-004".to_string(), "David Brown".to_string()),
245 ("EMP-005".to_string(), "Eve Johnson".to_string()),
246 ]
247 }
248
249 #[test]
250 fn test_enrollment_generation() {
251 let mut gen = BenefitEnrollmentGenerator::new(42);
252 let employees = test_employees();
253 let date = NaiveDate::from_ymd_opt(2024, 1, 1).unwrap();
254
255 let enrollments = gen.generate("C001", &employees, date, "USD");
256
257 assert!(!enrollments.is_empty());
258 for e in &enrollments {
259 assert_eq!(e.entity_code, "C001");
260 assert_eq!(e.currency, "USD");
261 assert!(e.employee_contribution > Decimal::ZERO);
262 assert!(!e.employee_name.is_empty());
263 assert!(!e.plan_name.is_empty());
264 }
265 }
266
267 #[test]
268 fn test_enrollment_plan_types() {
269 let mut gen = BenefitEnrollmentGenerator::new(77);
270 let employees: Vec<(String, String)> = (0..50)
271 .map(|i| (format!("EMP-{:03}", i), format!("Employee {}", i)))
272 .collect();
273 let date = NaiveDate::from_ymd_opt(2024, 1, 1).unwrap();
274
275 let enrollments = gen.generate("C001", &employees, date, "USD");
276
277 let health = enrollments
278 .iter()
279 .filter(|e| matches!(e.plan_type, BenefitPlanType::Health))
280 .count();
281 let dental = enrollments
282 .iter()
283 .filter(|e| matches!(e.plan_type, BenefitPlanType::Dental))
284 .count();
285
286 assert!(
288 health >= 30,
289 "Expected ~90% health enrollment, got {}/50",
290 health
291 );
292 assert!(dental > 0, "Expected some dental enrollments");
293 }
294
295 #[test]
296 fn test_enrollment_deterministic() {
297 let employees = test_employees();
298 let date = NaiveDate::from_ymd_opt(2024, 1, 1).unwrap();
299
300 let mut gen1 = BenefitEnrollmentGenerator::new(12345);
301 let e1 = gen1.generate("C001", &employees, date, "USD");
302 let mut gen2 = BenefitEnrollmentGenerator::new(12345);
303 let e2 = gen2.generate("C001", &employees, date, "USD");
304
305 assert_eq!(e1.len(), e2.len());
306 for (a, b) in e1.iter().zip(e2.iter()) {
307 assert_eq!(a.id, b.id);
308 assert_eq!(a.employee_id, b.employee_id);
309 assert_eq!(a.plan_type, b.plan_type);
310 assert_eq!(a.employee_contribution, b.employee_contribution);
311 }
312 }
313
314 #[test]
315 fn test_enrollment_active_rate() {
316 let mut gen = BenefitEnrollmentGenerator::new(55);
317 let employees: Vec<(String, String)> = (0..100)
318 .map(|i| (format!("EMP-{:03}", i), format!("Employee {}", i)))
319 .collect();
320 let date = NaiveDate::from_ymd_opt(2024, 1, 1).unwrap();
321
322 let enrollments = gen.generate("C001", &employees, date, "USD");
323 let active_count = enrollments.iter().filter(|e| e.is_active).count();
324 let active_pct = active_count as f64 / enrollments.len() as f64;
325
326 assert!(
327 active_pct > 0.85,
328 "Expected ~95% active rate, got {:.1}%",
329 active_pct * 100.0
330 );
331 }
332}