Batch - JobQueue improvements + Tag support

IMPLEMENTATION_COVERAGE.md

@@ -404,7 +404,7 @@
 
 ## batch
 <details>
-<summary>66% implemented</summary>
+<summary>79% implemented</summary>
 
 - [X] cancel_job
 - [X] create_compute_environment
@@ -421,12 +421,12 @@
 - [ ] describe_scheduling_policies
 - [X] list_jobs
 - [ ] list_scheduling_policies
-- [ ] list_tags_for_resource
+- [X] list_tags_for_resource
 - [X] register_job_definition
 - [X] submit_job
-- [ ] tag_resource
+- [X] tag_resource
 - [X] terminate_job
-- [ ] untag_resource
+- [X] untag_resource
 - [X] update_compute_environment
 - [X] update_job_queue
 - [ ] update_scheduling_policy

docs/docs/services/batch.rst

@@ -55,12 +55,12 @@ batch
 - [ ] describe_scheduling_policies
 - [X] list_jobs
 - [ ] list_scheduling_policies
-- [ ] list_tags_for_resource
+- [X] list_tags_for_resource
 - [X] register_job_definition
 - [X] submit_job
-- [ ] tag_resource
+- [X] tag_resource
 - [X] terminate_job
-- [ ] untag_resource
+- [X] untag_resource
 - [X] update_compute_environment
 - [X] update_job_queue
   

moto/batch/models.py

@@ -14,6 +14,7 @@ from moto.iam import iam_backends
 from moto.ec2 import ec2_backends
 from moto.ecs import ecs_backends
 from moto.logs import logs_backends
+from moto.utilities.tagging_service import TaggingService
 
 from .exceptions import InvalidParameterValueException, ClientException, ValidationError
 from .utils import (
@@ -24,6 +25,7 @@ from .utils import (
 )
 from moto.ec2.exceptions import InvalidSubnetIdError
 from moto.ec2.models import INSTANCE_TYPES as EC2_INSTANCE_TYPES
+from moto.ec2.models import INSTANCE_FAMILIES as EC2_INSTANCE_FAMILIES
 from moto.iam.exceptions import IAMNotFoundException
 from moto.core import ACCOUNT_ID as DEFAULT_ACCOUNT_ID
 from moto.core.utils import unix_time_millis, BackendDict
@@ -114,7 +116,15 @@ class ComputeEnvironment(CloudFormationModel):
 
 class JobQueue(CloudFormationModel):
     def __init__(
-        self, name, priority, state, environments, env_order_json, region_name
+        self,
+        name,
+        priority,
+        state,
+        environments,
+        env_order_json,
+        region_name,
+        backend,
+        tags=None,
     ):
         """
         :param name: Job queue name
@@ -137,6 +147,10 @@ class JobQueue(CloudFormationModel):
         self.env_order_json = env_order_json
         self.arn = make_arn_for_job_queue(DEFAULT_ACCOUNT_ID, name, region_name)
         self.status = "VALID"
+        self.backend = backend
+
+        if tags:
+            backend.tag_resource(self.arn, tags)
 
         self.jobs = []
 
@@ -148,6 +162,7 @@ class JobQueue(CloudFormationModel):
             "priority": self.priority,
             "state": self.state,
             "status": self.status,
+            "tags": self.backend.list_tags_for_resource(self.arn),
         }
 
         return result
@@ -184,6 +199,7 @@ class JobQueue(CloudFormationModel):
             priority=properties["Priority"],
             state=properties.get("State", "ENABLED"),
             compute_env_order=compute_envs,
+            backend=backend,
         )
         arn = queue[1]
 
@@ -732,6 +748,7 @@ class BatchBackend(BaseBackend):
     def __init__(self, region_name=None):
         super().__init__()
         self.region_name = region_name
+        self.tagger = TaggingService()
 
         self._compute_environments = {}
         self._job_queues = {}
@@ -1054,7 +1071,10 @@ class BatchBackend(BaseBackend):
             for instance_type in cr["instanceTypes"]:
                 if instance_type == "optimal":
                     pass  # Optimal should pick from latest of current gen
-                elif instance_type not in EC2_INSTANCE_TYPES:
+                elif (
+                    instance_type not in EC2_INSTANCE_TYPES
+                    and instance_type not in EC2_INSTANCE_FAMILIES
+                ):
                     raise InvalidParameterValueException(
                         "Instance type {0} does not exist".format(instance_type)
                     )
@@ -1104,6 +1124,12 @@ class BatchBackend(BaseBackend):
             if instance_type == "optimal":
                 instance_type = "m4.4xlarge"
 
+            if "." not in instance_type:
+                # instance_type can be a family of instance types (c2, t3, etc)
+                # We'll just use the first instance_type in this family
+                instance_type = [
+                    i for i in EC2_INSTANCE_TYPES.keys() if i.startswith(instance_type)
+                ][0]
             instance_vcpus.append(
                 (
                     EC2_INSTANCE_TYPES[instance_type]["VCpuInfo"]["DefaultVCpus"],
@@ -1190,7 +1216,9 @@ class BatchBackend(BaseBackend):
 
         return compute_env.name, compute_env.arn
 
-    def create_job_queue(self, queue_name, priority, state, compute_env_order):
+    def create_job_queue(
+        self, queue_name, priority, state, compute_env_order, tags=None
+    ):
         """
         Create a job queue
 
@@ -1249,6 +1277,8 @@ class BatchBackend(BaseBackend):
             env_objects,
             compute_env_order,
             self.region_name,
+            backend=self,
+            tags=tags,
         )
         self._job_queues[queue.arn] = queue
 
