Merge 35e24d5 into 9aef44b

kayak · Jul 5, 2018 · 03437ed · 03437ed
2 parents 9aef44b + 35e24d5
commit 03437ed
Show file tree

Hide file tree

Showing 2 changed files with 40 additions and 3 deletions.
diff --git a/fireant/slicer/queries/builder.py b/fireant/slicer/queries/builder.py
@@ -255,9 +255,9 @@ def fetch(self, limit=None, offset=None, hint=None, force_include=()) -> pd.Seri
                           str(query),
                           dimensions=self._dimensions)
 
-        display_key = getattr(dimension, 'display_key', None)
-        if display_key is not None:
-            return data[display_key]
+        df_key = format_key(getattr(dimension, 'display_key', None))
+        if df_key is not None:
+            return data[df_key]
 
         display_key = 'display'
         if hasattr(dimension, 'display_values'):

diff --git a/fireant/tests/slicer/queries/test_dimension_choices.py b/fireant/tests/slicer/queries/test_dimension_choices.py
@@ -1,5 +1,11 @@
 from unittest import TestCase
+from unittest.mock import (
+    ANY,
+    Mock,
+    patch,
+)
 
+from ..matchers import DimensionMatcher
 from ..mocks import slicer
 
 
@@ -61,3 +67,34 @@ def test_filter_choices(self):
                          'FROM "politics"."politician" '
                          'WHERE "political_party" IN (\'d\',\'r\') '
                          'GROUP BY "$candidate","$candidate_display"', str(query))
+
+
+# noinspection SqlDialectInspection,SqlNoDataSourceInspection
+@patch('fireant.slicer.queries.builder.fetch_data')
+class DimensionsChoicesFetchTests(TestCase):
+    def test_query_choices_for_cat_dimension(self, mock_fetch_data: Mock):
+        slicer.dimensions.political_party \
+            .choices \
+            .fetch()
+
+        mock_fetch_data.assert_called_once_with(ANY,
+                                                'SELECT '
+                                                '"political_party" "$political_party" '
+                                                'FROM "politics"."politician" '
+                                                'GROUP BY "$political_party" '
+                                                'ORDER BY "$political_party"',
+                                                dimensions=DimensionMatcher(slicer.dimensions.political_party))
+
+    def test_query_choices_for_uni_dimension(self, mock_fetch_data: Mock):
+        slicer.dimensions.candidate \
+            .choices \
+            .fetch()
+
+        mock_fetch_data.assert_called_once_with(ANY,
+                                                'SELECT '
+                                                '"candidate_id" "$candidate",'
+                                                '"candidate_name" "$candidate_display" '
+                                                'FROM "politics"."politician" '
+                                                'GROUP BY "$candidate","$candidate_display" '
+                                                'ORDER BY "$candidate_display"',
+                                                dimensions=DimensionMatcher(slicer.dimensions.candidate))