oxihuman_core/
experiment_tracker.rs

1// Copyright (C) 2026 COOLJAPAN OU (Team KitaSan)
2// SPDX-License-Identifier: Apache-2.0
3#![allow(dead_code)]
4
5//! Experiment variant tracker.
6
7use std::collections::HashMap;
8
9/// A recorded experiment assignment.
10#[derive(Debug, Clone)]
11pub struct ExperimentAssignment {
12    pub experiment_id: String,
13    pub variant: String,
14    pub user_id: String,
15}
16
17/// Tracks which experiment variant each user is assigned to.
18#[derive(Debug, Default)]
19pub struct ExperimentTracker {
20    /// experiment_id -> (user_id -> variant)
21    assignments: HashMap<String, HashMap<String, String>>,
22}
23
24impl ExperimentTracker {
25    pub fn new() -> Self {
26        Self::default()
27    }
28
29    pub fn assign(&mut self, experiment_id: &str, user_id: &str, variant: &str) {
30        self.assignments
31            .entry(experiment_id.to_string())
32            .or_default()
33            .insert(user_id.to_string(), variant.to_string());
34    }
35
36    pub fn get_variant(&self, experiment_id: &str, user_id: &str) -> Option<&str> {
37        self.assignments
38            .get(experiment_id)?
39            .get(user_id)
40            .map(String::as_str)
41    }
42
43    pub fn participant_count(&self, experiment_id: &str) -> usize {
44        self.assignments
45            .get(experiment_id)
46            .map(|m| m.len())
47            .unwrap_or(0)
48    }
49
50    pub fn variant_counts(&self, experiment_id: &str) -> HashMap<String, usize> {
51        let mut counts: HashMap<String, usize> = HashMap::new();
52        if let Some(users) = self.assignments.get(experiment_id) {
53            for v in users.values() {
54                *counts.entry(v.clone()).or_insert(0) += 1;
55            }
56        }
57        counts
58    }
59
60    pub fn experiment_count(&self) -> usize {
61        self.assignments.len()
62    }
63
64    pub fn all_assignments(&self, experiment_id: &str) -> Vec<ExperimentAssignment> {
65        let mut out = Vec::new();
66        if let Some(users) = self.assignments.get(experiment_id) {
67            for (user, variant) in users {
68                out.push(ExperimentAssignment {
69                    experiment_id: experiment_id.to_string(),
70                    variant: variant.clone(),
71                    user_id: user.clone(),
72                });
73            }
74        }
75        out
76    }
77}
78
79pub fn new_experiment_tracker() -> ExperimentTracker {
80    ExperimentTracker::new()
81}
82
83pub fn tracker_assign(tracker: &mut ExperimentTracker, exp: &str, user: &str, variant: &str) {
84    tracker.assign(exp, user, variant);
85}
86
87pub fn tracker_get_variant<'a>(
88    tracker: &'a ExperimentTracker,
89    exp: &str,
90    user: &str,
91) -> Option<&'a str> {
92    tracker.get_variant(exp, user)
93}
94
95pub fn tracker_participant_count(tracker: &ExperimentTracker, exp: &str) -> usize {
96    tracker.participant_count(exp)
97}
98
99pub fn tracker_experiment_count(tracker: &ExperimentTracker) -> usize {
100    tracker.experiment_count()
101}
102
103#[cfg(test)]
104mod tests {
105    use super::*;
106
107    #[test]
108    fn test_assign_and_get() {
109        let mut t = new_experiment_tracker();
110        tracker_assign(&mut t, "exp1", "user1", "control");
111        assert_eq!(tracker_get_variant(&t, "exp1", "user1"), Some("control"));
112    }
113
114    #[test]
115    fn test_unknown_user() {
116        let t = new_experiment_tracker();
117        assert_eq!(tracker_get_variant(&t, "exp", "nobody"), None);
118    }
119
120    #[test]
121    fn test_participant_count() {
122        let mut t = new_experiment_tracker();
123        tracker_assign(&mut t, "exp", "u1", "a");
124        tracker_assign(&mut t, "exp", "u2", "b");
125        assert_eq!(tracker_participant_count(&t, "exp"), 2);
126    }
127
128    #[test]
129    fn test_experiment_count() {
130        let mut t = new_experiment_tracker();
131        tracker_assign(&mut t, "exp1", "u1", "a");
132        tracker_assign(&mut t, "exp2", "u2", "b");
133        assert_eq!(tracker_experiment_count(&t), 2);
134    }
135
136    #[test]
137    fn test_variant_counts() {
138        let mut t = new_experiment_tracker();
139        tracker_assign(&mut t, "e", "u1", "ctrl");
140        tracker_assign(&mut t, "e", "u2", "ctrl");
141        tracker_assign(&mut t, "e", "u3", "treat");
142        let counts = t.variant_counts("e");
143        assert_eq!(counts["ctrl"], 2);
144        assert_eq!(counts["treat"], 1);
145    }
146
147    #[test]
148    fn test_overwrite_assignment() {
149        let mut t = new_experiment_tracker();
150        tracker_assign(&mut t, "e", "u1", "a");
151        tracker_assign(&mut t, "e", "u1", "b");
152        assert_eq!(tracker_get_variant(&t, "e", "u1"), Some("b"));
153    }
154
155    #[test]
156    fn test_all_assignments_count() {
157        let mut t = new_experiment_tracker();
158        tracker_assign(&mut t, "e", "u1", "a");
159        tracker_assign(&mut t, "e", "u2", "b");
160        assert_eq!(t.all_assignments("e").len(), 2);
161    }
162
163    #[test]
164    fn test_zero_participants_unknown_exp() {
165        let t = new_experiment_tracker();
166        assert_eq!(tracker_participant_count(&t, "no_exp"), 0);
167    }
168
169    #[test]
170    fn test_multiple_experiments_isolated() {
171        /* user in exp1 is not visible in exp2 */
172        let mut t = new_experiment_tracker();
173        tracker_assign(&mut t, "exp1", "u1", "x");
174        assert_eq!(tracker_get_variant(&t, "exp2", "u1"), None);
175    }
176}
oxihuman_core/experiment_tracker.rs

oxihuman_core/
experiment_tracker.rs