@@ -1516,5 +1546,15 @@ class BatchBackend(BaseBackend):
 
         job.terminate(reason)
 
+    def tag_resource(self, resource_arn, tags):
+        tags = self.tagger.convert_dict_to_tags_input(tags or {})
+        self.tagger.tag_resource(resource_arn, tags)
+
+    def list_tags_for_resource(self, resource_arn):
+        return self.tagger.get_tag_dict_for_resource(resource_arn)
+
+    def untag_resource(self, resource_arn, tag_keys):
+        self.tagger.untag_resource_using_names(resource_arn, tag_keys)
+
 
 batch_backends = BackendDict(BatchBackend, "batch")

moto/batch/responses.py

@@ -1,6 +1,6 @@
 from moto.core.responses import BaseResponse
 from .models import batch_backends
-from urllib.parse import urlsplit
+from urllib.parse import urlsplit, unquote
 
 from .exceptions import AWSError
 
@@ -114,6 +114,7 @@ class BatchResponse(BaseResponse):
         queue_name = self._get_param("jobQueueName")
         priority = self._get_param("priority")
         state = self._get_param("state")
+        tags = self._get_param("tags")
 
         try:
             name, arn = self.batch_backend.create_job_queue(
@@ -121,6 +122,7 @@ class BatchResponse(BaseResponse):
                 priority=priority,
                 state=state,
                 compute_env_order=compute_env_order,
+                tags=tags,
             )
         except AWSError as err:
             return err.response()
@@ -298,3 +300,16 @@ class BatchResponse(BaseResponse):
         self.batch_backend.cancel_job(job_id, reason)
 
         return ""
+
+    def tags(self):
+        resource_arn = unquote(self.path).split("/v1/tags/")[-1]
+        tags = self._get_param("tags")
+        if self.method == "POST":
+            self.batch_backend.tag_resource(resource_arn, tags)
+            return ""
+        if self.method == "GET":
+            tags = self.batch_backend.list_tags_for_resource(resource_arn)
+            return json.dumps({"tags": tags})
+        if self.method == "DELETE":
+            tag_keys = self.querystring.get("tagKeys")
+            self.batch_backend.untag_resource(resource_arn, tag_keys)

moto/batch/urls.py

@@ -19,4 +19,6 @@ url_paths = {
     "{0}/v1/listjobs": BatchResponse.dispatch,
     "{0}/v1/terminatejob": BatchResponse.dispatch,
     "{0}/v1/canceljob": BatchResponse.dispatch,
+    "{0}/v1/tags/(?P<arn_part_1>[^/]+)/(?P<arn_part_2>[^/]+)/?$": BatchResponse.dispatch,
+    "{0}/v1/tags/(?P<arn>[^/]+)/?$": BatchResponse.dispatch,
 }

moto/ec2/models.py

@@ -192,6 +192,7 @@ from .utils import (
 )
 
 INSTANCE_TYPES = load_resource(__name__, "resources/instance_types.json")
+INSTANCE_FAMILIES = list(set([i.split(".")[0] for i in INSTANCE_TYPES.keys()]))
 
 root = pathlib.Path(__file__).parent
 offerings_path = "resources/instance_type_offerings"

tests/terraform-tests.success.txt

@@ -9,6 +9,7 @@ TestAccAWSAPIGatewayV2VpcLink
 TestAccAWSAppsyncApiKey
 TestAccAWSAppsyncGraphqlApi
 TestAccAWSAvailabilityZones
+TestAccAWSBatchJobQueue
 TestAccAWSBillingServiceAccount
 TestAccAWSCallerIdentity
 TestAccAWSCloudTrailServiceAccount

tests/test_batch/test_batch_compute_envs.py

@@ -1,6 +1,7 @@
 from . import _get_clients, _setup
 import pytest
 import sure  # noqa # pylint: disable=unused-import
