diff --git a/moto/batch/models.py b/moto/batch/models.py index e336a60d..6eb02c39 100644 --- a/moto/batch/models.py +++ b/moto/batch/models.py @@ -454,6 +454,64 @@ class BatchBackend(BaseBackend): return result + def update_job_queue(self, queue_name, priority, state, compute_env_order): + """ + Update a job queue + + :param queue_name: Queue name + :type queue_name: str + :param priority: Queue priority + :type priority: int + :param state: Queue state + :type state: string + :param compute_env_order: Compute environment list + :type compute_env_order: list of dict + :return: Tuple of Name, ARN + :rtype: tuple of str + """ + if queue_name is None: + raise ClientException('jobQueueName must be provided') + + job_queue = self.get_job_queue(queue_name) + if job_queue is None: + raise ClientException('Job queue {0} does not exist'.format(queue_name)) + + if state is not None: + if state not in ('ENABLED', 'DISABLED'): + raise ClientException('state {0} must be one of ENABLED | DISABLED'.format(state)) + + job_queue.state = state + + if compute_env_order is not None: + if len(compute_env_order) == 0: + raise ClientException('At least 1 compute environment must be provided') + try: + # orders and extracts computeEnvironment names + ordered_compute_environments = [item['computeEnvironment'] for item in sorted(compute_env_order, key=lambda x: x['order'])] + env_objects = [] + # Check each ARN exists, then make a list of compute env's + for arn in ordered_compute_environments: + env = self.get_compute_environment_by_arn(arn) + if env is None: + raise ClientException('Compute environment {0} does not exist'.format(arn)) + env_objects.append(env) + except Exception: + raise ClientException('computeEnvironmentOrder is malformed') + + job_queue.env_order_json = compute_env_order + job_queue.environments = env_objects + + if priority is not None: + job_queue.priority = priority + + return queue_name, job_queue.arn + + def delete_job_queue(self, queue_name): + job_queue = self.get_job_queue(queue_name) + + if job_queue is not None: + del self._job_queues[job_queue.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 661b9c7c..7c870382 100644 --- a/moto/batch/responses.py +++ b/moto/batch/responses.py @@ -146,3 +146,35 @@ class BatchResponse(BaseResponse): result = {'jobQueues': queues} return json.dumps(result) + + # UpdateJobQueue + def updatejobqueue(self): + compute_env_order = self._get_param('computeEnvironmentOrder') + queue_name = self._get_param('jobQueue') + priority = self._get_param('priority') + state = self._get_param('state') + + try: + name, arn = self.batch_backend.update_job_queue( + queue_name=queue_name, + priority=priority, + state=state, + compute_env_order=compute_env_order + ) + except AWSError as err: + return err.response() + + result = { + 'jobQueueArn': arn, + 'jobQueueName': name + } + + return json.dumps(result) + + # DeleteJobQueue + def deletejobqueue(self): + queue_name = self._get_param('jobQueue') + + self.batch_backend.delete_job_queue(queue_name) + + return '' diff --git a/moto/batch/urls.py b/moto/batch/urls.py index 227e78ec..bc186bd2 100644 --- a/moto/batch/urls.py +++ b/moto/batch/urls.py @@ -12,4 +12,6 @@ url_paths = { '{0}/v1/updatecomputeenvironment': BatchResponse.dispatch, '{0}/v1/createjobqueue': BatchResponse.dispatch, '{0}/v1/describejobqueues': BatchResponse.dispatch, + '{0}/v1/updatejobqueue': BatchResponse.dispatch, + '{0}/v1/deletejobqueue': BatchResponse.dispatch } diff --git a/tests/test_batch/test_batch.py b/tests/test_batch/test_batch.py index b082d656..e7c4cf62 100644 --- a/tests/test_batch/test_batch.py +++ b/tests/test_batch/test_batch.py @@ -341,3 +341,83 @@ def test_job_queue_bad_arn(): ) except ClientError as err: err.response['Error']['Code'].should.equal('ClientException') + + +@mock_ec2 +@mock_ecs +@mock_iam +@mock_batch +def test_update_job_queue(): + 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' + resp = batch_client.create_compute_environment( + computeEnvironmentName=compute_name, + type='UNMANAGED', + state='ENABLED', + serviceRole=iam_arn + ) + arn = resp['computeEnvironmentArn'] + + resp = batch_client.create_job_queue( + jobQueueName='test_job_queue', + state='ENABLED', + priority=123, + computeEnvironmentOrder=[ + { + 'order': 123, + 'computeEnvironment': arn + }, + ] + ) + queue_arn = resp['jobQueueArn'] + + batch_client.update_job_queue( + jobQueue=queue_arn, + priority=5 + ) + + resp = batch_client.describe_job_queues() + resp.should.contain('jobQueues') + len(resp['jobQueues']).should.equal(1) + resp['jobQueues'][0]['priority'].should.equal(5) + + +@mock_ec2 +@mock_ecs +@mock_iam +@mock_batch +def test_update_job_queue(): + 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' + resp = batch_client.create_compute_environment( + computeEnvironmentName=compute_name, + type='UNMANAGED', + state='ENABLED', + serviceRole=iam_arn + ) + arn = resp['computeEnvironmentArn'] + + resp = batch_client.create_job_queue( + jobQueueName='test_job_queue', + state='ENABLED', + priority=123, + computeEnvironmentOrder=[ + { + 'order': 123, + 'computeEnvironment': arn + }, + ] + ) + queue_arn = resp['jobQueueArn'] + + batch_client.delete_job_queue( + jobQueue=queue_arn + ) + + resp = batch_client.describe_job_queues() + resp.should.contain('jobQueues') + len(resp['jobQueues']).should.equal(0)