moto/tests/test_sagemaker/test_sagemaker_cloudformation.py

import json
import boto3

import pytest
import sure  # noqa
from botocore.exceptions import ClientError

from moto import mock_cloudformation, mock_sagemaker
from moto.sts.models import ACCOUNT_ID


def _get_notebook_instance_template_string(
    resource_name="TestNotebook",
    instance_type="ml.c4.xlarge",
    role_arn="arn:aws:iam::{}:role/FakeRole".format(ACCOUNT_ID),
    include_outputs=True,
):
    template = {
        "AWSTemplateFormatVersion": "2010-09-09",
        "Resources": {
            resource_name: {
                "Type": "AWS::SageMaker::NotebookInstance",
                "Properties": {"InstanceType": instance_type, "RoleArn": role_arn},
            },
        },
    }
    if include_outputs:
        template["Outputs"] = {
            "Arn": {"Value": {"Ref": resource_name}},
            "Name": {"Value": {"Fn::GetAtt": [resource_name, "NotebookInstanceName"]}},
        }
    return json.dumps(template)


def _get_notebook_instance_lifecycle_config_template_string(
    resource_name="TestConfig", on_create=None, on_start=None, include_outputs=True,
):
    template = {
        "AWSTemplateFormatVersion": "2010-09-09",
        "Resources": {
            resource_name: {
                "Type": "AWS::SageMaker::NotebookInstanceLifecycleConfig",
                "Properties": {},
            },
        },
    }
    if on_create is not None:
        template["Resources"][resource_name]["Properties"]["OnCreate"] = [
            {"Content": on_create}
        ]
    if on_start is not None:
        template["Resources"][resource_name]["Properties"]["OnStart"] = [
            {"Content": on_start}
        ]
    if include_outputs:
        template["Outputs"] = {
            "Arn": {"Value": {"Ref": resource_name}},
            "Name": {
                "Value": {
                    "Fn::GetAtt": [
                        resource_name,
                        "NotebookInstanceLifecycleConfigName",
                    ]
                }
            },
        }
    return json.dumps(template)


@mock_cloudformation
@pytest.mark.parametrize(
    "stack_name,resource_name,template",
    [
        (
            "test_sagemaker_notebook_instance",
            "TestNotebookInstance",
            _get_notebook_instance_template_string(
                resource_name="TestNotebookInstance", include_outputs=False
            ),
        ),
        (
            "test_sagemaker_notebook_instance_lifecycle_config",
            "TestNotebookInstanceLifecycleConfig",
            _get_notebook_instance_lifecycle_config_template_string(
                resource_name="TestNotebookInstanceLifecycleConfig",
                include_outputs=False,
            ),
        ),
    ],
)
def test_sagemaker_cloudformation_create(stack_name, resource_name, template):
    cf = boto3.client("cloudformation", region_name="us-east-1")
    cf.create_stack(StackName=stack_name, TemplateBody=template)

    provisioned_resource = cf.list_stack_resources(StackName=stack_name)[
        "StackResourceSummaries"
    ][0]
    provisioned_resource["LogicalResourceId"].should.equal(resource_name)
    len(provisioned_resource["PhysicalResourceId"]).should.be.greater_than(0)


@mock_cloudformation
@mock_sagemaker
@pytest.mark.parametrize(
    "stack_name,template,describe_function_name,name_parameter,arn_parameter",
    [
        (
            "test_sagemaker_notebook_instance",
            _get_notebook_instance_template_string(),
            "describe_notebook_instance",
            "NotebookInstanceName",
            "NotebookInstanceArn",
        ),
        (
            "test_sagemaker_notebook_instance_lifecycle_config",
            _get_notebook_instance_lifecycle_config_template_string(),
            "describe_notebook_instance_lifecycle_config",
            "NotebookInstanceLifecycleConfigName",
            "NotebookInstanceLifecycleConfigArn",
        ),
    ],
)
def test_sagemaker_cloudformation_get_attr(
    stack_name, template, describe_function_name, name_parameter, arn_parameter
):
    cf = boto3.client("cloudformation", region_name="us-east-1")
    sm = boto3.client("sagemaker", region_name="us-east-1")

    # Create stack and get description for output values
    cf.create_stack(StackName=stack_name, TemplateBody=template)
    stack_description = cf.describe_stacks(StackName=stack_name)["Stacks"][0]
    outputs = {
        output["OutputKey"]: output["OutputValue"]
        for output in stack_description["Outputs"]
    }

    # Using the describe function, ensure output ARN matches resource ARN
    resource_description = getattr(sm, describe_function_name)(
        **{name_parameter: outputs["Name"]}
    )
    outputs["Arn"].should.equal(resource_description[arn_parameter])


