def test__sample_table_sampled_not_empty(self):
"""Test sample table when sampled previous sampled rows exist."""
# Setup
model = Mock(spec=SDVModel)
model.return_value = model
sampler = Mock(spec=Sampler)
sampler.model = model
sampler.model_kwargs = dict()
sampler._get_extension.return_value = {'child_rows': 5}
table_model_mock = Mock()
sampler.models = {'test': table_model_mock}
# model_mock = Mock()
# sampler._get_model.return_value = model_mock
sampler._sample_rows.return_value = pd.DataFrame(
{'value': [6, 7, 8, 9, 10]})
sampler.metadata.get_primary_key.return_value = 'id'
sampler.metadata.get_foreign_key.return_value = 'parent_id'
# Run
parent_row = pd.Series({'id': 1})
sampled = {
'test':
pd.DataFrame({
'value': [1, 2, 3, 4, 5],
'parent_id': [0, 0, 0, 0, 0]
})
}
Sampler._sample_table(sampler, 'test', 'parent', parent_row, sampled)
# Asserts
sampler._get_extension.assert_called_once_with(parent_row, 'test')
# sampler._get_model.assert_called_once_with({'child_rows': 5}, table_model_mock)
sampler._sample_rows.assert_called_once_with(model, 5, 'test')
assert sampler._sample_children.call_count == 1
assert sampler._sample_children.call_args[0][0] == 'test'
expected_sampled = pd.DataFrame({
'value': [1, 2, 3, 4, 5, 6, 7, 8, 9, 10],
'parent_id': [0, 0, 0, 0, 0, 1, 1, 1, 1, 1]
})
pd.testing.assert_frame_equal(
sampler._sample_children.call_args[0][1]['test'], expected_sampled)
def test__sample_table_sampled_empty(self):
"""Test sample table when sampled is still an empty dict."""
# Setup
model = Mock(spec=SDVModel)
model.return_value = model
sampler = Mock(spec=Sampler)
sampler.model = model
sampler.model_kwargs = dict()
sampler._get_extension.return_value = {'child_rows': 5}
table_model_mock = Mock()
sampler.models = {'test': table_model_mock}
sampler._sample_rows.return_value = pd.DataFrame(
{'value': [1, 2, 3, 4, 5]})
sampler.metadata.get_primary_key.return_value = 'id'
sampler.metadata.get_foreign_key.return_value = 'parent_id'
# Run
parent_row = pd.Series({'id': 0})
sampled = dict()
Sampler._sample_table(sampler, 'test', 'parent', parent_row, sampled)
# Asserts
sampler._get_extension.assert_called_once_with(parent_row, 'test')
sampler._sample_rows.assert_called_once_with(model, 5, 'test')
assert sampler._sample_children.call_count == 1
assert sampler._sample_children.call_args[0][0] == 'test'
expected_sampled = pd.DataFrame(
{
'value': [1, 2, 3, 4, 5],
'parent_id': [0, 0, 0, 0, 0]
},
columns=['value', 'parent_id'])
pd.testing.assert_frame_equal(
sampler._sample_children.call_args[0][1]['test'], expected_sampled)
