def test_save(self, pickle_mock, open_mock):
sdv = SDV()
sdv.save('save/path.pkl')
open_mock.assert_called_once_with('save/path.pkl', 'wb')
output = open_mock.return_value.__enter__.return_value
pickle_mock.dump.assert_called_once_with(sdv, output)
def test_sample_all_not_fitted(self):
"""Check that the sample_all raise an exception when is not fitted."""
# Run & asserts
sdv = Mock()
sdv.sampler = None
with pytest.raises(NotFittedError):
SDV.sample_all(sdv)
def test_sample_not_fitted(self):
"""Check that the sample raise an exception when is not fitted."""
# Setup
sdv = SDV()
# Run
with pytest.raises(NotFittedError):
sdv.sample('DEMO', 5)
def test_sample_not_fitted(self):
"""Check that the sample raise an exception when is not fitted."""
# Setup
sdv = Mock()
sdv.sampler = None
# Run
with pytest.raises(NotFittedError):
SDV.sample(sdv, 'DEMO', 5)
def test_sample_not_fitted(self):
"""Check that the sample raise an exception when is not fitted."""
# Run and asserts
sdv = Mock()
sdv.sampler = None
table_name = 'DEMO'
num_rows = 5
with pytest.raises(NotFittedError):
SDV.sample(sdv, table_name, num_rows)
def test__validate_dataset_structure_raise_error(self):
"""Test that a ValueError is raised because the bad structure"""
# Setup
sdv = Mock()
sdv.metadata.get_tables.return_value = ['foo', 'bar', 'tar']
sdv.metadata.get_parents.side_effect = [[], [], ['foo', 'bar']]
# Run
with pytest.raises(ValueError):
SDV._validate_dataset_structure(sdv)
def test__validate_dataset_structure_no_error(self):
"""Test that any error is raised with a supported structure"""
# Setup
sdv = Mock()
sdv.metadata.get_tables.return_value = ['foo', 'bar', 'tar']
sdv.metadata.get_parents.side_effect = [[], ['foo'], ['bar']]
# Run
SDV._validate_dataset_structure(sdv)
# Asserts
assert sdv.metadata.get_parents.call_count == 3
def test_load(self, pickle_mock, open_mock):
returned = SDV.load('save/path.pkl')
open_mock.assert_called_once_with('save/path.pkl', 'rb')
output = open_mock.return_value.__enter__.return_value
pickle_mock.load.assert_called_once_with(output)
assert returned is pickle_mock.load.return_value
def run_demo(folder_name):
"""Runs the demo for specified folder"""
start = timer()
meta_file = os.path.join('demo', folder_name,
folder_name.capitalize() + '_manual_meta.json')
sdv = SDV(meta_file)
sdv.fit()
sampled_rows = {}
LOGGER.info('Parent map: %s', sdv.dn.parent_map)
LOGGER.info('Transformed data: %s', sdv.dn.transformed_data)
table_list = table_dict[folder_name]
for table in table_list:
sampled_rows[table] = sdv.sample_rows(table, 1)
LOGGER.info('Sampled row from %s: %s', table, sampled_rows[table])
end = timer()
LOGGER.info('Total time: %s seconds', round(end - start))
def test____init__users_params(self):
"""Create default instance"""
# Run
sdv = SDV(model='test', model_kwargs={'a': 2})
# Asserts
assert sdv.model == 'test'
assert sdv.model_kwargs == {'a': 2}
def run_demo(folder_name):
"""Runs the demo for specified folder"""
start = timer()
meta_file = os.path.join('demo', folder_name,
folder_name.capitalize() + '_manual_meta.json')
sdv = SDV(meta_file)
sdv.fit()
sampled = sdv.sample_all()
LOGGER.info('Parent map: %s', sdv.dn.parent_map)
LOGGER.info('Transformed data: %s', sdv.dn.transformed_data)
for name, table in sampled.items():
LOGGER.info('Sampled row from %s: %s', name, table.head(3).T)
end = timer()
LOGGER.info('Total time: %s seconds', round(end - start))
def test____init__default(self):
"""Create default instance"""
# Run
sdv = SDV()
# Asserts
assert sdv.model == DEFAULT_MODEL
assert sdv.model_kwargs == DEFAULT_MODEL_KWARGS
assert sdv.model_kwargs is not DEFAULT_MODEL_KWARGS
def test_sample_all_fitted(self):
"""Check that the sample_all is called"""
# Run
sdv = Mock()
sdv.sampler.sample_all.return_value = 'test'
result = SDV.sample_all(sdv)
# Asserts
sdv.sampler.sample_all.assert_called_once_with(
5, reset_primary_keys=False)
assert result == 'test'
def test_sample_fitted(self):
"""Check that the sample is called."""
# Sample
sdv = Mock()
sdv._model_instance.sample.return_value = 'test'
# Run
result = SDV.sample(sdv, 'DEMO', 5)
# Asserts
assert result == 'test'
sdv._model_instance.sample.assert_called_once_with(
'DEMO', 5, sample_children=True, reset_primary_keys=False)
def test_sample_fitted(self):
"""Check that the sample is called."""
# Run
sdv = Mock()
table_name = 'DEMO'
num_rows = 5
sdv.sampler.sample.return_value = 'test'
result = SDV.sample(sdv, table_name, num_rows)
# Asserts
sdv.sampler.sample.assert_called_once_with('DEMO',
5,
sample_children=True,
reset_primary_keys=False)
assert result == 'test'