diff --git a/moto/ec2/exceptions.py b/moto/ec2/exceptions.py index f17bcbaa1..d54e75868 100644 --- a/moto/ec2/exceptions.py +++ b/moto/ec2/exceptions.py @@ -615,3 +615,13 @@ class InvalidVpcEndPointIdError(EC2ClientError): "InvalidVpcEndPointId.NotFound", "The VpcEndPoint ID '{0}' does not exist".format(vpc_end_point_id), ) + + +class InvalidTaggableResourceType(EC2ClientError): + def __init__(self, resource_type): + super(InvalidTaggableResourceType, self).__init__( + "InvalidParameterValue", + "'{}' is not a valid taggable resource type for this operation.".format( + resource_type + ), + ) diff --git a/moto/ec2/models.py b/moto/ec2/models.py index fa2dce801..dbbfb2838 100644 --- a/moto/ec2/models.py +++ b/moto/ec2/models.py @@ -109,6 +109,7 @@ from .exceptions import ( IncorrectStateIamProfileAssociationError, InvalidAssociationIDIamProfileAssociationError, InvalidVpcEndPointIdError, + InvalidTaggableResourceType, ) from .utils import ( EC2_RESOURCE_TO_PREFIX, @@ -1522,10 +1523,26 @@ class AmiBackend(object): ami_id = ami["ami_id"] self.amis[ami_id] = Ami(self, **ami) - def create_image(self, instance_id, name=None, description=None, context=None): + def create_image( + self, + instance_id, + name=None, + description=None, + context=None, + tag_specifications=None, + ): # TODO: check that instance exists and pull info from it. ami_id = random_ami_id() instance = self.get_instance(instance_id) + tags = [] + for tag_specification in tag_specifications: + resource_type = tag_specification["ResourceType"] + if resource_type == "image": + tags += tag_specification["Tag"] + elif resource_type == "snapshot": + raise NotImplementedError() + else: + raise InvalidTaggableResourceType(resource_type) ami = Ami( self, @@ -1536,6 +1553,8 @@ class AmiBackend(object): description=description, owner_id=OWNER_ID, ) + for tag in tags: + ami.add_tag(tag["Key"], tag["Value"]) self.amis[ami_id] = ami return ami diff --git a/moto/ec2/responses/amis.py b/moto/ec2/responses/amis.py index 178d583e0..079e2cc06 100755 --- a/moto/ec2/responses/amis.py +++ b/moto/ec2/responses/amis.py @@ -8,9 +8,14 @@ class AmisResponse(BaseResponse): name = self.querystring.get("Name")[0] description = self._get_param("Description", if_none="") instance_id = self._get_param("InstanceId") + tag_specifications = self._get_multi_param("TagSpecification") if self.is_not_dryrun("CreateImage"): image = self.ec2_backend.create_image( - instance_id, name, description, context=self + instance_id, + name, + description, + context=self, + tag_specifications=tag_specifications, ) template = self.response_template(CREATE_IMAGE_RESPONSE) return template.render(image=image) diff --git a/requirements-dev.txt b/requirements-dev.txt index 2df056d85..566f2fed2 100644 --- a/requirements-dev.txt +++ b/requirements-dev.txt @@ -9,7 +9,7 @@ flask flask-cors boto>=2.45.0 boto3>=1.4.4 -botocore>=1.18.17 +botocore>=1.19.30 six>=1.9 prompt-toolkit==2.0.10 # 3.x is not available with python2 click==6.7 diff --git a/tests/test_ec2/test_amis.py b/tests/test_ec2/test_amis.py index 57e64cb55..8ea4cb865 100644 --- a/tests/test_ec2/test_amis.py +++ b/tests/test_ec2/test_amis.py @@ -876,3 +876,48 @@ def test_ami_snapshots_have_correct_owner(): for snapshot in snapshots_rseponse["Snapshots"]: assert owner_id == snapshot["OwnerId"] + + +@mock_ec2 +def test_create_image_with_tag_specification(): + ec2 = boto3.resource("ec2", region_name="us-west-1") + client = boto3.client("ec2", region_name="us-west-1") + tag_specifications = [ + { + "ResourceType": "image", + "Tags": [ + { + "Key": "Base_AMI_Name", + "Value": "Deep Learning Base AMI (Amazon Linux 2) Version 31.0", + }, + {"Key": "OS_Version", "Value": "AWS Linux 2",}, + ], + }, + ] + instance = ec2.create_instances(ImageId=EXAMPLE_AMI_ID, MinCount=1, MaxCount=1)[0] + image_id = client.create_image( + InstanceId=instance.instance_id, + Name="test-image", + Description="test ami", + TagSpecifications=tag_specifications, + )["ImageId"] + + image = client.describe_images(ImageIds=[image_id])["Images"][0] + image["Tags"].should.equal(tag_specifications[0]["Tags"]) + + with pytest.raises(ClientError) as ex: + client.create_image( + InstanceId=instance.instance_id, + Name="test-image", + Description="test ami", + TagSpecifications=[ + { + "ResourceType": "invalid-resource-type", + "Tags": [{"Key": "key", "Value": "value"}], + } + ], + ) + ex.value.response["Error"]["Code"].should.equal("InvalidParameterValue") + ex.value.response["Error"]["Message"].should.equal( + "'invalid-resource-type' is not a valid taggable resource type for this operation." + )