diff --git a/moto/batch/exceptions.py b/moto/batch/exceptions.py index cd6031a95..a71e54ce3 100644 --- a/moto/batch/exceptions.py +++ b/moto/batch/exceptions.py @@ -30,3 +30,8 @@ class ValidationError(AWSError): class InternalFailure(AWSError): CODE = 'InternalFailure' STATUS = 500 + + +class ClientException(AWSError): + CODE = 'ClientException' + STATUS = 400 diff --git a/moto/batch/models.py b/moto/batch/models.py index 7ed75e749..8572a46c0 100644 --- a/moto/batch/models.py +++ b/moto/batch/models.py @@ -9,7 +9,7 @@ from moto.iam import iam_backends from moto.ec2 import ec2_backends from moto.ecs import ecs_backends -from .exceptions import InvalidParameterValueException, InternalFailure +from .exceptions import InvalidParameterValueException, InternalFailure, ClientException from .utils import make_arn_for_compute_env from moto.ec2.exceptions import InvalidSubnetIdError from moto.ec2.models import INSTANCE_TYPES as EC2_INSTANCE_TYPES @@ -31,12 +31,14 @@ class ComputeEnvironment(BaseModel): self.instances = [] self.ecs_arn = None + self.ecs_name = None def add_instance(self, instance): self.instances.append(instance) - def set_ecs_arn(self, arn): + def set_ecs(self, arn, name): self.ecs_arn = arn + self.ecs_name = name class BatchBackend(BaseBackend): @@ -75,7 +77,7 @@ class BatchBackend(BaseBackend): self.__dict__ = {} self.__init__(region_name) - def get_compute_environment(self, arn): + def get_compute_environment_arn(self, arn): return self._compute_environments.get(arn) def get_compute_environment_by_name(self, name): @@ -84,6 +86,20 @@ class BatchBackend(BaseBackend): return comp_env return None + def get_compute_environment(self, identifier): + """ + Get compute environment by name or ARN + :param identifier: Name or ARN + :type identifier: str + + :return: Compute Environment or None + :rtype: ComputeEnvironment or None + """ + env = self.get_compute_environment_arn(identifier) + if env is None: + env = self.get_compute_environment_by_name(identifier) + return env + def describe_compute_environments(self, environments=None, max_results=None, next_token=None): envs = set() if environments is not None: @@ -173,7 +189,7 @@ class BatchBackend(BaseBackend): # Should be of format P2OnDemand_Batch_UUID cluster_name = 'OnDemand_Batch_' + str(uuid.uuid4()) ecs_cluster = self.ecs_backend.create_cluster(cluster_name) - new_comp_env.set_ecs_arn(ecs_cluster.arn) + new_comp_env.set_ecs(ecs_cluster.arn, cluster_name) return compute_environment_name, new_comp_env.arn @@ -271,6 +287,52 @@ class BatchBackend(BaseBackend): return instances + def delete_compute_environment(self, compute_environment_name): + if compute_environment_name is None: + raise InvalidParameterValueException('Missing computeEnvironment parameter') + + compute_env = self.get_compute_environment(compute_environment_name) + + if compute_env is not None: + # Pop ComputeEnvironment + self._compute_environments.pop(compute_env.arn) + + # Delete ECS cluster + self.ecs_backend.delete_cluster(compute_env.ecs_name) + + if compute_env.type == 'MANAGED': + # Delete compute envrionment + instance_ids = [instance.id for instance in compute_env.instances] + self.ec2_backend.terminate_instances(instance_ids) + + def update_compute_environment(self, compute_environment_name, state, compute_resources, service_role): + # Validate + compute_env = self.get_compute_environment(compute_environment_name) + if compute_env is None: + raise ClientException('Compute environment {0} does not exist') + + # Look for IAM role + if service_role is not None: + try: + role = self.iam_backend.get_role_by_arn(service_role) + except IAMNotFoundException: + raise InvalidParameterValueException('Could not find IAM role {0}'.format(service_role)) + + compute_env.service_role = role + + if state is not None: + if state not in ('ENABLED', 'DISABLED'): + raise InvalidParameterValueException('state {0} must be one of ENABLED | DISABLED'.format(state)) + + compute_env.state = state + + if compute_resources is not None: + # TODO Implement resizing of instances based on changing vCpus + # compute_resources CAN contain desiredvCpus, maxvCpus, minvCpus, and can contain none of them. + pass + + return compute_env.name, compute_env.arn + available_regions = boto3.session.Session().get_available_regions("batch") batch_backends = {region: BatchBackend(region_name=region) for region in available_regions} diff --git a/moto/batch/responses.py b/moto/batch/responses.py index 590cc27a4..86ee4fdfe 100644 --- a/moto/batch/responses.py +++ b/moto/batch/responses.py @@ -76,3 +76,37 @@ class BatchResponse(BaseResponse): result = {'computeEnvironments': envs} return json.dumps(result) + + # DeleteComputeEnvironment + def deletecomputeenvironment(self): + compute_environment = self._get_param('computeEnvironment') + + try: + self.batch_backend.delete_compute_environment(compute_environment) + except AWSError as err: + return err.response() + + return '' + + def updatecomputeenvironment(self): + compute_env_name = self._get_param('computeEnvironment') + compute_resource = self._get_param('computeResources') + service_role = self._get_param('serviceRole') + state = self._get_param('state') + + try: + name, arn = self.batch_backend.update_compute_environment( + compute_environment_name=compute_env_name, + compute_resources=compute_resource, + service_role=service_role, + state=state + ) + except AWSError as err: + return err.response() + + result = { + 'computeEnvironmentArn': arn, + 'computeEnvironmentName': name + } + + return json.dumps(result) diff --git a/moto/batch/urls.py b/moto/batch/urls.py index 18de99199..ef8f7927a 100644 --- a/moto/batch/urls.py +++ b/moto/batch/urls.py @@ -8,4 +8,6 @@ url_bases = [ url_paths = { '{0}/v1/createcomputeenvironment$': BatchResponse.dispatch, '{0}/v1/describecomputeenvironments$': BatchResponse.dispatch, + '{0}/v1/deletecomputeenvironment': BatchResponse.dispatch, + '{0}/v1/updatecomputeenvironment': BatchResponse.dispatch, } diff --git a/tests/test_batch/test_batch.py b/tests/test_batch/test_batch.py index aceb95804..159f255c3 100644 --- a/tests/test_batch/test_batch.py +++ b/tests/test_batch/test_batch.py @@ -156,3 +156,112 @@ def test_describe_compute_environment(): ) len(resp['computeEnvironments']).should.equal(0) + +@mock_ec2 +@mock_ecs +@mock_iam +@mock_batch +def test_delete_unmanaged_compute_environment(): + ec2_client, iam_client, ecs_client, batch_client = _get_clients() + vpc_id, subnet_id, sg_id, iam_arn = _setup(ec2_client, iam_client) + + compute_name = 'test_compute_env' + batch_client.create_compute_environment( + computeEnvironmentName=compute_name, + type='UNMANAGED', + state='ENABLED', + serviceRole=iam_arn + ) + + batch_client.delete_compute_environment( + computeEnvironment=compute_name, + ) + + resp = batch_client.describe_compute_environments() + len(resp['computeEnvironments']).should.equal(0) + + resp = ecs_client.list_clusters() + len(resp.get('clusterArns', [])).should.equal(0) + + +@mock_ec2 +@mock_ecs +@mock_iam +@mock_batch +def test_delete_managed_compute_environment(): + ec2_client, iam_client, ecs_client, batch_client = _get_clients() + vpc_id, subnet_id, sg_id, iam_arn = _setup(ec2_client, iam_client) + + compute_name = 'test_compute_env' + 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, + 'tags': { + 'string': 'string' + }, + 'bidPercentage': 123, + 'spotIamFleetRole': 'string' + }, + serviceRole=iam_arn + ) + + batch_client.delete_compute_environment( + computeEnvironment=compute_name, + ) + + resp = batch_client.describe_compute_environments() + len(resp['computeEnvironments']).should.equal(0) + + resp = ec2_client.describe_instances() + resp.should.contain('Reservations') + len(resp['Reservations']).should.equal(3) + for reservation in resp['Reservations']: + reservation['Instances'][0]['State']['Name'].should.equal('terminated') + + resp = ecs_client.list_clusters() + len(resp.get('clusterArns', [])).should.equal(0) + + +@mock_ec2 +@mock_ecs +@mock_iam +@mock_batch +def test_update_unmanaged_compute_environment_state(): + ec2_client, iam_client, ecs_client, batch_client = _get_clients() + vpc_id, subnet_id, sg_id, iam_arn = _setup(ec2_client, iam_client) + + compute_name = 'test_compute_env' + 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' + ) + + resp = batch_client.describe_compute_environments() + len(resp['computeEnvironments']).should.equal(1) + resp['computeEnvironments'][0]['state'].should.equal('DISABLED')