def test__sample_with_previous(self):
"""Check _sample with previous"""
# Setup
get_extension_mock = Mock()
get_extension_mock.return_value = {'child_rows': 0.999}
get_model_mock = Mock()
get_model_mock.return_value = None
sample_valid_rows_mock = Mock()
sample_valid_rows_mock.return_value = pd.DataFrame({'foo': [0, 1]})
sample_children_mock = Mock()
dn_mock = Mock()
dn_mock.foreign_keys = {
('DEMO', 'p_name'): ('parent_id', 'foreign_key')
}
# Run
sampler_mock = Mock()
sampler_mock._get_extension = get_extension_mock
sampler_mock._get_model = get_model_mock
sampler_mock._sample_valid_rows = sample_valid_rows_mock
sampler_mock._sample_children = sample_children_mock
sampler_mock.dn = dn_mock
table_name = 'DEMO'
parent_name = 'p_name'
parent_row = {'parent_id': 'foo'}
sampled = {'DEMO': pd.DataFrame({'bar': [1, 2]})}
Sampler._sample(sampler_mock, table_name, parent_name, parent_row,
sampled)
# Asserts
exp_dataframe_sampled = pd.DataFrame({
'bar': [1, 2, np.NaN, np.NaN],
'foo': [np.NaN, np.NaN, 0, 1],
'foreign_key': [np.NaN, np.NaN, 'foo', 'foo']
})
args_sample_children, kwargs_sample_children = sample_children_mock.call_args
exp_arg_table_name, exp_arg_sampled = args_sample_children
get_extension_mock.assert_called_once_with({'parent_id': 'foo'},
'DEMO', 'p_name')
get_model_mock.assert_called_once_with({'child_rows': 0.999})
sample_valid_rows_mock.assert_called_once_with(None, 1, 'DEMO')
assert exp_arg_table_name == 'DEMO'
pd.testing.assert_frame_equal(exp_arg_sampled['DEMO'],
exp_dataframe_sampled)
def test__sample_without_previous(self):
"""Check _sample without previous"""
# Setup
get_extension_mock = Mock()
get_extension_mock.return_value = {'child_rows': 0.999}
get_model_mock = Mock()
get_model_mock.return_value = None
sample_valid_rows_mock = Mock()
sample_valid_rows_mock.return_value = {}
sample_children_mock = Mock()
dn_mock = Mock()
dn_mock.foreign_keys = {
('DEMO', 'p_name'): ('parent_id', 'foreign_key')
}
# Run
sampler_mock = Mock()
sampler_mock._get_extension = get_extension_mock
sampler_mock._get_model = get_model_mock
sampler_mock._sample_valid_rows = sample_valid_rows_mock
sampler_mock._sample_children = sample_children_mock
sampler_mock.dn = dn_mock
table_name = 'DEMO'
parent_name = 'p_name'
parent_row = {'parent_id': 'foo'}
sampled = {}
Sampler._sample(sampler_mock, table_name, parent_name, parent_row,
sampled)
# Asserts
get_extension_mock.assert_called_once_with({'parent_id': 'foo'},
'DEMO', 'p_name')
get_model_mock.assert_called_once_with({'child_rows': 0.999})
sample_valid_rows_mock.assert_called_once_with(None, 1, 'DEMO')
sample_children_mock.assert_called_once_with(
'DEMO', {'DEMO': {
'foreign_key': 'foo'
}})