+from botocore.exceptions import ClientError
 from moto import mock_batch, mock_iam, mock_ec2, mock_ecs, settings
 from uuid import uuid4
 
@@ -55,6 +56,93 @@ def test_create_managed_compute_environment():
     all_clusters.should.contain(our_env["ecsClusterArn"])
 
 
+@mock_ec2
+@mock_ecs
+@mock_iam
+@mock_batch
+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]
+    our_env["computeResources"]["instanceTypes"].should.equal(["t2"])
+
+
+@mock_ec2
+@mock_ecs
+@mock_iam
+@mock_batch
+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"]
+    err["Code"].should.equal("InvalidParameterValue")
+    err["Message"].should.equal("Instance type unknown does not exist")
+
+
 @mock_ec2
 @mock_ecs
 @mock_iam

tests/test_batch/test_batch_tags.py

@@ -0,0 +1,137 @@
+from . import _get_clients, _setup
+
+import sure  # noqa # pylint: disable=unused-import
+from moto import mock_batch, mock_iam, mock_ec2, mock_ecs
+from uuid import uuid4
+
+
+@mock_ec2
+@mock_ecs
+@mock_iam
+@mock_batch
+def test_create_job_queue_with_tags():
+    ec2_client, iam_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,
+    )
+    arn = resp["computeEnvironmentArn"]
+
+    jq_name = str(uuid4())[0:6]
+    resp = batch_client.create_job_queue(
+        jobQueueName=jq_name,
+        state="ENABLED",
+        priority=123,
+        computeEnvironmentOrder=[{"order": 123, "computeEnvironment": arn}],
+        tags={"k1": "v1", "k2": "v2"},
+    )
+    resp.should.contain("jobQueueArn")
+    resp.should.contain("jobQueueName")
+    queue_arn = resp["jobQueueArn"]
+
+    my_queue = batch_client.describe_job_queues(jobQueues=[queue_arn])["jobQueues"][0]
+    my_queue.should.have.key("tags").equals({"k1": "v1", "k2": "v2"})
+
+
+@mock_ec2
+@mock_ecs
+@mock_iam
+@mock_batch
+def test_list_tags():
+    ec2_client, iam_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,
+    )
+    arn = resp["computeEnvironmentArn"]
+
+    jq_name = str(uuid4())[0:6]
+    resp = batch_client.create_job_queue(
+        jobQueueName=jq_name,
+        state="ENABLED",
+        priority=123,
+        computeEnvironmentOrder=[{"order": 123, "computeEnvironment": arn}],
+        tags={"k1": "v1", "k2": "v2"},
+    )
+    resp.should.contain("jobQueueArn")
+    resp.should.contain("jobQueueName")
+    queue_arn = resp["jobQueueArn"]
+
+    my_queue = batch_client.list_tags_for_resource(resourceArn=queue_arn)
+    my_queue.should.have.key("tags").equals({"k1": "v1", "k2": "v2"})
+
+
+@mock_ec2
+@mock_ecs
+@mock_iam
+@mock_batch
+def test_tag_job_queue():
+    ec2_client, iam_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,
+    )
+    arn = resp["computeEnvironmentArn"]
+
+    jq_name = str(uuid4())[0:6]
+    resp = batch_client.create_job_queue(
+        jobQueueName=jq_name,
+        state="ENABLED",
+        priority=123,
+        computeEnvironmentOrder=[{"order": 123, "computeEnvironment": arn}],
+    )
+    queue_arn = resp["jobQueueArn"]
+
+    batch_client.tag_resource(resourceArn=queue_arn, tags={"k1": "v1", "k2": "v2"})
+
+    my_queue = batch_client.list_tags_for_resource(resourceArn=queue_arn)
+    my_queue.should.have.key("tags").equals({"k1": "v1", "k2": "v2"})
+
+
+@mock_ec2
+@mock_ecs
+@mock_iam
+@mock_batch
+def test_untag_job_queue():
+    ec2_client, iam_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,
+    )
+    arn = resp["computeEnvironmentArn"]
+
+    jq_name = str(uuid4())[0:6]
+    resp = batch_client.create_job_queue(
+        jobQueueName=jq_name,
+        state="ENABLED",
+        priority=123,
+        computeEnvironmentOrder=[{"order": 123, "computeEnvironment": arn}],
+        tags={"k1": "v1", "k2": "v2"},
+    )
+    queue_arn = resp["jobQueueArn"]
+
+    batch_client.tag_resource(resourceArn=queue_arn, tags={"k3": "v3"})
+    batch_client.untag_resource(resourceArn=queue_arn, tagKeys=["k2"])
+
+    my_queue = batch_client.list_tags_for_resource(resourceArn=queue_arn)
+    my_queue.should.have.key("tags").equals({"k1": "v1", "k3": "v3"})