Implement total cube count functionality

- Add unit tests for 'count' API method - Add integration tests for 'count' method (using fixtures) - Implement count functionality
Crunch-io · Dec 19, 2017 · 315c67c · 315c67c
1 parent 360278f
commit 315c67c
Show file tree

Hide file tree

Showing 3 changed files with 92 additions and 0 deletions.
diff --git a/src/cr/cube/crunch_cube.py b/src/cr/cube/crunch_cube.py
@@ -727,3 +727,11 @@ def has_means(self):
     def filter_annotation(self):
         '''Get cube's filter annotation.'''
         return self._cube.get('filter_names', [])
+
+    def count(self, weighted=True):
+        '''Get cube's count with automatic weighted/unweighted selection.'''
+        if weighted and self.is_weighted:
+            return sum(
+                self._cube['result']['measures'].get('count', {}).get('data')
+            )
+        return self._cube['result']['n']
diff --git a/tests/integration/test_crunch_cube.py b/tests/integration/test_crunch_cube.py
@@ -1673,3 +1673,15 @@ def test_ca_margin_with_hs(self):
         expected = np.array([6, 6, 6])
         actual = cube.margin(include_transforms_for_dims=[0, 1], axis=1)
         np.testing.assert_almost_equal(actual, expected)
+
+    def test_count_unweighted(self):
+        cube = CrunchCube(FIXT_ADMIT_X_GENDER_WEIGHTED)
+        expected = 4526
+        actual = cube.count(weighted=False)
+        self.assertEqual(actual, expected)
+
+    def test_count_weighted(self):
+        cube = CrunchCube(FIXT_ADMIT_X_GENDER_WEIGHTED)
+        expected = 4451.955438803242
+        actual = cube.count(weighted=True)
+        self.assertEqual(actual, expected)
diff --git a/tests/unit/test_crunch_cube.py b/tests/unit/test_crunch_cube.py
@@ -275,3 +275,75 @@ def test_y_offset_ca_as_zero_index_cube(self, mock_dims, mock_as_array):
         # 'expand' argument expands the CA dim as slices
         actual = CrunchCube({}).y_offset(expand=True)
         self.assertEqual(actual, expected)
+
+    @patch('cr.cube.crunch_cube.CrunchCube.is_weighted', False)
+    def test_n_unweighted_and_has_no_weight(self):
+        unweighted_count = Mock()
+        weighted_counts = Mock()
+        actual = CrunchCube({
+            'result': {
+                'n': unweighted_count,
+                'measures': {
+                    'count': {
+                        'data': weighted_counts,
+                    },
+                },
+            },
+        }).count(weighted=False)
+
+        expected = unweighted_count
+        self.assertEqual(actual, expected)
+
+    @patch('cr.cube.crunch_cube.CrunchCube.is_weighted', True)
+    def test_n_unweighted_and_has_weight(self):
+        unweighted_count = Mock()
+        weighted_counts = Mock()
+        actual = CrunchCube({
+            'result': {
+                'n': unweighted_count,
+                'measures': {
+                    'count': {
+                        'data': weighted_counts,
+                    },
+                },
+            },
+        }).count(weighted=False)
+
+        expected = unweighted_count
+        self.assertEqual(actual, expected)
+
+    @patch('cr.cube.crunch_cube.CrunchCube.is_weighted', False)
+    def test_n_weighted_and_has_no_weight(self):
+        unweighted_count = Mock()
+        weighted_counts = Mock()
+        actual = CrunchCube({
+            'result': {
+                'n': unweighted_count,
+                'measures': {
+                    'count': {
+                        'data': weighted_counts,
+                    },
+                },
+            },
+        }).count(weighted=True)
+
+        expected = unweighted_count
+        self.assertEqual(actual, expected)
+
+    @patch('cr.cube.crunch_cube.CrunchCube.is_weighted', True)
+    def test_n_weighted_and_has_weight(self):
+        unweighted_count = Mock()
+        weighted_counts = [1, 2, 3, 4]
+        actual = CrunchCube({
+            'result': {
+                'n': unweighted_count,
+                'measures': {
+                    'count': {
+                        'data': weighted_counts,
+                    },
+                },
+            },
+        }).count(weighted=True)
+
+        expected = sum(weighted_counts)
+        self.assertEqual(actual, expected)