From 4e45152d86fa4d499a3a6fcdaed6dea4839314ef Mon Sep 17 00:00:00 2001
From: Bert Blommers <bblommers@users.noreply.github.com>
Date: Sun, 10 Oct 2021 19:16:28 +0000
Subject: [PATCH] EC2:create_volume() - Add support for VolumeType (#4264)

---
 moto/ec2/models.py                         | 31 ++++++++++++++++++----
 moto/ec2/responses/elastic_block_store.py  | 12 ++++++---
 tests/test_ec2/test_elastic_block_store.py | 23 ++++++++++++++++
 3 files changed, 58 insertions(+), 8 deletions(-)

diff --git a/moto/ec2/models.py b/moto/ec2/models.py
index c2952f875..0b6c3e2fb 100644
--- a/moto/ec2/models.py
+++ b/moto/ec2/models.py
@@ -685,7 +685,11 @@ class Instance(TaggedEC2Resource, BotoInstance, CloudFormationModel):
         kms_key_id=None,
     ):
         volume = self.ec2_backend.create_volume(
-            size, self.region_name, snapshot_id, encrypted, kms_key_id
+            size=size,
+            zone_name=self.region_name,
+            snapshot_id=snapshot_id,
+            encrypted=encrypted,
+            kms_key_id=kms_key_id,
         )
         self.ec2_backend.attach_volume(
             volume.id, self.id, device_path, delete_on_termination
@@ -694,7 +698,7 @@ class Instance(TaggedEC2Resource, BotoInstance, CloudFormationModel):
     def setup_defaults(self):
         # Default have an instance with root volume should you not wish to
         # override with attach volume cmd.
-        volume = self.ec2_backend.create_volume(8, "us-east-1a")
+        volume = self.ec2_backend.create_volume(size=8, zone_name="us-east-1a")
         self.ec2_backend.attach_volume(volume.id, self.id, "/dev/sda1", True)
 
     def teardown_defaults(self):
@@ -1572,7 +1576,7 @@ class Ami(TaggedEC2Resource):
             self.launch_permission_groups.add("all")
 
         # AWS auto-creates these, we should reflect the same.
-        volume = self.ec2_backend.create_volume(15, region_name)
+        volume = self.ec2_backend.create_volume(size=15, zone_name=region_name)
         self.ebs_snapshot = self.ec2_backend.create_snapshot(
             volume.id, "Auto-created snapshot for AMI %s" % self.id, owner_id
         )
@@ -3209,8 +3213,10 @@ class Volume(TaggedEC2Resource, CloudFormationModel):
         snapshot_id=None,
         encrypted=False,
         kms_key_id=None,
+        volume_type=None,
     ):
         self.id = volume_id
+        self.volume_type = volume_type or "gp2"
         self.size = size
         self.zone = zone
         self.create_time = utc_date_and_time()
@@ -3331,7 +3337,13 @@ class EBSBackend(object):
         super().__init__()
 
     def create_volume(
-        self, size, zone_name, snapshot_id=None, encrypted=False, kms_key_id=None
+        self,
+        size,
+        zone_name,
+        snapshot_id=None,
+        encrypted=False,
+        kms_key_id=None,
+        volume_type=None,
     ):
         if kms_key_id and not encrypted:
             raise InvalidParameterDependency("KmsKeyId", "Encrypted")
@@ -3345,7 +3357,16 @@ class EBSBackend(object):
                 size = snapshot.volume.size
             if snapshot.encrypted:
                 encrypted = snapshot.encrypted
-        volume = Volume(self, volume_id, size, zone, snapshot_id, encrypted, kms_key_id)
+        volume = Volume(
+            self,
+            volume_id=volume_id,
+            size=size,
+            zone=zone,
+            snapshot_id=snapshot_id,
+            encrypted=encrypted,
+            kms_key_id=kms_key_id,
+            volume_type=volume_type,
+        )
         self.volumes[volume_id] = volume
         return volume
 
