from uuid import uuid4 import pytest from botocore.exceptions import ClientError from moto import mock_aws, settings from . import _get_clients, _setup @mock_aws def test_create_managed_compute_environment(): ec2_client, iam_client, ecs_client, _, batch_client = _get_clients() _, subnet_id, sg_id, iam_arn = _setup(ec2_client, iam_client) compute_name = str(uuid4()) resp = batch_client.create_compute_environment( computeEnvironmentName=compute_name, type="MANAGED", state="ENABLED", computeResources={ "type": "EC2", "minvCpus": 5, "maxvCpus": 10, "desiredvCpus": 5, "instanceTypes": ["t2.small", "t2.medium"], "imageId": "some_image_id", "subnets": [subnet_id], "securityGroupIds": [sg_id], "ec2KeyPair": "string", "instanceRole": iam_arn.replace("role", "instance-profile"), "tags": {"string": "string"}, "bidPercentage": 123, "spotIamFleetRole": "string", }, serviceRole=iam_arn, ) assert "computeEnvironmentArn" in resp assert resp["computeEnvironmentName"] == compute_name our_env = batch_client.describe_compute_environments( computeEnvironments=[compute_name] )["computeEnvironments"][0] # Given a t2.medium is 2 vcpu and t2.small is 1, therefore 2 mediums and 1 small should be created if not settings.TEST_SERVER_MODE: # Can't verify this in ServerMode, as other tests may have created instances resp = ec2_client.describe_instances() assert "Reservations" in resp assert len(resp["Reservations"]) == 3 # Should have created 1 ECS cluster all_clusters = ecs_client.list_clusters()["clusterArns"] assert our_env["ecsClusterArn"] in all_clusters @mock_aws def test_create_managed_compute_environment_with_instance_family(): """ The InstanceType parameter can have multiple values: instance_type t2.small instance_family t2 <-- What we're testing here 'optimal' unknown value """ ec2_client, iam_client, _, _, batch_client = _get_clients() _, subnet_id, sg_id, iam_arn = _setup(ec2_client, iam_client) compute_name = str(uuid4()) batch_client.create_compute_environment( computeEnvironmentName=compute_name, type="MANAGED", state="ENABLED", computeResources={ "type": "EC2", "minvCpus": 5, "maxvCpus": 10, "desiredvCpus": 5, "instanceTypes": ["t2"], "imageId": "some_image_id", "subnets": [subnet_id], "securityGroupIds": [sg_id], "ec2KeyPair": "string", "instanceRole": iam_arn.replace("role", "instance-profile"), "tags": {"string": "string"}, "bidPercentage": 123, "spotIamFleetRole": "string", }, serviceRole=iam_arn, ) our_env = batch_client.describe_compute_environments( computeEnvironments=[compute_name] )["computeEnvironments"][0] assert our_env["computeResources"]["instanceTypes"] == ["t2"] @mock_aws def test_create_managed_compute_environment_with_unknown_instance_type(): """ The InstanceType parameter can have multiple values: instance_type t2.small instance_family t2 'optimal' unknown value <-- What we're testing here """ ec2_client, iam_client, _, _, batch_client = _get_clients() _, subnet_id, sg_id, iam_arn = _setup(ec2_client, iam_client) compute_name = str(uuid4()) with pytest.raises(ClientError) as exc: batch_client.create_compute_environment( computeEnvironmentName=compute_name, type="MANAGED", state="ENABLED", computeResources={ "type": "EC2", "minvCpus": 5, "maxvCpus": 10, "desiredvCpus": 5, "instanceTypes": ["unknown"], "imageId": "some_image_id", "subnets": [subnet_id], "securityGroupIds": [sg_id], "ec2KeyPair": "string", "instanceRole": iam_arn.replace("role", "instance-profile"), "tags": {"string": "string"}, "bidPercentage": 123, "spotIamFleetRole": "string", }, serviceRole=iam_arn, ) err = exc.value.response["Error"] assert err["Code"] == "InvalidParameterValue" assert err["Message"] == "Instance type unknown does not exist" @mock_aws def test_create_unmanaged_compute_environment(): ec2_client, iam_client, ecs_client, _, batch_client = _get_clients() _, _, _, iam_arn = _setup(ec2_client, iam_client) compute_name = str(uuid4()) resp = batch_client.create_compute_environment( computeEnvironmentName=compute_name, type="UNMANAGED", state="ENABLED", serviceRole=iam_arn, ) assert "computeEnvironmentArn" in resp assert resp["computeEnvironmentName"] == compute_name our_env = batch_client.describe_compute_environments( computeEnvironments=[compute_name] )["computeEnvironments"][0] assert "ecsClusterArn" in our_env # Its unmanaged so no instances should be created if not settings.TEST_SERVER_MODE: # Can't verify this in ServerMode, as other tests may have created instances resp = ec2_client.describe_instances() assert "Reservations" in resp assert len(resp["Reservations"]) == 0 # Should have created 1 ECS cluster all_clusters = ecs_client.list_clusters()["clusterArns"] assert our_env["ecsClusterArn"] in all_clusters # TODO create 1000s of tests to test complex option combinations of create environment @mock_aws def test_describe_compute_environment(): ec2_client, iam_client, _, _, batch_client = _get_clients() _, _, _, iam_arn = _setup(ec2_client, iam_client) compute_name = str(uuid4()) compute_arn = ( f"arn:aws:batch:eu-central-1:123456789012:compute-environment/{compute_name}" ) batch_client.create_compute_environment( computeEnvironmentName=compute_name, type="UNMANAGED", state="ENABLED", serviceRole=iam_arn, ) all_envs = batch_client.describe_compute_environments()["computeEnvironments"] our_envs = [e for e in all_envs if e["computeEnvironmentName"] == compute_name] assert len(our_envs) == 1 assert our_envs[0]["computeEnvironmentName"] == compute_name assert our_envs[0]["computeEnvironmentArn"] == compute_arn assert "ecsClusterArn" in our_envs[0] assert our_envs[0]["state"] == "ENABLED" assert our_envs[0]["status"] == "VALID" # Test filtering resp = batch_client.describe_compute_environments(computeEnvironments=["test1"]) assert len(resp["computeEnvironments"]) == 0 @mock_aws def test_delete_unmanaged_compute_environment(): ec2_client, iam_client, ecs_client, _, batch_client = _get_clients() _, _, _, iam_arn = _setup(ec2_client, iam_client) compute_name = str(uuid4()) batch_client.create_compute_environment( computeEnvironmentName=compute_name, type="UNMANAGED", state="ENABLED", serviceRole=iam_arn, ) our_env = batch_client.describe_compute_environments( computeEnvironments=[compute_name] )["computeEnvironments"][0] batch_client.delete_compute_environment(computeEnvironment=compute_name) all_envs = batch_client.describe_compute_environments()["computeEnvironments"] all_names = [e["computeEnvironmentName"] for e in all_envs] assert compute_name not in all_names cluster = ecs_client.describe_clusters(clusters=[our_env["ecsClusterArn"]])[ "clusters" ][0] assert cluster["status"] == "INACTIVE" @mock_aws def test_delete_managed_compute_environment(): ec2_client, iam_client, ecs_client, _, batch_client = _get_clients() _, subnet_id, sg_id, iam_arn = _setup(ec2_client, iam_client) compute_name = str(uuid4()) batch_client.create_compute_environment( computeEnvironmentName=compute_name, type="MANAGED", state="ENABLED", computeResources={ "type": "EC2", "minvCpus": 5, "maxvCpus": 10, "desiredvCpus": 5, "instanceTypes": ["t2.small", "t2.medium"], "imageId": "some_image_id", "subnets": [subnet_id], "securityGroupIds": [sg_id], "ec2KeyPair": "string", "instanceRole": iam_arn.replace("role", "instance-profile"), "tags": {"string": "string"}, "bidPercentage": 123, "spotIamFleetRole": "string", }, serviceRole=iam_arn, ) our_env = batch_client.describe_compute_environments( computeEnvironments=[compute_name] )["computeEnvironments"][0] batch_client.delete_compute_environment(computeEnvironment=compute_name) all_envs = batch_client.describe_compute_environments()["computeEnvironments"] all_names = [e["computeEnvironmentName"] for e in all_envs] assert compute_name not in all_names if not settings.TEST_SERVER_MODE: # Too many instances to know which one is ours in ServerMode resp = ec2_client.describe_instances() assert "Reservations" in resp assert len(resp["Reservations"]) == 3 for reservation in resp["Reservations"]: assert reservation["Instances"][0]["State"]["Name"] == "terminated" cluster = ecs_client.describe_clusters(clusters=[our_env["ecsClusterArn"]])[ "clusters" ][0] assert cluster["status"] == "INACTIVE" @mock_aws def test_update_unmanaged_compute_environment_state(): ec2_client, iam_client, _, _, batch_client = _get_clients() _, _, _, iam_arn = _setup(ec2_client, iam_client) compute_name = str(uuid4()) batch_client.create_compute_environment( computeEnvironmentName=compute_name, type="UNMANAGED", state="ENABLED", serviceRole=iam_arn, ) batch_client.update_compute_environment( computeEnvironment=compute_name, state="DISABLED" ) all_envs = batch_client.describe_compute_environments()["computeEnvironments"] our_envs = [e for e in all_envs if e["computeEnvironmentName"] == compute_name] assert len(our_envs) == 1 assert our_envs[0]["state"] == "DISABLED" @mock_aws def test_update_iam_role(): ec2_client, iam_client, _, _, batch_client = _get_clients() _, _, _, iam_arn = _setup(ec2_client, iam_client) iam_arn2 = iam_client.create_role(RoleName="r", AssumeRolePolicyDocument="sp")[ "Role" ]["Arn"] compute_name = str(uuid4()) batch_client.create_compute_environment( computeEnvironmentName=compute_name, type="UNMANAGED", state="ENABLED", serviceRole=iam_arn, ) batch_client.update_compute_environment( computeEnvironment=compute_name, serviceRole=iam_arn2 ) all_envs = batch_client.describe_compute_environments()["computeEnvironments"] our_envs = [e for e in all_envs if e["computeEnvironmentName"] == compute_name] assert len(our_envs) == 1 assert our_envs[0]["serviceRole"] == iam_arn2 with pytest.raises(ClientError) as exc: batch_client.update_compute_environment( computeEnvironment=compute_name, serviceRole="unknown" ) err = exc.value.response["Error"] assert err["Code"] == "InvalidParameterValue" @pytest.mark.parametrize("compute_env_type", ["FARGATE", "FARGATE_SPOT"]) @mock_aws def test_create_fargate_managed_compute_environment(compute_env_type): ec2_client, iam_client, ecs_client, _, batch_client = _get_clients() _, subnet_id, sg_id, iam_arn = _setup(ec2_client, iam_client) compute_name = str(uuid4()) resp = batch_client.create_compute_environment( computeEnvironmentName=compute_name, type="MANAGED", state="ENABLED", computeResources={ "type": compute_env_type, "maxvCpus": 10, "subnets": [subnet_id], "securityGroupIds": [sg_id], }, serviceRole=iam_arn, ) assert "computeEnvironmentArn" in resp assert resp["computeEnvironmentName"] == compute_name our_env = batch_client.describe_compute_environments( computeEnvironments=[compute_name] )["computeEnvironments"][0] assert our_env["computeResources"]["type"] == compute_env_type # Should have created 1 ECS cluster all_clusters = ecs_client.list_clusters()["clusterArns"] assert our_env["ecsClusterArn"] in all_clusters @mock_aws def test_create_ec2_managed_compute_environment__without_required_params(): ec2_client, iam_client, _, _, batch_client = _get_clients() _, subnet_id, _, iam_arn = _setup(ec2_client, iam_client) with pytest.raises(ClientError) as exc: batch_client.create_compute_environment( computeEnvironmentName="ec2-env", type="MANAGED", state="ENABLED", computeResources={"type": "EC2", "maxvCpus": 1, "subnets": [subnet_id]}, serviceRole=iam_arn, ) err = exc.value.response["Error"] assert err["Code"] == "ClientException" assert ( "Error executing request, Exception : Instance role is required." in err["Message"] ) with pytest.raises(ClientError) as exc: batch_client.create_compute_environment( computeEnvironmentName="ec2-env", type="MANAGED", state="ENABLED", computeResources={ "type": "EC2", "maxvCpus": 1, "subnets": [subnet_id], "instanceRole": iam_arn.replace("role", "instance-profile"), }, serviceRole=iam_arn, ) err = exc.value.response["Error"] assert err["Code"] == "ClientException" assert ( "Error executing request, Exception : Resource minvCpus is required." in err["Message"] )