def test_model_package_enable_network_isolation_with_no_product_id(sagemaker_session): sagemaker_session.sagemaker_client.describe_model_package = Mock( return_value=DESCRIBE_MODEL_PACKAGE_RESPONSE) model_package = ModelPackage(role='role', model_package_arn='my-model-package', sagemaker_session=sagemaker_session) assert model_package.enable_network_isolation() is False
def test_model_package_enable_network_isolation_with_product_id(sagemaker_session): model_package_response = copy.deepcopy(DESCRIBE_MODEL_PACKAGE_RESPONSE) model_package_response['InferenceSpecification']['Containers'].append( { 'Image': '1.dkr.ecr.us-east-2.amazonaws.com/some-container:latest', 'ModelDataUrl': 's3://bucket/output/model.tar.gz', 'ProductId': 'some-product-id' } ) sagemaker_session.sagemaker_client.describe_model_package = Mock( return_value=model_package_response) model_package = ModelPackage(role='role', model_package_arn='my-model-package', sagemaker_session=sagemaker_session) assert model_package.enable_network_isolation() is True
def test_model_package_enable_network_isolation_with_product_id(sagemaker_session): model_package_response = copy.deepcopy(DESCRIBE_MODEL_PACKAGE_RESPONSE) model_package_response["InferenceSpecification"]["Containers"].append( { "Image": "1.dkr.ecr.us-east-2.amazonaws.com/some-container:latest", "ModelDataUrl": "s3://bucket/output/model.tar.gz", "ProductId": "some-product-id", } ) sagemaker_session.sagemaker_client.describe_model_package = Mock( return_value=model_package_response ) model_package = ModelPackage( role="role", model_package_arn="my-model-package", sagemaker_session=sagemaker_session ) assert model_package.enable_network_isolation() is True
def test_model_package_enable_network_isolation_with_no_product_id( sagemaker_session): model_package = ModelPackage(role="role", model_package_arn="my-model-package", sagemaker_session=sagemaker_session) assert model_package.enable_network_isolation() is False