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)