feat: ability to override the analysis unit for an experiment (#1274)

Author: mikewilli

lib/metric-config-parser/metric_config_parser/experiment.py

@@ -166,6 +166,9 @@ def analysis_unit(self) -> AnalysisUnit | None:
         """Retrieve the appropriate analysis unit, which is
         derived from the experiment's randomization unit.
         """
+        if self.experiment_spec.analysis_unit is not None:
+            return AnalysisUnit(self.experiment_spec.analysis_unit)
+
         if self.randomization_unit and self.randomization_unit == RandomizationUnit.GROUP_ID:
             return AnalysisUnit.PROFILE_GROUP
 
@@ -268,6 +271,11 @@ def _validate_dataset_id(instance: Any, attribute, value):
         raise ValueError("dataset_id must be set to a custom dataset for private experiments")
 
 
+def _validate_analysis_unit(instance: Any, attribute, value):
+    if value is not None:
+        AnalysisUnit(value)
+
+
 @attr.s(auto_attribs=True, kw_only=True)
 class ExperimentSpec:
     """Describes the interface for overriding experiment details."""
@@ -283,6 +291,7 @@ class ExperimentSpec:
     is_private: bool = False
     dataset_id: str | None = attr.ib(default=None, validator=_validate_dataset_id)
     sample_size: int | None = None
+    analysis_unit: str | None = attr.ib(default=None, validator=_validate_analysis_unit)
 
     def resolve(
         self,

lib/metric-config-parser/metric_config_parser/tests/conftest.py

@@ -7,7 +7,7 @@
 from git import Repo
 
 from metric_config_parser.config import ConfigCollection
-from metric_config_parser.experiment import Branch, Experiment
+from metric_config_parser.experiment import Branch, BucketConfig, Experiment
 
 TEST_DIR = Path(__file__).parent
 
@@ -27,6 +27,9 @@ def experiments():
             reference_branch="b",
             is_high_population=False,
             app_name="firefox_desktop",
+            bucket_config=BucketConfig(
+                randomization_unit="group_id", namespace="namespace", start=0, count=10
+            ),
         ),
         Experiment(
             experimenter_slug="test_slug",

lib/metric-config-parser/metric_config_parser/tests/test_experiment.py

@@ -73,6 +73,31 @@ def test_control_branch(self, experiments, config_collection):
         configured = spec.resolve(experiments[0], config_collection)
         assert configured.experiment.reference_branch == "a"
 
+    def test_analysis_unit(self, experiments, config_collection):
+        trivial = AnalysisSpec().resolve(experiments[0], config_collection)
+        assert trivial.experiment.analysis_unit == AnalysisUnit.PROFILE_GROUP
+
+        conf = dedent(
+            """
+            [experiment]
+            analysis_unit = "client_id"
+            """
+        )
+        spec = AnalysisSpec.from_dict(toml.loads(conf))
+        configured = spec.resolve(experiments[0], config_collection)
+        assert configured.experiment.analysis_unit == AnalysisUnit.CLIENT
+
+    def test_analysis_unit_invalid(self):
+        conf = dedent(
+            """
+            [experiment]
+            analysis_unit = "not_valid_id"
+            """
+        )
+
+        with pytest.raises(ClassValidationError):
+            AnalysisSpec.from_dict(toml.loads(conf))
+
     def test_recognizes_segments(self, experiments, config_collection):
         conf = dedent(
             """

lib/metric-config-parser/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "mozilla-metric-config-parser"
-version = "2025.11.1"
+version = "2025.12.1"
 authors = [{ name = "Mozilla Corporation", email = "[email protected]" }]
 description = "Parses metric configuration files."
 readme = "README.md"