def test_create_sagemaker_model_generates_model_name(name_from_base, sagemaker_session): model_package_name = "my-model-package" model_package = ModelPackage(role="role", model_package_arn=model_package_name, sagemaker_session=sagemaker_session) model_package._create_sagemaker_model() name_from_base.assert_called_with(model_package_name) assert name_from_base.return_value == model_package.name
def test_create_sagemaker_model_generates_model_name_each_time( name_from_base, sagemaker_session): model_package_name = "my-model-package" model_package = ModelPackage(role="role", model_package_arn=model_package_name, sagemaker_session=sagemaker_session) model_package._create_sagemaker_model() model_package._create_sagemaker_model() name_from_base.assert_called_with(model_package_name) assert 2 == name_from_base.call_count
def test_create_sagemaker_model_uses_model_name(name_from_base, sagemaker_session): model_name = "my-model" model_package_name = "my-model-package" model_package = ModelPackage( role="role", name=model_name, model_package_arn=model_package_name, sagemaker_session=sagemaker_session, ) model_package._create_sagemaker_model() assert model_name == model_package.name name_from_base.assert_not_called() sagemaker_session.create_model.assert_called_with( model_name, "role", {"ModelPackageName": model_package_name}, vpc_config=None, enable_network_isolation=False, )