def test__sample_rows(self):
"""Test sample rows from model"""
# Setup
sampler = Mock(spec=Sampler)
sampler._get_primary_keys.return_value = ('pk', [1, 2, 3, 4])
model = Mock()
model.sample.return_value = dict()
# Run
result = Sampler._sample_rows(sampler,
model,
num_rows=5,
table_name='test')
# Asserts
assert result == {'pk': [1, 2, 3, 4]}
sampler._get_primary_keys.assert_called_once_with('test', 5)
model.sample.called_once_with(5)
def test__sample_rows(self):
"""Test sample rows from model"""
# Setup
primary_keys = ('pk', [1, 2, 3, 4])
model_sample = dict()
# Run
sampler = Mock()
sampler._get_primary_keys.return_value = primary_keys
model = Mock()
model.sample.return_value = model_sample
num_rows = 5
table_name = 'test'
result = Sampler._sample_rows(sampler, model, num_rows, table_name)
# Asserts
expected = {'pk': [1, 2, 3, 4]}
assert result == expected
sampler._get_primary_keys.assert_called_once_with('test', 5)
model.sample.called_once_with(5)
