Merge 3bfd7bd into e53e57e

fireant/queries/pagination.py

@@ -124,7 +124,14 @@ def _apply_pagination(df):
         # there might be missing combinations of index values.
         dfx = df.reset_index(level=0, drop=True)
         value_in_index = sorted_dimension_values.isin(dfx.index)
-        index_slice = sorted_dimension_values[value_in_index]
+        index_slice = sorted_dimension_values[value_in_index].values
+
+        """
+        In the case of bool dimensions, convert index_slice to an array of literal `True`, because pandas `.loc` handles
+        lists of bool as a mask.
+        """
+        if bool in {type(x) for x in sorted_dimension_values}:
+            index_slice |= True
 
         # Need to include nulls so append them to the end of the sorted data frame
         isnull = _index_isnull(dfx)

fireant/tests/dataset/mocks.py

@@ -355,6 +355,10 @@ def PoliticsRow(timestamp, candidate_id, candidate_name, political_party, electi
     .groupby([f('timestamp'), f('political_party')]) \
     .sum()
 
+dimx2_date_bool_df = mock_politics_database[[f('timestamp'), f('winner'), f('votes')]] \
+    .groupby([f('timestamp'), f('winner')]) \
+    .sum()
+
 dimx2_date_num_df = mock_politics_database[[f('timestamp'), f('candidate-id'), f('votes'), f('wins')]] \
     .groupby([f('timestamp'), f('candidate-id')]) \
     .sum()

fireant/tests/queries/test_fetch_operations.py

@@ -0,0 +1,84 @@
+from unittest import TestCase
+
+from unittest.mock import (
+    Mock,
+    call,
+    patch,
+)
+
+import fireant as f
+from fireant.tests.dataset.mocks import (
+    ElectionOverElection,
+    mock_dataset,
+)
+from fireant.utils import alias_selector
+
+
+@patch('fireant.queries.builder.scrub_totals_from_share_results', side_effect=lambda *args, **kwargs: args[0])
+@patch('fireant.queries.builder.paginate', side_effect=lambda *args, **kwargs: args[0])
+@patch('fireant.queries.builder.fetch_data')
+class QueryBuilderOperationsTests(TestCase):
+    def test_operations_evaluated(self, mock_fetch_data: Mock, *mocks):
+        mock_operation = Mock(name='mock_operation ', spec=f.Operation)
+        mock_operation.alias, mock_operation.definition = 'mock_operation', mock_dataset.table.abc
+        mock_operation.metrics = []
+
+        mock_widget = f.Widget(mock_operation)
+        mock_widget.transform = Mock()
+
+        mock_df = {}
+        mock_fetch_data.return_value = mock_df
+
+        # Need to keep widget the last call in the chain otherwise the object gets cloned and the assertion won't work
+        mock_dataset.query \
+            .dimension(mock_dataset.fields.timestamp) \
+            .widget(mock_widget) \
+            .fetch()
+
+        mock_operation.apply.assert_called_once_with(mock_df, None)
+
+    def test_operations_evaluated_for_each_reference(self, mock_fetch_data: Mock, *mocks):
+        eoe = ElectionOverElection(mock_dataset.fields.timestamp)
+
+        mock_operation = Mock(name='mock_operation ', spec=f.Operation)
+        mock_operation.alias, mock_operation.definition = 'mock_operation', mock_dataset.table.abc
+        mock_operation.metrics = []
+
+        mock_widget = f.Widget(mock_operation)
+        mock_widget.transform = Mock()
+
+        mock_df = {}
+        mock_fetch_data.return_value = mock_df
+
+        # Need to keep widget the last call in the chain otherwise the object gets cloned and the assertion won't work
+        mock_dataset.query \
+            .dimension(mock_dataset.fields.timestamp) \
+            .reference(eoe) \
+            .widget(mock_widget) \
+            .fetch()
+
+        mock_operation.apply.assert_has_calls([
+            call(mock_df, None),
+            call(mock_df, eoe),
+        ])
+
+    def test_operations_results_stored_in_data_frame(self, mock_fetch_data: Mock, *mocks):
+        mock_operation = Mock(name='mock_operation ', spec=f.Operation)
+        mock_operation.alias, mock_operation.definition = 'mock_operation', mock_dataset.table.abc
+        mock_operation.metrics = []
+
+        mock_widget = f.Widget(mock_operation)
+        mock_widget.transform = Mock()
+
+        mock_df = {}
+        mock_fetch_data.return_value = mock_df
+
+        # Need to keep widget the last call in the chain otherwise the object gets cloned and the assertion won't work
+        mock_dataset.query \
+            .dimension(mock_dataset.fields.timestamp) \
+            .widget(mock_widget) \
+            .fetch()
+
+        f_op_key = alias_selector(mock_operation.alias)
+        self.assertIn(f_op_key, mock_df)
+        self.assertEqual(mock_df[f_op_key], mock_operation.apply.return_value)