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____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 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____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 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 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))