diff --git a/moto/ec2/responses/elastic_block_store.py b/moto/ec2/responses/elastic_block_store.py
index ac6492bf5..c38c1e233 100644
--- a/moto/ec2/responses/elastic_block_store.py
+++ b/moto/ec2/responses/elastic_block_store.py
@@ -44,13 +44,19 @@ class ElasticBlockStore(BaseResponse):
         size = self._get_param("Size")
         zone = self._get_param("AvailabilityZone")
         snapshot_id = self._get_param("SnapshotId")
+        volume_type = self._get_param("VolumeType")
         tags = self._parse_tag_specification("TagSpecification")
         volume_tags = tags.get("volume", {})
         encrypted = self._get_bool_param("Encrypted", if_none=False)
         kms_key_id = self._get_param("KmsKeyId")
         if self.is_not_dryrun("CreateVolume"):
             volume = self.ec2_backend.create_volume(
-                size, zone, snapshot_id, encrypted, kms_key_id
+                size=size,
+                zone_name=zone,
+                snapshot_id=snapshot_id,
+                encrypted=encrypted,
+                kms_key_id=kms_key_id,
+                volume_type=volume_type,
             )
             volume.add_tags(volume_tags)
             template = self.response_template(CREATE_VOLUME_RESPONSE)
@@ -183,7 +189,7 @@ CREATE_VOLUME_RESPONSE = """<CreateVolumeResponse xmlns="http://ec2.amazonaws.co
       {% endfor %}
     </tagSet>
   {% endif %}
-  <volumeType>standard</volumeType>
+  <volumeType>{{ volume.volume_type }}</volumeType>
 </CreateVolumeResponse>"""
 
 DESCRIBE_VOLUMES_RESPONSE = """<DescribeVolumesResponse xmlns="http://ec2.amazonaws.com/doc/2013-10-15/">
@@ -229,7 +235,7 @@ DESCRIBE_VOLUMES_RESPONSE = """<DescribeVolumesResponse xmlns="http://ec2.amazon
                  {% endfor %}
                </tagSet>
              {% endif %}
-             <volumeType>standard</volumeType>
+             <volumeType>{{ volume.volume_type }}</volumeType>
           </item>
       {% endfor %}
    </volumeSet>
diff --git a/tests/test_ec2/test_elastic_block_store.py b/tests/test_ec2/test_elastic_block_store.py
index f3c7f908f..8baf0e86d 100644
--- a/tests/test_ec2/test_elastic_block_store.py
+++ b/tests/test_ec2/test_elastic_block_store.py
@@ -1543,3 +1543,26 @@ def test_kms_key_id_property_hidden_when_volume_not_encrypted():
     )
     volume.encrypted.should.be.false
     volume.kms_key_id.should.be.none
+
+
+@mock_ec2
+def test_create_volume_with_standard_type():
+    ec2 = boto3.client("ec2", region_name="us-east-1")
+    volume = ec2.create_volume(AvailabilityZone="us-east-1a", Size=100, Iops=1000)
+    volume["VolumeType"].should.equal("gp2")
+
+    volume = ec2.describe_volumes(VolumeIds=[volume["VolumeId"]])["Volumes"][0]
+    volume["VolumeType"].should.equal("gp2")
+
+
+@pytest.mark.parametrize("volume_type", ["gp2", "gp3", "io1", "io2", "standard"])
+@mock_ec2
+def test_create_volume_with_non_standard_type(volume_type):
+    ec2 = boto3.client("ec2", region_name="us-east-1")
+    volume = ec2.create_volume(
+        AvailabilityZone="us-east-1a", Size=100, Iops=1000, VolumeType=volume_type
+    )
+    volume["VolumeType"].should.equal(volume_type)
+
+    volume = ec2.describe_volumes(VolumeIds=[volume["VolumeId"]])["Volumes"][0]
+    volume["VolumeType"].should.equal(volume_type)