Fix:EC2-Create-Subnet availability Zone Id support (#3198)

* Fix:EC2-Create-Subnet availability Zone Id support

* Linting

* Fix:fixed build errors

* linting

Co-authored-by: Bert Blommers <bblommers@users.noreply.github.com>
Co-authored-by: Bert Blommers <info@bertblommers.nl>
Co-authored-by: usmankb <usman@krazybee.com>

moto/ec2/models.py

@@ -3402,7 +3402,14 @@ class SubnetBackend(object):
                 return subnets[subnet_id]
         raise InvalidSubnetIdError(subnet_id)
 
-    def create_subnet(self, vpc_id, cidr_block, availability_zone, context=None):
+    def create_subnet(
+        self,
+        vpc_id,
+        cidr_block,
+        availability_zone=None,
+        availability_zone_id=None,
+        context=None,
+    ):
         subnet_id = random_subnet_id()
         vpc = self.get_vpc(
             vpc_id
@@ -3430,15 +3437,25 @@ class SubnetBackend(object):
         # consider it the default
         default_for_az = str(availability_zone not in self.subnets).lower()
         map_public_ip_on_launch = default_for_az
-        if availability_zone is None:
+
+        if availability_zone is None and not availability_zone_id:
             availability_zone = "us-east-1a"
         try:
-            availability_zone_data = next(
-                zone
-                for zones in RegionsAndZonesBackend.zones.values()
-                for zone in zones
-                if zone.name == availability_zone
-            )
+            if availability_zone:
+                availability_zone_data = next(
+                    zone
+                    for zones in RegionsAndZonesBackend.zones.values()
+                    for zone in zones
+                    if zone.name == availability_zone
+                )
+            elif availability_zone_id:
+                availability_zone_data = next(
+                    zone
+                    for zones in RegionsAndZonesBackend.zones.values()
+                    for zone in zones
+                    if zone.zone_id == availability_zone_id
+                )
+
         except StopIteration:
             raise InvalidAvailabilityZoneError(
                 availability_zone,

moto/ec2/responses/subnets.py

@@ -9,12 +9,14 @@ class Subnets(BaseResponse):
     def create_subnet(self):
         vpc_id = self._get_param("VpcId")
         cidr_block = self._get_param("CidrBlock")
-        availability_zone = self._get_param(
-            "AvailabilityZone",
-            if_none=random.choice(self.ec2_backend.describe_availability_zones()).name,
-        )
+        availability_zone = self._get_param("AvailabilityZone")
+        availability_zone_id = self._get_param("AvailabilityZoneId")
+        if not availability_zone and not availability_zone_id:
+            availability_zone = random.choice(
+                self.ec2_backend.describe_availability_zones()
+            ).name
         subnet = self.ec2_backend.create_subnet(
-            vpc_id, cidr_block, availability_zone, context=self
+            vpc_id, cidr_block, availability_zone, availability_zone_id, context=self
         )
         template = self.response_template(CREATE_SUBNET_RESPONSE)
         return template.render(subnet=subnet)

tests/test_ec2/test_subnets.py

@@ -75,6 +75,18 @@ def test_subnet_should_have_proper_availability_zone_set():
     subnetA.availability_zone.should.equal("us-west-1b")
 
 
+@mock_ec2
+def test_availability_zone_in_create_subnet():
+    ec2 = boto3.resource("ec2", region_name="us-west-1")
+
+    vpc = ec2.create_vpc(CidrBlock="172.31.0.0/16")
+
+    subnet = ec2.create_subnet(
+        VpcId=vpc.id, CidrBlock="172.31.48.0/20", AvailabilityZoneId="use1-az6"
+    )
+    subnet.availability_zone_id.should.equal("use1-az6")
+
+
 @mock_ec2
 def test_default_subnet():
     ec2 = boto3.resource("ec2", region_name="us-west-1")
@@ -612,7 +624,15 @@ def test_run_instances_should_attach_to_default_subnet():
     # Assert subnet is created appropriately
     subnets = client.describe_subnets()["Subnets"]
     default_subnet_id = subnets[0]["SubnetId"]
-    instances["Instances"][0]["NetworkInterfaces"][0]["SubnetId"].should.equal(
-        default_subnet_id
+    if len(subnets) > 1:
+        default_subnet_id1 = subnets[1]["SubnetId"]
+    assert (
+        instances["Instances"][0]["NetworkInterfaces"][0]["SubnetId"]
+        == default_subnet_id
+        or instances["Instances"][0]["NetworkInterfaces"][0]["SubnetId"]
+        == default_subnet_id1
+    )
+    assert (
+        subnets[0]["AvailableIpAddressCount"] == 4090
+        or subnets[1]["AvailableIpAddressCount"] == 4090
     )
-    subnets[0]["AvailableIpAddressCount"].should.equal(4090)