Ejemplo n.º 1
0
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
Ejemplo n.º 2
0
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
Ejemplo n.º 3
0
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