moto/tests/test_ecs/test_ecs_tasksets.py

from botocore.exceptions import ClientError
import boto3
import sure  # noqa # pylint: disable=unused-import

from moto import mock_ecs
import pytest


cluster_name = "test_ecs_cluster"
service_name = "test_ecs_service"
task_def_name = "test_ecs_task"


@mock_ecs
def test_create_task_set():
    client = boto3.client("ecs", region_name="us-east-1")
    _ = client.create_cluster(clusterName=cluster_name)
    create_task_def(client)
    _ = client.create_service(
        cluster=cluster_name,
        serviceName=service_name,
        taskDefinition=task_def_name,
        desiredCount=2,
        deploymentController={"type": "EXTERNAL"},
    )
    load_balancers = [
        {
            "targetGroupArn": "arn:aws:elasticloadbalancing:us-east-1:01234567890:targetgroup/c26b93c1bc35466ba792d5b08fe6a5bc/ec39113f8831453a",
            "containerName": "hello_world",
            "containerPort": 8080,
        },
    ]

    task_set = client.create_task_set(
        cluster=cluster_name,
        service=service_name,
        taskDefinition=task_def_name,
        loadBalancers=load_balancers,
    )["taskSet"]

    cluster_arn = client.describe_clusters(clusters=[cluster_name])["clusters"][0][
        "clusterArn"
    ]
    service_arn = client.describe_services(
        cluster=cluster_name, services=[service_name]
    )["services"][0]["serviceArn"]
    task_set["clusterArn"].should.equal(cluster_arn)
    task_set["serviceArn"].should.equal(service_arn)
    task_set["taskDefinition"].should.match(f"{task_def_name}:1$")
    task_set["scale"].should.equal({"value": 100.0, "unit": "PERCENT"})
    task_set["loadBalancers"][0]["targetGroupArn"].should.equal(
        "arn:aws:elasticloadbalancing:us-east-1:01234567890:targetgroup/"
        "c26b93c1bc35466ba792d5b08fe6a5bc/ec39113f8831453a"
    )
    task_set["loadBalancers"][0]["containerPort"].should.equal(8080)
    task_set["loadBalancers"][0]["containerName"].should.equal("hello_world")
    task_set["launchType"].should.equal("EC2")
    task_set["platformVersion"].should.equal("LATEST")


@mock_ecs
def test_create_task_set_errors():
    # given
    client = boto3.client("ecs", region_name="us-east-1")
    _ = client.create_cluster(clusterName=cluster_name)
    create_task_def(client)
    _ = client.create_service(
        cluster=cluster_name,
        serviceName=service_name,
        taskDefinition=task_def_name,
        desiredCount=2,
        deploymentController={"type": "EXTERNAL"},
    )

    # not existing launch type
    # when
    with pytest.raises(ClientError) as e:
        client.create_task_set(
            cluster=cluster_name,
            service=service_name,
            taskDefinition=task_def_name,
            launchType="SOMETHING",
        )

    # then
    ex = e.value
    ex.operation_name.should.equal("CreateTaskSet")
    ex.response["ResponseMetadata"]["HTTPStatusCode"].should.equal(400)
    ex.response["Error"]["Code"].should.contain("ClientException")
    ex.response["Error"]["Message"].should.equal(
        "launch type should be one of [EC2,FARGATE]"
    )


@mock_ecs
def test_describe_task_sets():
    client = boto3.client("ecs", region_name="us-east-1")
    _ = client.create_cluster(clusterName=cluster_name)
    create_task_def(client)
    _ = client.create_service(
        cluster=cluster_name,
        serviceName=service_name,
        taskDefinition=task_def_name,
        desiredCount=2,
        deploymentController={"type": "EXTERNAL"},
    )

    load_balancers = [
        {
            "targetGroupArn": "arn:aws:elasticloadbalancing:us-east-1:01234567890:targetgroup/c26b93c1bc35466ba792d5b08fe6a5bc/ec39113f8831453a",
            "containerName": "hello_world",
            "containerPort": 8080,
        }
    ]

    _ = client.create_task_set(
        cluster=cluster_name,
        service=service_name,
        taskDefinition=task_def_name,
        loadBalancers=load_balancers,
    )
    task_sets = client.describe_task_sets(cluster=cluster_name, service=service_name)[
        "taskSets"
    ]
    assert "tags" not in task_sets[0]

    task_sets = client.describe_task_sets(
        cluster=cluster_name, service=service_name, include=["TAGS"]
    )["taskSets"]

    cluster_arn = client.describe_clusters(clusters=[cluster_name])["clusters"][0][
        "clusterArn"
    ]

    service_arn = client.describe_services(
        cluster=cluster_name, services=[service_name]
    )["services"][0]["serviceArn"]

    task_sets.should.have.length_of(1)
    task_sets[0].should.have.key("tags")
    task_sets[0]["taskDefinition"].should.match(f"{task_def_name}:1$")
    task_sets[0]["clusterArn"].should.equal(cluster_arn)
    task_sets[0]["serviceArn"].should.equal(service_arn)
    task_sets[0]["serviceArn"].should.match(f"{service_name}$")
    task_sets[0]["scale"].should.equal({"value": 100.0, "unit": "PERCENT"})
    task_sets[0]["taskSetArn"].should.match(f"{task_sets[0]['id']}$")
    task_sets[0]["loadBalancers"][0]["targetGroupArn"].should.equal(
        "arn:aws:elasticloadbalancing:us-east-1:01234567890:targetgroup/"
        "c26b93c1bc35466ba792d5b08fe6a5bc/ec39113f8831453a"
    )
    task_sets[0]["loadBalancers"][0]["containerPort"].should.equal(8080)
    task_sets[0]["loadBalancers"][0]["containerName"].should.equal("hello_world")
    task_sets[0]["launchType"].should.equal("EC2")


