def test__sample_with_conditions_empty_transformed_conditions(): """Test that None is passed to ``_sample_batch`` if transformed conditions are empty. The ``Sample`` method is expected to: - Return sampled data and pass None to ``sample_batch`` as the ``transformed_conditions``. Input: - Number of rows to sample - Conditions Output: - Sampled data """ # Setup model = GaussianCopula() data = pd.DataFrame({ 'column1': list(range(100)), 'column2': list(range(100)), 'column3': list(range(100)) }) conditions = {'column1': 25} conditions_series = pd.Series([25, 25, 25, 25, 25], name='column1') model._sample_batch = Mock() sampled = pd.DataFrame({ 'column1': [28, 28], 'column2': [37, 37], 'column3': [93, 93], }) model._sample_batch.return_value = sampled model.fit(data) model._metadata = Mock() model._metadata.get_fields.return_value = ['column1', 'column2', 'column3'] model._metadata.transform.return_value = pd.DataFrame() model._metadata.make_ids_unique.side_effect = lambda x: x # Run output = model._sample_with_conditions(pd.DataFrame([conditions] * 5), 100, None) # Assert expected_output = pd.DataFrame({ 'column1': [28, 28], 'column2': [37, 37], 'column3': [93, 93], }) _, args, kwargs = model._metadata.transform.mock_calls[0] pd.testing.assert_series_equal(args[0]['column1'], conditions_series) model._metadata.transform.assert_called_once() model._sample_batch.assert_called_with(5, 100, None, conditions, None, 0.01, None, None) pd.testing.assert_frame_equal(output, expected_output)
def test__sample_with_conditions_transform_conditions_correctly(): """Test that transformed conditions are batched correctly. The ``Sample`` method is expected to: - Return sampled data and call ``_sample_batch`` for every unique transformed condition group. Input: - Number of rows to sample - Conditions Output: - Sampled data """ # Setup model = GaussianCopula() data = pd.DataFrame({ 'column1': list(range(100)), 'column2': list(range(100)), 'column3': list(range(100)) }) condition_values = [25, 25, 25, 30, 30] conditions_series = pd.Series([25, 25, 25, 30, 30], name='column1') model._sample_batch = Mock() expected_outputs = [ pd.DataFrame({ 'column1': [25, 25, 25], 'column2': [37, 37, 37], 'column3': [93, 93, 93], }), pd.DataFrame({ 'column1': [30], 'column2': [37], 'column3': [93], }), pd.DataFrame({ 'column1': [30], 'column2': [37], 'column3': [93], }) ] model._sample_batch.side_effect = expected_outputs model.fit(data) model._metadata = Mock() model._metadata.get_fields.return_value = ['column1', 'column2', 'column3'] model._metadata.transform.return_value = pd.DataFrame( [[50], [50], [50], [60], [70]], columns=['transformed_column']) # Run model._sample_with_conditions(pd.DataFrame({'column1': condition_values}), 100, None) # Assert _, args, kwargs = model._metadata.transform.mock_calls[0] pd.testing.assert_series_equal(args[0]['column1'], conditions_series) model._metadata.transform.assert_called_once() model._sample_batch.assert_any_call( 3, 100, None, {'column1': 25}, {'transformed_column': 50}, 0.01, None, None, ) model._sample_batch.assert_any_call( 1, 100, None, {'column1': 30}, {'transformed_column': 60}, 0.01, None, None, ) model._sample_batch.assert_any_call( 1, 100, None, {'column1': 30}, {'transformed_column': 70}, 0.01, None, None, )