def test_sample_rows_not_fitted(self):
"""Check that the sample_rows raise an exception when is not fitted."""
# Setup
# Run and asserts
sdv_mock = mock.Mock()
sdv_mock.sampler = None
table_name = 'DEMO'
num_rows = 5
with self.assertRaises(NotFittedError):
SDV.sample_rows(sdv_mock, table_name, num_rows)
def test_sample_rows_fitted(self):
"""Check that the sample_rows is called."""
# Setup
# Run
sdv_mock = mock.Mock()
sdv_mock.sampler = mock.Mock()
table_name = 'DEMO'
num_rows = 5
SDV.sample_rows(sdv_mock, table_name, num_rows)
# Asserts
sdv_mock.sampler.sample_rows.assert_called_once_with(
'DEMO', 5, sample_children=True, reset_primary_keys=False)