From 51778f9735d239ce0f04d55b32ce6ddb956516e1 Mon Sep 17 00:00:00 2001
From: Ernesto Arbitrio <ernesto@crunch.io>
Date: Fri, 12 Jul 2024 16:03:46 +0200
Subject: [PATCH] add median measure in enum

---
 src/cr/cube/enums.py    |  3 +++
 tests/unit/test_enum.py | 38 ++++++++++++++++++++++++++++++++++++--
 2 files changed, 39 insertions(+), 2 deletions(-)

diff --git a/src/cr/cube/enums.py b/src/cr/cube/enums.py
index 5d97e908c..a9fd02dcd 100644
--- a/src/cr/cube/enums.py
+++ b/src/cr/cube/enums.py
@@ -121,6 +121,7 @@ class MEASURE(enum.Enum):
     COLUMN_STDDEV = "col_std_dev"
     COLUMN_STDERR = "col_std_err"
     MEAN = "mean"
+    MEDIAN = "median"
     PAIRWISE_T_TEST = "pairwise_t_test"
     POPULATION = "population"
     POPULATION_MOE = "population_moe"
@@ -170,6 +171,7 @@ class CUBE_MEASURE(enum.Enum):
 NUMERIC_CUBE_MEASURES = frozenset(
     (
         CUBE_MEASURE.MEAN,
+        CUBE_MEASURE.MEDIAN,
         CUBE_MEASURE.SUM,
         CUBE_MEASURE.STDDEV,
         CUBE_MEASURE.UNWEIGHTED_VALID_COUNT,
@@ -180,6 +182,7 @@ class CUBE_MEASURE(enum.Enum):
 NUMERIC_MEASURES = frozenset(
     (
         MEASURE.MEAN,
+        MEASURE.MEDIAN,
         MEASURE.SUM,
         MEASURE.STDDEV,
         MEASURE.WEIGHTED_VALID_COUNT,
diff --git a/tests/unit/test_enum.py b/tests/unit/test_enum.py
index 9f35169b9..b21aa4511 100644
--- a/tests/unit/test_enum.py
+++ b/tests/unit/test_enum.py
@@ -1,8 +1,14 @@
 # encoding: utf-8
 
-"""Unit test suite for cr.cube.cube_slice module."""
+"""Unit test suite for cr.cube.enums module."""
 
-from cr.cube.enums import _DimensionType
+from cr.cube.enums import (
+    CUBE_MEASURE,
+    MEASURE,
+    NUMERIC_CUBE_MEASURES,
+    NUMERIC_MEASURES,
+    _DimensionType,
+)
 
 
 class Describe_DimensionType:
@@ -20,3 +26,31 @@ def it_knows_its_name(self):
         dimension_type = _DimensionType("WORM_HOLE")
         name = dimension_type.name
         assert name == "WORM_HOLE"
+
+
+class TestCubeMeasures:
+    def test_numeric_cube_measures_intersection(self):
+        intersection = NUMERIC_CUBE_MEASURES & {m for m in CUBE_MEASURE}
+        expected_intersection = {
+            MEASURE.MEAN,
+            MEASURE.MEDIAN,
+            MEASURE.SUM,
+            MEASURE.STDDEV,
+            MEASURE.UNWEIGHTED_VALID_COUNT,
+            MEASURE.WEIGHTED_VALID_COUNT,
+        }
+        assert sorted(list([m.value for m in intersection])) == sorted(
+            list([m.value for m in expected_intersection])
+        )
+
+    def test_numeric_cube_measures_difference(self):
+        difference = {m.value for m in NUMERIC_MEASURES} - {
+            m.value for m in NUMERIC_CUBE_MEASURES
+        }
+        expected_difference = {
+            "smoothed_mean",
+            "total_share_sum",
+            "row_share_sum",
+            "col_share_sum",
+        }
+        assert list(sorted(difference)) == list(sorted(expected_difference))