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