@mock_ecs
def test_delete_task_set():
    client = boto3.client("ecs", region_name="us-east-1")
    _ = client.create_cluster(clusterName=cluster_name)
    create_task_def(client)
    _ = client.create_service(
        cluster=cluster_name,
        serviceName=service_name,
        taskDefinition=task_def_name,
        desiredCount=2,
        deploymentController={"type": "EXTERNAL"},
    )

    task_set = client.create_task_set(
        cluster=cluster_name, service=service_name, taskDefinition=task_def_name
    )["taskSet"]

    task_sets = client.describe_task_sets(
        cluster=cluster_name, service=service_name, taskSets=[task_set["taskSetArn"]]
    )["taskSets"]

    assert len(task_sets) == 1

    task_sets = client.describe_task_sets(
        cluster=cluster_name, service=service_name, taskSets=[task_set["taskSetArn"]]
    )["taskSets"]

    assert len(task_sets) == 1

    response = client.delete_task_set(
        cluster=cluster_name, service=service_name, taskSet=task_set["taskSetArn"]
    )
    assert response["taskSet"]["taskSetArn"] == task_set["taskSetArn"]

    task_sets = client.describe_task_sets(
        cluster=cluster_name, service=service_name, taskSets=[task_set["taskSetArn"]]
    )["taskSets"]

    assert len(task_sets) == 0

    with pytest.raises(ClientError):
        _ = client.delete_task_set(
            cluster=cluster_name, service=service_name, taskSet=task_set["taskSetArn"]
        )


@mock_ecs
def test_delete_task_set__using_partial_arn():
    client = boto3.client("ecs", region_name="us-east-1")
    _ = client.create_cluster(clusterName=cluster_name)
    create_task_def(client)
    _ = client.create_service(
        cluster=cluster_name,
        serviceName=service_name,
        taskDefinition=task_def_name,
        desiredCount=2,
        deploymentController={"type": "EXTERNAL"},
    )

    task_set = client.create_task_set(
        cluster=cluster_name, service=service_name, taskDefinition=task_def_name
    )["taskSet"]

    # Partial ARN match
    # arn:aws:ecs:us-east-1:123456789012:task-set/test_ecs_cluster/test_ecs_service/ecs-svc/386233676373827416
    # --> ecs-svc/386233676373827416
    partial_arn = "/".join(task_set["taskSetArn"].split("/")[-2:])
    task_sets = client.describe_task_sets(
        cluster=cluster_name, service=service_name, taskSets=[partial_arn]
    )["taskSets"]

    assert len(task_sets) == 1

    response = client.delete_task_set(
        cluster=cluster_name, service=service_name, taskSet=partial_arn
    )
    assert response["taskSet"]["taskSetArn"] == task_set["taskSetArn"]


@mock_ecs
def test_update_service_primary_task_set():
    client = boto3.client("ecs", region_name="us-east-1")
    _ = client.create_cluster(clusterName=cluster_name)
    create_task_def(client)
    _ = client.create_service(
        cluster=cluster_name,
        serviceName=service_name,
        desiredCount=2,
        deploymentController={"type": "EXTERNAL"},
    )

    task_set = client.create_task_set(
        cluster=cluster_name, service=service_name, taskDefinition=task_def_name
    )["taskSet"]

    service = client.describe_services(cluster=cluster_name, services=[service_name],)[
        "services"
    ][0]

    _ = client.update_service_primary_task_set(
        cluster=cluster_name,
        service=service_name,
        primaryTaskSet=task_set["taskSetArn"],
    )

    service = client.describe_services(cluster=cluster_name, services=[service_name],)[
        "services"
    ][0]
    assert service["taskSets"][0]["status"] == "PRIMARY"
    assert service["taskDefinition"] == service["taskSets"][0]["taskDefinition"]

    another_task_set = client.create_task_set(
        cluster=cluster_name, service=service_name, taskDefinition=task_def_name
    )["taskSet"]
    service = client.describe_services(cluster=cluster_name, services=[service_name],)[
        "services"
    ][0]
    assert service["taskSets"][1]["status"] == "ACTIVE"

    _ = client.update_service_primary_task_set(
        cluster=cluster_name,
        service=service_name,
        primaryTaskSet=another_task_set["taskSetArn"],
    )
    service = client.describe_services(cluster=cluster_name, services=[service_name],)[
        "services"
    ][0]
    assert service["taskSets"][0]["status"] == "ACTIVE"
    assert service["taskSets"][1]["status"] == "PRIMARY"
    assert service["taskDefinition"] == service["taskSets"][1]["taskDefinition"]


