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(
    instance_type="ml.c4.xlarge",
    role_arn="arn:aws:iam::{}:role/FakeRole".format(ACCOUNT_ID),
    include_outputs=True,
):
    template = {
        "AWSTemplateFormatVersion": "2010-09-09",
        "Resources": {
            "TestNotebookInstance": {
                "Type": "AWS::SageMaker::NotebookInstance",
                "Properties": {"InstanceType": instance_type, "RoleArn": role_arn},
            },
        },
    }
    if include_outputs:
        template["Outputs"] = {
            "NotebookInstanceArn": {"Value": {"Ref": "TestNotebookInstance"}},
            "NotebookInstanceName": {
                "Value": {
                    "Fn::GetAtt": ["TestNotebookInstance", "NotebookInstanceName"]
                },
            },
        }
    return json.dumps(template)


@mock_cloudformation
def test_sagemaker_cloudformation_create_notebook_instance():
    cf = boto3.client("cloudformation", region_name="us-east-1")

    stack_name = "test_sagemaker_notebook_instance"
    template = _get_notebook_instance_template_string(include_outputs=False)
    cf.create_stack(StackName=stack_name, TemplateBody=template)

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


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

    stack_name = "test_sagemaker_notebook_instance"
    template = _get_notebook_instance_template_string()
    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"]
    }
    notebook_instance_name = outputs["NotebookInstanceName"]
    notebook_instance_arn = outputs["NotebookInstanceArn"]

    notebook_instance_description = sm.describe_notebook_instance(
        NotebookInstanceName=notebook_instance_name,
    )
    notebook_instance_arn.should.equal(
        notebook_instance_description["NotebookInstanceArn"]
    )


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

    # Create stack with notebook instance and verify existence
    stack_name = "test_sagemaker_notebook_instance"
    template = _get_notebook_instance_template_string()
    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"]
    }
    notebook_instance = sm.describe_notebook_instance(
        NotebookInstanceName=outputs["NotebookInstanceName"],
    )
    outputs["NotebookInstanceArn"].should.equal(
        notebook_instance["NotebookInstanceArn"]
    )

    # Delete the stack and verify notebook instance has also been deleted
    # TODO replace exception check with `list_notebook_instances` method when implemented
    cf.delete_stack(StackName=stack_name)
    with pytest.raises(ClientError) as ce:
        sm.describe_notebook_instance(
            NotebookInstanceName=outputs["NotebookInstanceName"]
        )
    ce.value.response["Error"]["Message"].should.contain("RecordNotFound")


@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["NotebookInstanceName"]
    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["NotebookInstanceName"]
    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"])