1use serde_json::Value;
2use std::collections::HashMap;
3use std::time::{SystemTime, UNIX_EPOCH};
4
5use crate::assigner::VariantAssigner;
6use crate::matcher::AudienceMatcher;
7use crate::models::*;
8use crate::utils::{array_equals_shallow, hash_unit};
9
10pub type EventLogger = Box<dyn Fn(&Context, &str, Option<Value>) + Send + Sync>;
11
12struct Experiment {
13 data: ExperimentData,
14 variables: Vec<HashMap<String, Value>>,
15}
16
17pub struct Context {
18 units: HashMap<String, String>,
19 attrs: Vec<Attribute>,
20 data: ContextData,
21 assignments: HashMap<String, Assignment>,
22 exposures: Vec<Exposure>,
23 goals: Vec<Goal>,
24 overrides: HashMap<String, i32>,
25 cassignments: HashMap<String, i32>,
26 state: ContextState,
27 pending: usize,
28 attrs_seq: u64,
29 index: HashMap<String, Experiment>,
30 index_variables: HashMap<String, Vec<String>>,
31 assigners: HashMap<String, VariantAssigner>,
32 hashes: HashMap<String, String>,
33 audience_matcher: AudienceMatcher,
34 event_logger: Option<EventLogger>,
35}
36
37impl Context {
38 pub fn new(data: ContextData) -> Self {
39 let mut ctx = Self {
40 units: HashMap::new(),
41 attrs: Vec::new(),
42 data: ContextData::default(),
43 assignments: HashMap::new(),
44 exposures: Vec::new(),
45 goals: Vec::new(),
46 overrides: HashMap::new(),
47 cassignments: HashMap::new(),
48 state: ContextState::Loading,
49 pending: 0,
50 attrs_seq: 0,
51 index: HashMap::new(),
52 index_variables: HashMap::new(),
53 assigners: HashMap::new(),
54 hashes: HashMap::new(),
55 audience_matcher: AudienceMatcher::new(),
56 event_logger: None,
57 };
58 ctx.init(data);
59 ctx.state = ContextState::Ready;
60 ctx
61 }
62
63 pub fn set_event_logger(&mut self, logger: EventLogger) {
64 self.event_logger = Some(logger);
65 }
66
67 fn init(&mut self, data: ContextData) {
68 self.data = data;
69 self.index.clear();
70 self.index_variables.clear();
71
72 for experiment in &self.data.experiments {
73 let mut variables: Vec<HashMap<String, Value>> = Vec::new();
74
75 for variant in &experiment.variants {
76 let parsed: HashMap<String, Value> = variant
77 .config
78 .as_ref()
79 .and_then(|c| {
80 if c.is_empty() {
81 None
82 } else {
83 serde_json::from_str(c).ok()
84 }
85 })
86 .unwrap_or_default();
87
88 for key in parsed.keys() {
89 self.index_variables
90 .entry(key.clone())
91 .or_default()
92 .push(experiment.name.clone());
93 }
94
95 variables.push(parsed);
96 }
97
98 self.index.insert(
99 experiment.name.clone(),
100 Experiment {
101 data: experiment.clone(),
102 variables,
103 },
104 );
105 }
106 }
107
108 pub fn is_ready(&self) -> bool {
109 self.state == ContextState::Ready
110 }
111
112 pub fn is_failed(&self) -> bool {
113 self.state == ContextState::Failed
114 }
115
116 pub fn is_finalized(&self) -> bool {
117 self.state == ContextState::Finalized
118 }
119
120 pub fn is_finalizing(&self) -> bool {
121 self.state == ContextState::Finalizing
122 }
123
124 pub fn pending(&self) -> usize {
125 self.pending
126 }
127
128 pub fn data(&self) -> &ContextData {
129 &self.data
130 }
131
132 pub fn set_unit(&mut self, unit_type: &str, uid: &str) -> Result<(), String> {
133 if self.is_finalized() {
134 return Err("ABSmartly Context is finalized.".to_string());
135 }
136 if self.is_finalizing() {
137 return Err("ABSmartly Context is finalizing.".to_string());
138 }
139
140 let uid = uid.trim();
141 if uid.is_empty() {
142 return Err(format!("Unit '{}' UID must not be blank.", unit_type));
143 }
144
145 if let Some(existing) = self.units.get(unit_type) {
146 if existing != uid {
147 return Err(format!("Unit '{}' UID already set.", unit_type));
148 }
149 }
150
151 self.units.insert(unit_type.to_string(), uid.to_string());
152 Ok(())
153 }
154
155 pub fn get_unit(&self, unit_type: &str) -> Option<&String> {
156 self.units.get(unit_type)
157 }
158
159 pub fn get_units(&self) -> &HashMap<String, String> {
160 &self.units
161 }
162
163 pub fn set_attribute(&mut self, name: &str, value: impl Into<Value>) -> Result<(), String> {
164 if self.is_finalized() {
165 return Err("ABSmartly Context is finalized.".to_string());
166 }
167 if self.is_finalizing() {
168 return Err("ABSmartly Context is finalizing.".to_string());
169 }
170
171 self.attrs.push(Attribute {
172 name: name.to_string(),
173 value: value.into(),
174 set_at: now_millis(),
175 });
176 self.attrs_seq += 1;
177 Ok(())
178 }
179
180 pub fn get_attribute(&self, name: &str) -> Option<&Value> {
181 self.attrs
182 .iter()
183 .rev()
184 .find(|a| a.name == name)
185 .map(|a| &a.value)
186 }
187
188 pub fn get_attributes(&self) -> HashMap<String, Value> {
189 let mut attrs = HashMap::new();
190 for attr in &self.attrs {
191 attrs.insert(attr.name.clone(), attr.value.clone());
192 }
193 attrs
194 }
195
196 pub fn set_override(&mut self, experiment_name: &str, variant: i32) {
197 self.overrides.insert(experiment_name.to_string(), variant);
198 }
199
200 pub fn set_custom_assignment(&mut self, experiment_name: &str, variant: i32) -> Result<(), String> {
201 if self.is_finalized() {
202 return Err("ABSmartly Context is finalized.".to_string());
203 }
204 if self.is_finalizing() {
205 return Err("ABSmartly Context is finalizing.".to_string());
206 }
207 self.cassignments
208 .insert(experiment_name.to_string(), variant);
209 Ok(())
210 }
211
212 pub fn peek(&mut self, experiment_name: &str) -> i32 {
213 self.assign(experiment_name).variant
214 }
215
216 pub fn treatment(&mut self, experiment_name: &str) -> i32 {
217 let assignment = self.assign(experiment_name);
218 let variant = assignment.variant;
219
220 if !assignment.exposed {
221 if let Some(assignment) = self.assignments.get_mut(experiment_name) {
222 assignment.exposed = true;
223 }
224 self.queue_exposure(experiment_name);
225 }
226
227 variant
228 }
229
230 pub fn track(&mut self, goal_name: &str, properties: impl Into<Value>) -> Result<(), String> {
231 if self.is_finalized() {
232 return Err("ABSmartly Context is finalized.".to_string());
233 }
234 if self.is_finalizing() {
235 return Err("ABSmartly Context is finalizing.".to_string());
236 }
237
238 let properties_map: Option<HashMap<String, Value>> = match properties.into() {
239 Value::Object(map) => Some(map.into_iter().collect()),
240 Value::Null => None,
241 _ => None,
242 };
243
244 let goal = Goal {
245 name: goal_name.to_string(),
246 properties: properties_map,
247 achieved_at: now_millis(),
248 };
249
250 self.log_event("goal", Some(serde_json::to_value(&goal).unwrap_or_default()));
251 self.goals.push(goal);
252 self.pending += 1;
253
254 Ok(())
255 }
256
257 pub fn variable_value(&mut self, key: &str, default_value: impl Into<Value>) -> Value {
258 if let Some(experiment_names) = self.index_variables.get(key).cloned() {
259 for exp_name in experiment_names {
260 let assignment = self.assign(&exp_name);
261 if let Some(variables) = &assignment.variables {
262 if !assignment.exposed {
263 if let Some(a) = self.assignments.get_mut(&exp_name) {
264 a.exposed = true;
265 }
266 self.queue_exposure(&exp_name);
267 }
268
269 if let Some(value) = variables.get(key) {
270 if assignment.assigned || assignment.overridden {
271 return value.clone();
272 }
273 }
274 }
275 }
276 }
277 default_value.into()
278 }
279
280 pub fn peek_variable_value(&mut self, key: &str, default_value: impl Into<Value>) -> Value {
281 if let Some(experiment_names) = self.index_variables.get(key).cloned() {
282 for exp_name in experiment_names {
283 let assignment = self.assign(&exp_name);
284 if let Some(variables) = &assignment.variables {
285 if let Some(value) = variables.get(key) {
286 if assignment.assigned || assignment.overridden {
287 return value.clone();
288 }
289 }
290 }
291 }
292 }
293 default_value.into()
294 }
295
296 pub fn variable_keys(&self) -> HashMap<String, Vec<String>> {
297 let mut result = HashMap::new();
298 for (key, exp_names) in &self.index_variables {
299 result.insert(key.clone(), exp_names.clone());
300 }
301 result
302 }
303
304 pub fn custom_field_value(&self, experiment_name: &str, field_name: &str) -> Option<Value> {
305 if let Some(exp) = self.index.get(experiment_name) {
306 if let Some(ref custom_fields) = exp.data.custom_field_values {
307 if let Some(field) = custom_fields.iter().find(|f| f.name == field_name) {
308 return match field.field_type.as_str() {
309 "text" | "string" => Some(Value::String(field.value.clone())),
310 "number" => field.value.parse::<f64>().ok().map(|n| {
311 serde_json::Number::from_f64(n)
312 .map(Value::Number)
313 .unwrap_or(Value::Null)
314 }),
315 "json" => {
316 if field.value == "null" {
317 Some(Value::Null)
318 } else if field.value.is_empty() {
319 Some(Value::String(String::new()))
320 } else {
321 serde_json::from_str(&field.value).ok()
322 }
323 }
324 "boolean" => Some(Value::Bool(field.value == "true")),
325 _ => None,
326 };
327 }
328 }
329 }
330 None
331 }
332
333 pub fn custom_field_value_type(&self, experiment_name: &str, field_name: &str) -> Option<String> {
334 if let Some(exp) = self.index.get(experiment_name) {
335 if let Some(ref custom_fields) = exp.data.custom_field_values {
336 if let Some(field) = custom_fields.iter().find(|f| f.name == field_name) {
337 return Some(field.field_type.clone());
338 }
339 }
340 }
341 None
342 }
343
344 pub fn custom_field_keys(&self) -> Vec<String> {
345 let mut keys = std::collections::HashSet::new();
346 for exp in &self.data.experiments {
347 if let Some(ref custom_fields) = exp.custom_field_values {
348 for field in custom_fields {
349 keys.insert(field.name.clone());
350 }
351 }
352 }
353 keys.into_iter().collect()
354 }
355
356 pub fn experiments(&self) -> Vec<String> {
357 self.data.experiments.iter().map(|e| e.name.clone()).collect()
358 }
359
360 pub fn refresh(&mut self, new_data: ContextData) {
361 self.assignments.clear();
362 self.init(new_data);
363 self.log_event("refresh", Some(serde_json::to_value(&self.data).unwrap_or_default()));
364 }
365
366 pub fn publish(&mut self) {
367 if self.pending == 0 {
368 return;
369 }
370
371 let params = self.build_publish_params();
372 self.log_event("publish", Some(serde_json::to_value(¶ms).unwrap_or_default()));
373
374 self.pending = 0;
375 self.exposures.clear();
376 self.goals.clear();
377 }
378
379 pub fn finalize(&mut self) {
380 if self.is_finalized() {
381 return;
382 }
383
384 self.state = ContextState::Finalizing;
385
386 if self.pending > 0 {
387 self.publish();
388 }
389
390 self.state = ContextState::Finalized;
391 self.log_event("finalize", None);
392 }
393
394 fn assign(&mut self, experiment_name: &str) -> Assignment {
395 let has_custom = self.cassignments.contains_key(experiment_name);
396 let has_override = self.overrides.contains_key(experiment_name);
397 let has_experiment = self.index.contains_key(experiment_name);
398
399 if let Some(cached) = self.assignments.get(experiment_name) {
400 if has_override {
401 if cached.overridden && cached.variant == self.overrides[experiment_name] {
402 return cached.clone();
403 }
404 } else if !has_experiment {
405 if !cached.assigned {
406 return cached.clone();
407 }
408 } else if !has_custom || self.cassignments[experiment_name] == cached.variant {
409 if let Some(exp) = self.index.get(experiment_name) {
410 if self.experiment_matches(&exp.data, cached)
411 && self.audience_matches(&exp.data, cached)
412 {
413 return cached.clone();
414 }
415 }
416 }
417 }
418
419 let exp_data_opt = self.index.get(experiment_name).map(|e| e.data.clone());
420
421 let mut assignment = Assignment {
422 eligible: true,
423 ..Default::default()
424 };
425
426 if has_override {
427 if let Some(ref exp_data) = exp_data_opt {
428 assignment.id = exp_data.id;
429 assignment.unit_type = exp_data.unit_type.clone();
430 }
431
432 assignment.overridden = true;
433 assignment.variant = self.overrides[experiment_name];
434 } else if let Some(ref exp_data) = exp_data_opt {
435 if !exp_data.audience.is_empty() {
436 let attrs = self.get_attributes();
437 let result = self.audience_matcher.evaluate(&exp_data.audience, &attrs);
438 if let Some(matched) = result {
439 assignment.audience_mismatch = !matched;
440 }
441 }
442
443 if exp_data.audience_strict && assignment.audience_mismatch {
444 assignment.variant = 0;
445 } else if exp_data.full_on_variant == 0 {
446 if let Some(ref unit_type) = exp_data.unit_type {
447 if self.units.contains_key(unit_type) {
448 let unit_hash = self.unit_hash(unit_type);
449
450 if let Some(ref hash) = unit_hash {
451 let assigner = self
452 .assigners
453 .entry(unit_type.clone())
454 .or_insert_with(|| VariantAssigner::new(hash));
455
456 let eligible = assigner.assign(
457 &exp_data.traffic_split,
458 exp_data.traffic_seed_hi,
459 exp_data.traffic_seed_lo,
460 ) == 1;
461
462 assignment.assigned = true;
463 assignment.eligible = eligible;
464
465 if eligible {
466 if has_custom {
467 assignment.variant = self.cassignments[experiment_name];
468 assignment.custom = true;
469 } else {
470 assignment.variant = assigner.assign(
471 &exp_data.split,
472 exp_data.seed_hi,
473 exp_data.seed_lo,
474 ) as i32;
475 }
476 } else {
477 assignment.variant = 0;
478 }
479 }
480 }
481 }
482 } else {
483 assignment.assigned = true;
484 assignment.eligible = true;
485 assignment.variant = exp_data.full_on_variant as i32;
486 assignment.full_on = true;
487 }
488
489 assignment.unit_type = exp_data.unit_type.clone();
490 assignment.id = exp_data.id;
491 assignment.iteration = exp_data.iteration;
492 assignment.traffic_split = Some(exp_data.traffic_split.clone());
493 assignment.full_on_variant = exp_data.full_on_variant;
494 assignment.attrs_seq = self.attrs_seq;
495 }
496
497 if let Some(exp) = self.index.get(experiment_name) {
498 if (assignment.variant as usize) < exp.variables.len() {
499 assignment.variables = Some(exp.variables[assignment.variant as usize].clone());
500 }
501 }
502
503 self.assignments
504 .insert(experiment_name.to_string(), assignment.clone());
505 assignment
506 }
507
508 fn experiment_matches(&self, experiment: &ExperimentData, assignment: &Assignment) -> bool {
509 experiment.id == assignment.id
510 && experiment.unit_type == assignment.unit_type
511 && experiment.iteration == assignment.iteration
512 && experiment.full_on_variant == assignment.full_on_variant
513 && assignment
514 .traffic_split
515 .as_ref()
516 .map_or(false, |ts| array_equals_shallow(&experiment.traffic_split, ts))
517 }
518
519 fn audience_matches(&self, experiment: &ExperimentData, assignment: &Assignment) -> bool {
520 if !experiment.audience.is_empty() && self.attrs_seq > assignment.attrs_seq {
521 let attrs = self.get_attributes();
522 let result = self.audience_matcher.evaluate(&experiment.audience, &attrs);
523 if let Some(matched) = result {
524 return matched == !assignment.audience_mismatch;
525 }
526 }
527 true
528 }
529
530 fn queue_exposure(&mut self, experiment_name: &str) {
531 if let Some(assignment) = self.assignments.get(experiment_name) {
532 let exposure = Exposure {
533 id: assignment.id,
534 name: experiment_name.to_string(),
535 exposed_at: now_millis(),
536 unit: assignment.unit_type.clone(),
537 variant: assignment.variant,
538 assigned: assignment.assigned,
539 eligible: assignment.eligible,
540 overridden: assignment.overridden,
541 full_on: assignment.full_on,
542 custom: assignment.custom,
543 audience_mismatch: assignment.audience_mismatch,
544 };
545
546 self.log_event("exposure", Some(serde_json::to_value(&exposure).unwrap_or_default()));
547 self.exposures.push(exposure);
548 self.pending += 1;
549 }
550 }
551
552 fn unit_hash(&mut self, unit_type: &str) -> Option<String> {
553 if let Some(hash) = self.hashes.get(unit_type) {
554 return Some(hash.clone());
555 }
556
557 if let Some(unit) = self.units.get(unit_type) {
558 let hash = hash_unit(unit);
559 self.hashes.insert(unit_type.to_string(), hash.clone());
560 return Some(hash);
561 }
562
563 None
564 }
565
566 fn build_publish_params(&self) -> PublishParams {
567 let units: Vec<Unit> = self
568 .units
569 .iter()
570 .map(|(unit_type, _)| Unit {
571 unit_type: unit_type.clone(),
572 uid: self.hashes.get(unit_type).cloned(),
573 })
574 .collect();
575
576 PublishParams {
577 published_at: now_millis(),
578 units,
579 hashed: true,
580 exposures: if self.exposures.is_empty() {
581 None
582 } else {
583 Some(self.exposures.clone())
584 },
585 goals: if self.goals.is_empty() {
586 None
587 } else {
588 Some(self.goals.clone())
589 },
590 attributes: if self.attrs.is_empty() {
591 None
592 } else {
593 Some(self.attrs.clone())
594 },
595 }
596 }
597
598 fn log_event(&self, event_name: &str, data: Option<Value>) {
599 if let Some(ref logger) = self.event_logger {
600 logger(self, event_name, data);
601 }
602 }
603}
604
605fn now_millis() -> i64 {
606 SystemTime::now()
607 .duration_since(UNIX_EPOCH)
608 .map(|d| d.as_millis() as i64)
609 .unwrap_or(0)
610}
611
612#[cfg(test)]
613mod tests {
614 use super::*;
615 use serde_json::json;
616
617 fn make_experiment(name: &str, variants: Vec<&str>, split: Vec<f64>) -> ExperimentData {
618 ExperimentData {
619 id: 1,
620 name: name.to_string(),
621 unit_type: Some("session_id".to_string()),
622 iteration: 1,
623 seed_hi: 0,
624 seed_lo: 0,
625 split,
626 traffic_seed_hi: 0,
627 traffic_seed_lo: 0,
628 traffic_split: vec![0.0, 1.0],
629 full_on_variant: 0,
630 audience: String::new(),
631 audience_strict: false,
632 variants: variants
633 .iter()
634 .map(|c| Variant {
635 config: if c.is_empty() { None } else { Some(c.to_string()) },
636 })
637 .collect(),
638 variables: HashMap::new(),
639 custom_field_values: None,
640 }
641 }
642
643 fn make_context_data(experiments: Vec<ExperimentData>) -> ContextData {
644 ContextData { experiments }
645 }
646
647 #[test]
648 fn test_context_is_ready_after_creation() {
649 let data = make_context_data(vec![]);
650 let context = Context::new(data);
651 assert!(context.is_ready());
652 assert!(!context.is_failed());
653 assert!(!context.is_finalized());
654 }
655
656 #[test]
657 fn test_context_set_unit() {
658 let data = make_context_data(vec![]);
659 let mut context = Context::new(data);
660
661 assert!(context.set_unit("session_id", "user123").is_ok());
662 assert_eq!(context.get_unit("session_id"), Some(&"user123".to_string()));
663 }
664
665 #[test]
666 fn test_context_set_unit_cannot_change() {
667 let data = make_context_data(vec![]);
668 let mut context = Context::new(data);
669
670 assert!(context.set_unit("session_id", "user123").is_ok());
671 assert!(context.set_unit("session_id", "user456").is_err());
672 }
673
674 #[test]
675 fn test_context_set_unit_same_value_ok() {
676 let data = make_context_data(vec![]);
677 let mut context = Context::new(data);
678
679 assert!(context.set_unit("session_id", "user123").is_ok());
680 assert!(context.set_unit("session_id", "user123").is_ok());
681 }
682
683 #[test]
684 fn test_context_set_unit_blank_not_allowed() {
685 let data = make_context_data(vec![]);
686 let mut context = Context::new(data);
687
688 assert!(context.set_unit("session_id", "").is_err());
689 assert!(context.set_unit("session_id", " ").is_err());
690 }
691
692 #[test]
693 fn test_context_set_attribute() {
694 let data = make_context_data(vec![]);
695 let mut context = Context::new(data);
696
697 assert!(context.set_attribute("country", json!("US")).is_ok());
698 assert_eq!(context.get_attribute("country"), Some(&json!("US")));
699 }
700
701 #[test]
702 fn test_context_peek_returns_zero_for_nonexistent() {
703 let data = make_context_data(vec![]);
704 let mut context = Context::new(data);
705
706 assert_eq!(context.peek("nonexistent_experiment"), 0);
707 }
708
709 #[test]
710 fn test_context_treatment_returns_zero_for_nonexistent() {
711 let data = make_context_data(vec![]);
712 let mut context = Context::new(data);
713
714 assert_eq!(context.treatment("nonexistent_experiment"), 0);
715 }
716
717 #[test]
718 fn test_context_treatment_with_experiment() {
719 let exp = make_experiment("test_exp", vec!["{}", r#"{"button":"red"}"#], vec![0.5, 0.5]);
720 let data = make_context_data(vec![exp]);
721 let mut context = Context::new(data);
722
723 context.set_unit("session_id", "test_user").unwrap();
724 let variant = context.treatment("test_exp");
725 assert!(variant == 0 || variant == 1);
726 }
727
728 #[test]
729 fn test_context_peek_does_not_queue_exposure() {
730 let exp = make_experiment("test_exp", vec!["{}", r#"{"button":"red"}"#], vec![0.5, 0.5]);
731 let data = make_context_data(vec![exp]);
732 let mut context = Context::new(data);
733
734 context.set_unit("session_id", "test_user").unwrap();
735 context.peek("test_exp");
736 assert_eq!(context.pending(), 0);
737 }
738
739 #[test]
740 fn test_context_treatment_queues_exposure() {
741 let exp = make_experiment("test_exp", vec!["{}", r#"{"button":"red"}"#], vec![0.5, 0.5]);
742 let data = make_context_data(vec![exp]);
743 let mut context = Context::new(data);
744
745 context.set_unit("session_id", "test_user").unwrap();
746 context.treatment("test_exp");
747 assert_eq!(context.pending(), 1);
748 }
749
750 #[test]
751 fn test_context_treatment_only_queues_once() {
752 let exp = make_experiment("test_exp", vec!["{}", r#"{"button":"red"}"#], vec![0.5, 0.5]);
753 let data = make_context_data(vec![exp]);
754 let mut context = Context::new(data);
755
756 context.set_unit("session_id", "test_user").unwrap();
757 context.treatment("test_exp");
758 context.treatment("test_exp");
759 context.treatment("test_exp");
760 assert_eq!(context.pending(), 1);
761 }
762
763 #[test]
764 fn test_context_set_override() {
765 let exp = make_experiment("test_exp", vec!["{}", r#"{"button":"red"}"#], vec![0.5, 0.5]);
766 let data = make_context_data(vec![exp]);
767 let mut context = Context::new(data);
768
769 context.set_override("test_exp", 1);
770 context.set_unit("session_id", "test_user").unwrap();
771
772 assert_eq!(context.treatment("test_exp"), 1);
773 }
774
775 #[test]
776 fn test_context_set_custom_assignment() {
777 let exp = make_experiment("test_exp", vec!["{}", r#"{"button":"red"}"#], vec![0.5, 0.5]);
778 let data = make_context_data(vec![exp]);
779 let mut context = Context::new(data);
780
781 context.set_custom_assignment("test_exp", 1).unwrap();
782 context.set_unit("session_id", "test_user").unwrap();
783
784 let variant = context.treatment("test_exp");
785 assert_eq!(variant, 1);
786 }
787
788 #[test]
789 fn test_context_track() {
790 let data = make_context_data(vec![]);
791 let mut context = Context::new(data);
792
793 assert!(context.track("purchase", json!({"amount": 99.99})).is_ok());
794
795 assert_eq!(context.pending(), 1);
796 }
797
798 #[test]
799 fn test_context_track_without_properties() {
800 let data = make_context_data(vec![]);
801 let mut context = Context::new(data);
802
803 assert!(context.track("click", ()).is_ok());
804 assert_eq!(context.pending(), 1);
805 }
806
807 #[test]
808 fn test_context_publish_clears_pending() {
809 let exp = make_experiment("test_exp", vec!["{}", r#"{"button":"red"}"#], vec![0.5, 0.5]);
810 let data = make_context_data(vec![exp]);
811 let mut context = Context::new(data);
812
813 context.set_unit("session_id", "test_user").unwrap();
814 context.treatment("test_exp");
815 context.track("click", ()).unwrap();
816
817 assert_eq!(context.pending(), 2);
818 context.publish();
819 assert_eq!(context.pending(), 0);
820 }
821
822 #[test]
823 fn test_context_finalize() {
824 let data = make_context_data(vec![]);
825 let mut context = Context::new(data);
826
827 context.finalize();
828 assert!(context.is_finalized());
829 }
830
831 #[test]
832 fn test_context_cannot_track_after_finalize() {
833 let data = make_context_data(vec![]);
834 let mut context = Context::new(data);
835
836 context.finalize();
837 assert!(context.track("click", ()).is_err());
838 }
839
840 #[test]
841 fn test_context_cannot_set_unit_after_finalize() {
842 let data = make_context_data(vec![]);
843 let mut context = Context::new(data);
844
845 context.finalize();
846 assert!(context.set_unit("session_id", "user123").is_err());
847 }
848
849 #[test]
850 fn test_context_cannot_set_attribute_after_finalize() {
851 let data = make_context_data(vec![]);
852 let mut context = Context::new(data);
853
854 context.finalize();
855 assert!(context.set_attribute("country", json!("US")).is_err());
856 }
857
858 #[test]
859 fn test_context_variable_value_returns_default_for_nonexistent() {
860 let data = make_context_data(vec![]);
861 let mut context = Context::new(data);
862
863 let value = context.variable_value("nonexistent", json!("default"));
864 assert_eq!(value, json!("default"));
865 }
866
867 #[test]
868 fn test_context_peek_variable_value_returns_default_for_nonexistent() {
869 let data = make_context_data(vec![]);
870 let mut context = Context::new(data);
871
872 let value = context.peek_variable_value("nonexistent", json!(42));
873 assert_eq!(value, json!(42));
874 }
875
876 #[test]
877 fn test_context_experiments_list() {
878 let exp1 = make_experiment("exp1", vec!["{}"], vec![1.0]);
879 let exp2 = make_experiment("exp2", vec!["{}"], vec![1.0]);
880 let data = make_context_data(vec![exp1, exp2]);
881 let context = Context::new(data);
882
883 let experiments = context.experiments();
884 assert!(experiments.contains(&"exp1".to_string()));
885 assert!(experiments.contains(&"exp2".to_string()));
886 }
887
888 #[test]
889 fn test_context_full_on_variant() {
890 let mut exp = make_experiment("test_exp", vec!["{}", r#"{"button":"red"}"#], vec![0.5, 0.5]);
891 exp.full_on_variant = 1;
892 let data = make_context_data(vec![exp]);
893 let mut context = Context::new(data);
894
895 context.set_unit("session_id", "any_user").unwrap();
896 assert_eq!(context.treatment("test_exp"), 1);
897 }
898
899 #[test]
900 fn test_context_audience_mismatch_strict() {
901 let mut exp = make_experiment("test_exp", vec!["{}", r#"{"button":"red"}"#], vec![0.0, 1.0]);
902 exp.audience = r#"{"filter":[{"eq":[{"var":"country"},{"value":"US"}]}]}"#.to_string();
903 exp.audience_strict = true;
904 let data = make_context_data(vec![exp]);
905 let mut context = Context::new(data);
906
907 context.set_unit("session_id", "test_user").unwrap();
908 context.set_attribute("country", json!("UK")).unwrap();
909
910 assert_eq!(context.treatment("test_exp"), 0);
911 }
912
913 #[test]
914 fn test_context_audience_match() {
915 let mut exp = make_experiment("test_exp", vec!["{}", r#"{"button":"red"}"#], vec![0.0, 1.0]);
916 exp.audience = r#"{"filter":[{"eq":[{"var":"country"},{"value":"US"}]}]}"#.to_string();
917 exp.audience_strict = true;
918 let data = make_context_data(vec![exp]);
919 let mut context = Context::new(data);
920
921 context.set_unit("session_id", "test_user").unwrap();
922 context.set_attribute("country", json!("US")).unwrap();
923
924 assert_eq!(context.treatment("test_exp"), 1);
925 }
926
927 #[test]
928 fn test_context_refresh() {
929 let exp1 = make_experiment("exp1", vec!["{}"], vec![1.0]);
930 let data1 = make_context_data(vec![exp1]);
931 let mut context = Context::new(data1);
932
933 let exp2 = make_experiment("exp2", vec!["{}"], vec![1.0]);
934 let data2 = make_context_data(vec![exp2]);
935 context.refresh(data2);
936
937 let experiments = context.experiments();
938 assert!(experiments.contains(&"exp2".to_string()));
939 assert!(!experiments.contains(&"exp1".to_string()));
940 }
941
942 #[test]
943 fn test_ergonomic_set_attribute_with_string() {
944 let data = make_context_data(vec![]);
945 let mut context = Context::new(data);
946
947 assert!(context.set_attribute("country", "US").is_ok());
948 assert_eq!(context.get_attribute("country"), Some(&json!("US")));
949 }
950
951 #[test]
952 fn test_ergonomic_set_attribute_with_number() {
953 let data = make_context_data(vec![]);
954 let mut context = Context::new(data);
955
956 assert!(context.set_attribute("age", 25).is_ok());
957 assert_eq!(context.get_attribute("age"), Some(&json!(25)));
958 }
959
960 #[test]
961 fn test_ergonomic_set_attribute_with_bool() {
962 let data = make_context_data(vec![]);
963 let mut context = Context::new(data);
964
965 assert!(context.set_attribute("premium", true).is_ok());
966 assert_eq!(context.get_attribute("premium"), Some(&json!(true)));
967 }
968
969 #[test]
970 fn test_ergonomic_variable_value_with_string_default() {
971 let data = make_context_data(vec![]);
972 let mut context = Context::new(data);
973
974 let value = context.variable_value("nonexistent", "default_value");
975 assert_eq!(value, json!("default_value"));
976 }
977
978 #[test]
979 fn test_ergonomic_variable_value_with_number_default() {
980 let data = make_context_data(vec![]);
981 let mut context = Context::new(data);
982
983 let value = context.variable_value("nonexistent", 42);
984 assert_eq!(value, json!(42));
985 }
986
987 #[test]
988 fn test_ergonomic_peek_variable_value_with_bool_default() {
989 let data = make_context_data(vec![]);
990 let mut context = Context::new(data);
991
992 let value = context.peek_variable_value("nonexistent", false);
993 assert_eq!(value, json!(false));
994 }
995
996 #[test]
997 fn test_ergonomic_track_with_json_properties() {
998 let data = make_context_data(vec![]);
999 let mut context = Context::new(data);
1000
1001 assert!(context.track("purchase", json!({
1002 "item_count": 1,
1003 "total_amount": 99.99
1004 })).is_ok());
1005 assert_eq!(context.pending(), 1);
1006 }
1007
1008 #[test]
1009 fn test_ergonomic_track_with_unit_no_properties() {
1010 let data = make_context_data(vec![]);
1011 let mut context = Context::new(data);
1012
1013 assert!(context.track("click", ()).is_ok());
1014 assert_eq!(context.pending(), 1);
1015 }
1016}