@mock_ecs
def test_update_task_set():
    client = boto3.client("ecs", region_name="us-east-1")
    _ = client.create_cluster(clusterName=cluster_name)
    create_task_def(client)
    _ = client.create_service(
        cluster=cluster_name,
        serviceName=service_name,
        desiredCount=2,
        deploymentController={"type": "EXTERNAL"},
    )

    task_set = client.create_task_set(
        cluster=cluster_name, service=service_name, taskDefinition=task_def_name
    )["taskSet"]

    another_task_set = client.create_task_set(
        cluster=cluster_name, service=service_name, taskDefinition=task_def_name
    )["taskSet"]
    assert another_task_set["scale"]["unit"] == "PERCENT"
    assert another_task_set["scale"]["value"] == 100.0

    client.update_task_set(
        cluster=cluster_name,
        service=service_name,
        taskSet=task_set["taskSetArn"],
        scale={"value": 25.0, "unit": "PERCENT"},
    )

    updated_task_set = client.describe_task_sets(
        cluster=cluster_name, service=service_name, taskSets=[task_set["taskSetArn"]]
    )["taskSets"][0]
    assert updated_task_set["scale"]["value"] == 25.0
    assert updated_task_set["scale"]["unit"] == "PERCENT"


@mock_ecs
def test_create_task_sets_with_tags():
    client = boto3.client("ecs", region_name="us-east-1")
    _ = client.create_cluster(clusterName=cluster_name)
    create_task_def(client)
    _ = client.create_service(
        cluster=cluster_name,
        serviceName=service_name,
        taskDefinition=task_def_name,
        desiredCount=2,
        deploymentController={"type": "EXTERNAL"},
    )

    load_balancers = [
        {
            "targetGroupArn": "arn:aws:elasticloadbalancing:us-east-1:01234567890:targetgroup/c26b93c1bc35466ba792d5b08fe6a5bc/ec39113f8831453a",
            "containerName": "hello_world",
            "containerPort": 8080,
        }
    ]

    _ = client.create_task_set(
        cluster=cluster_name,
        service=service_name,
        taskDefinition=task_def_name,
        loadBalancers=load_balancers,
        tags=[{"key": "k1", "value": "v1"}, {"key": "k2", "value": "v2"}],
    )

    task_set = client.describe_task_sets(
        cluster=cluster_name, service=service_name, include=["TAGS"]
    )["taskSets"][0]
    task_set.should.have.key("tags").equals(
        [{"key": "k1", "value": "v1"}, {"key": "k2", "value": "v2"}]
    )

    client.tag_resource(
        resourceArn=task_set["taskSetArn"], tags=[{"key": "k3", "value": "v3"}]
    )

    task_set = client.describe_task_sets(
        cluster=cluster_name, service=service_name, include=["TAGS"]
    )["taskSets"][0]
    task_set.should.have.key("tags")
    task_set["tags"].should.have.length_of(3)
    task_set["tags"].should.contain({"key": "k1", "value": "v1"})
    task_set["tags"].should.contain({"key": "k2", "value": "v2"})
    task_set["tags"].should.contain({"key": "k3", "value": "v3"})

    client.untag_resource(resourceArn=task_set["taskSetArn"], tagKeys=["k2"])

    resp = client.list_tags_for_resource(resourceArn=task_set["taskSetArn"])
    resp.should.have.key("tags")
    resp["tags"].should.have.length_of(2)
    resp["tags"].should.contain({"key": "k1", "value": "v1"})
    resp["tags"].should.contain({"key": "k3", "value": "v3"})


def create_task_def(client):
    client.register_task_definition(
        family=task_def_name,
        containerDefinitions=[
            {
                "name": "hello_world",
                "image": "docker/hello-world:latest",
                "cpu": 1024,
                "memory": 400,
                "essential": True,
                "environment": [
                    {"name": "AWS_ACCESS_KEY_ID", "value": "SOME_ACCESS_KEY"}
                ],
                "logConfiguration": {"logDriver": "json-file"},
            }
        ],
    )