@mock_cloudformation
@mock_sagemaker
@pytest.mark.parametrize(
    "stack_name,template,describe_function_name,name_parameter,arn_parameter,error_message",
    [
        (
            "test_sagemaker_notebook_instance",
            _get_notebook_instance_template_string(),
            "describe_notebook_instance",
            "NotebookInstanceName",
            "NotebookInstanceArn",
            "RecordNotFound",
        ),
        (
            "test_sagemaker_notebook_instance_lifecycle_config",
            _get_notebook_instance_lifecycle_config_template_string(),
            "describe_notebook_instance_lifecycle_config",
            "NotebookInstanceLifecycleConfigName",
            "NotebookInstanceLifecycleConfigArn",
            "Notebook Instance Lifecycle Config does not exist",
        ),
    ],
)
def test_sagemaker_cloudformation_notebook_instance_delete(
    stack_name,
    template,
    describe_function_name,
    name_parameter,
    arn_parameter,
    error_message,
):
    cf = boto3.client("cloudformation", region_name="us-east-1")
    sm = boto3.client("sagemaker", region_name="us-east-1")

    # Create stack and verify existence
    cf.create_stack(StackName=stack_name, TemplateBody=template)
    stack_description = cf.describe_stacks(StackName=stack_name)["Stacks"][0]
    outputs = {
        output["OutputKey"]: output["OutputValue"]
        for output in stack_description["Outputs"]
    }
    resource_description = getattr(sm, describe_function_name)(
        **{name_parameter: outputs["Name"]}
    )
    outputs["Arn"].should.equal(resource_description[arn_parameter])

    # Delete the stack and verify resource has also been deleted
    cf.delete_stack(StackName=stack_name)
    with pytest.raises(ClientError) as ce:
        getattr(sm, describe_function_name)(**{name_parameter: outputs["Name"]})
    ce.value.response["Error"]["Message"].should.contain(error_message)


@mock_cloudformation
@mock_sagemaker
def test_sagemaker_cloudformation_notebook_instance_update():
    cf = boto3.client("cloudformation", region_name="us-east-1")
    sm = boto3.client("sagemaker", region_name="us-east-1")

    # Set up template for stack with initial and update instance types
    stack_name = "test_sagemaker_notebook_instance"
    initial_instance_type = "ml.c4.xlarge"
    updated_instance_type = "ml.c4.4xlarge"
    initial_template_json = _get_notebook_instance_template_string(
        instance_type=initial_instance_type
    )
    updated_template_json = _get_notebook_instance_template_string(
        instance_type=updated_instance_type
    )

    # Create stack with initial template and check attributes
    cf.create_stack(StackName=stack_name, TemplateBody=initial_template_json)
    stack_description = cf.describe_stacks(StackName=stack_name)["Stacks"][0]
    outputs = {
        output["OutputKey"]: output["OutputValue"]
        for output in stack_description["Outputs"]
    }
    initial_notebook_name = outputs["Name"]
    notebook_instance_description = sm.describe_notebook_instance(
        NotebookInstanceName=initial_notebook_name,
    )
    initial_instance_type.should.equal(notebook_instance_description["InstanceType"])

    # Update stack with new instance type and check attributes
    cf.update_stack(StackName=stack_name, TemplateBody=updated_template_json)
    stack_description = cf.describe_stacks(StackName=stack_name)["Stacks"][0]
    outputs = {
        output["OutputKey"]: output["OutputValue"]
        for output in stack_description["Outputs"]
    }
    updated_notebook_name = outputs["Name"]
    updated_notebook_name.should.equal(initial_notebook_name)

    notebook_instance_description = sm.describe_notebook_instance(
        NotebookInstanceName=updated_notebook_name,
    )
    updated_instance_type.should.equal(notebook_instance_description["InstanceType"])


@mock_cloudformation
@mock_sagemaker
def test_sagemaker_cloudformation_notebook_instance_lifecycle_config_update():
    cf = boto3.client("cloudformation", region_name="us-east-1")
    sm = boto3.client("sagemaker", region_name="us-east-1")

    # Set up template for stack with initial and update instance types
    stack_name = "test_sagemaker_notebook_instance_lifecycle_config"
    initial_on_create_script = "echo Hello World"
    updated_on_create_script = "echo Goodbye World"
    initial_template_json = _get_notebook_instance_lifecycle_config_template_string(
        on_create=initial_on_create_script
    )
    updated_template_json = _get_notebook_instance_lifecycle_config_template_string(
        on_create=updated_on_create_script
    )

    # Create stack with initial template and check attributes
    cf.create_stack(StackName=stack_name, TemplateBody=initial_template_json)
    stack_description = cf.describe_stacks(StackName=stack_name)["Stacks"][0]
    outputs = {
        output["OutputKey"]: output["OutputValue"]
        for output in stack_description["Outputs"]
    }
    initial_config_name = outputs["Name"]
    notebook_lifecycle_config_description = sm.describe_notebook_instance_lifecycle_config(
        NotebookInstanceLifecycleConfigName=initial_config_name,
    )
    len(notebook_lifecycle_config_description["OnCreate"]).should.equal(1)
    initial_on_create_script.should.equal(
        notebook_lifecycle_config_description["OnCreate"][0]["Content"]
    )

    # Update stack with new instance type and check attributes
    cf.update_stack(StackName=stack_name, TemplateBody=updated_template_json)
    stack_description = cf.describe_stacks(StackName=stack_name)["Stacks"][0]
    outputs = {
        output["OutputKey"]: output["OutputValue"]
        for output in stack_description["Outputs"]
    }
    updated_config_name = outputs["Name"]
    updated_config_name.should.equal(initial_config_name)

    notebook_lifecycle_config_description = sm.describe_notebook_instance_lifecycle_config(
        NotebookInstanceLifecycleConfigName=updated_config_name,
    )
    len(notebook_lifecycle_config_description["OnCreate"]).should.equal(1)
    updated_on_create_script.should.equal(
        notebook_lifecycle_config_description["OnCreate"][0]["Content"]
    )