Added DeleteComputeEnvironment and UpdateComputeEnvironment

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

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}

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)

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,
 }

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')