def test__get_model(self): """Test get model""" # Setup sampler = Mock(spec=Sampler) sampler._unflatten_dict.return_value = {'unflatten': 'dict'} sampler._unflatten_gaussian_copula.return_value = { 'unflatten': 'gaussian' } table_model = Mock() table_model.to_dict.return_value = { 'distribution': 'copulas.multivariate.gaussian.GaussianMultivariate' } # Run extension = {'extension': 'dict'} Sampler._get_model(sampler, extension, table_model) # Asserts sampler._unflatten_dict.assert_called_once_with({'extension': 'dict'}) expected_unflatten_gaussian_call = { 'unflatten': 'dict', 'fitted': True, 'distribution': 'copulas.multivariate.gaussian.GaussianMultivariate' } sampler._unflatten_gaussian_copula.assert_called_once_with( expected_unflatten_gaussian_call) table_model.from_dict.assert_called_once_with( {'unflatten': 'gaussian'})
def test__get_model(self, qualified_name): """Retrieve the model with parameters""" # Setup unflatten_dict_mock = Mock() unflatten_dict_mock.return_value = dict() qualified_name.return_value = 'copulas.multivariate.gaussian.GaussianMultivariate' unflatten_gaussian_mock = Mock() unflatten_gaussian_mock.return_value = None model_mock = Mock() model_mock.from_dict.return_value = None modeler_mock = Mock() modeler_mock.model = model_mock # Run sampler_mock = Mock() sampler_mock._unflatten_dict = unflatten_dict_mock sampler_mock.modeler = modeler_mock sampler_mock._unflatten_gaussian_copula = unflatten_gaussian_mock Sampler._get_model(sampler_mock, None) # Asserts exp_unflatten_gaussian_called = { 'fitted': True, 'type': 'copulas.multivariate.gaussian.GaussianMultivariate' } qualified_name.assert_called_once_with(modeler_mock.model) unflatten_dict_mock.assert_called_once_with(None) unflatten_gaussian_mock.assert_called_once_with( exp_unflatten_gaussian_called) model_mock.from_dict.assert_called_once_with(None)