def test_get_package_versions_valid_model(get_distribution_mock):
"""Test the ``get_package_versions`` method with model.
Expect that the package versions for `sdv` and `rdt` are returned.
Setup:
- Patch the versions of sdv and rdt
- Create a model named GaussianCopula
Input:
- model
Output:
- a dict mapping libraries to library versions.
"""
# Setup
expected = {
'sdv': get_distribution_mock.return_value.version,
'rdt': get_distribution_mock.return_value.version,
'copulas': get_distribution_mock.return_value.version,
}
model = Mock(spec=GaussianCopula)
# Run
versions = get_package_versions(model=model)
# Assert
assert get_distribution_mock.has_calls([call('sdv'), call('rdt'), call('copulas')])
assert versions == expected
def test_get_package_versions_error(get_distribution_mock):
"""Test the ``get_package_versions`` method with model.
Expect that the package versions for `sdv` and `rdt` are returned, and that the unavailable
packages are ignored.
Setup:
- Patch the versions of sdv and rdt
- Create a model named GaussianCopula
Input:
- model
Output:
- a dict mapping libraries to library versions.
"""
# Setup
dist_mock = Mock()
get_distribution_mock.side_effect = [
dist_mock, pkg_resources.ResolutionError(), dist_mock]
expected = {
'sdv': dist_mock.version,
'copulas': dist_mock.version,
}
model = Mock(spec=GaussianCopula)
# Run
versions = get_package_versions(model=model)
# Assert
assert get_distribution_mock.has_calls([call('sdv'), call('rdt'), call('copulas')])
assert versions == expected
def save(self, path):
"""Save this model instance to the given path using pickle.
Args:
path (str):
Path where the SDV instance will be serialized.
"""
self._package_versions = get_package_versions(getattr(self, '_model', None))
with open(path, 'wb') as output:
pickle.dump(self, output)
def test_get_package_versions_no_model(get_distribution_mock):
"""Test the ``get_package_versions`` method when there is no model.
Expect that the package versions for `sdv` and `rdt` are returned.
Setup:
- Patch the versions of sdv and rdt
Input:
- model is None
Output:
- a dict mapping libraries to library versions.
"""
# Setup
expected = {
'sdv': get_distribution_mock.return_value.version,
'rdt': get_distribution_mock.return_value.version,
}
# Run
versions = get_package_versions(model=None)
# Assert
assert get_distribution_mock.has_calls([call('sdv'), call('rdt')])
assert versions == expected