Merge pull request #341 from spulec/add-availabilityZone-filter-to-subnet

Improve availability zone Subnet support

moto/ec2/models.py

@@ -1695,41 +1695,66 @@ class VPCPeeringConnectionBackend(object):
 
 
 class Subnet(TaggedEC2Resource):
-    def __init__(self, ec2_backend, subnet_id, vpc_id, cidr_block):
+    def __init__(self, ec2_backend, subnet_id, vpc_id, cidr_block, availability_zone):
         self.ec2_backend = ec2_backend
         self.id = subnet_id
         self.vpc_id = vpc_id
         self.cidr_block = cidr_block
+        self._availability_zone = availability_zone
 
     @classmethod
     def create_from_cloudformation_json(cls, resource_name, cloudformation_json, region_name):
         properties = cloudformation_json['Properties']
 
         vpc_id = properties['VpcId']
+        cidr_block = properties['CidrBlock']
+        availability_zone = properties.get('AvailabilityZone')
         ec2_backend = ec2_backends[region_name]
         subnet = ec2_backend.create_subnet(
             vpc_id=vpc_id,
-            cidr_block=properties['CidrBlock']
+            cidr_block=cidr_block,
+            availability_zone=availability_zone,
         )
         return subnet
 
     @property
     def availability_zone(self):
-        # This could probably be smarter, but there doesn't appear to be a
+        if self._availability_zone is None:
-        # way to pull AZs for a region in boto
+            # This could probably be smarter, but there doesn't appear to be a
-        return self.ec2_backend.region_name + "a"
+            # way to pull AZs for a region in boto
+            return self.ec2_backend.region_name + "a"
+        else:
+            return self._availability_zone
 
     @property
     def physical_resource_id(self):
         return self.id
 
     def get_filter_value(self, filter_name):
+        """
+        API Version 2014-10-01 defines the following filters for DescribeSubnets:
+
+        * availabilityZone
+        * available-ip-address-count
+        * cidrBlock
+        * defaultForAz
+        * state
+        * subnet-id
+        * tag:key=value
+        * tag-key
+        * tag-value
+        * vpc-id
+
+        Taken from: http://docs.aws.amazon.com/AWSEC2/latest/APIReference/API_DescribeSubnets.html
+        """
         if filter_name in ['cidr', 'cidrBlock', 'cidr-block']:
             return self.cidr_block
         elif filter_name == 'vpc-id':
             return self.vpc_id
         elif filter_name == 'subnet-id':
             return self.id
+        elif filter_name == 'availabilityZone':
+            return self.availability_zone
 
         filter_value = super(Subnet, self).get_filter_value(filter_name)
 
@@ -1756,9 +1781,9 @@ class SubnetBackend(object):
             raise InvalidSubnetIdError(subnet_id)
         return subnet
 
-    def create_subnet(self, vpc_id, cidr_block):
+    def create_subnet(self, vpc_id, cidr_block, availability_zone=None):
         subnet_id = random_subnet_id()
-        subnet = Subnet(self, subnet_id, vpc_id, cidr_block)
+        subnet = Subnet(self, subnet_id, vpc_id, cidr_block, availability_zone)
         self.get_vpc(vpc_id)  # Validate VPC exists
 
         # AWS associates a new subnet with the default Network ACL

moto/ec2/responses/subnets.py

@@ -7,7 +7,15 @@ class Subnets(BaseResponse):
     def create_subnet(self):
         vpc_id = self.querystring.get('VpcId')[0]
         cidr_block = self.querystring.get('CidrBlock')[0]
-        subnet = self.ec2_backend.create_subnet(vpc_id, cidr_block)
+        if 'AvailabilityZone' in self.querystring:
+            availability_zone = self.querystring['AvailabilityZone'][0]
+        else:
+            availability_zone = None
+        subnet = self.ec2_backend.create_subnet(
+          vpc_id,
+          cidr_block,
+          availability_zone,
+        )
         template = self.response_template(CREATE_SUBNET_RESPONSE)
         return template.render(subnet=subnet)
 
@@ -33,7 +41,7 @@ CREATE_SUBNET_RESPONSE = """
     <vpcId>{{ subnet.vpc_id }}</vpcId>
     <cidrBlock>{{ subnet.cidr_block }}</cidrBlock>
     <availableIpAddressCount>251</availableIpAddressCount>
-    <availabilityZone>us-east-1a</availabilityZone>
+    <availabilityZone>{{ subnet.availability_zone }}</availabilityZone>
     <tagSet>
       {% for tag in subnet.get_tags() %}
         <item>
@@ -64,7 +72,7 @@ DESCRIBE_SUBNETS_RESPONSE = """
         <vpcId>{{ subnet.vpc_id }}</vpcId>
         <cidrBlock>{{ subnet.cidr_block }}</cidrBlock>
         <availableIpAddressCount>251</availableIpAddressCount>
-        <availabilityZone>us-east-1a</availabilityZone>
+        <availabilityZone>{{ subnet.availability_zone }}</availabilityZone>
         <tagSet>
           {% for tag in subnet.get_tags() %}
             <item>

tests/test_cloudformation/test_cloudformation_stack_integration.py

@@ -1189,3 +1189,32 @@ def test_security_group_ingress_separate_from_security_group_by_id_using_vpc():
     security_group1.rules[0].ip_protocol.should.equal('tcp')
     security_group1.rules[0].from_port.should.equal('80')
     security_group1.rules[0].to_port.should.equal('8080')
+
+
+@mock_cloudformation
+@mock_ec2
+def test_subnets_should_be_created_with_availability_zone():
+    vpc_conn = boto.vpc.connect_to_region('us-west-1')
+    vpc = vpc_conn.create_vpc("10.0.0.0/16")
+
+    subnet_template = {
+       "AWSTemplateFormatVersion" : "2010-09-09",
+       "Resources" : {
+          "testSubnet" : {
+             "Type" : "AWS::EC2::Subnet",
+             "Properties" : {
+                "VpcId" : vpc.id,
+                "CidrBlock" : "10.0.0.0/24",
+                "AvailabilityZone" : "us-west-1b",
+             }
+          }
+       }
+    }
+    cf_conn = boto.cloudformation.connect_to_region("us-west-1")
+    template_json = json.dumps(subnet_template)
+    cf_conn.create_stack(
+        "test_stack",
+        template_body=template_json,
+    )
+    subnet = vpc_conn.get_all_subnets(filters={'cidrBlock': '10.0.0.0/24'})[0]
+    subnet.availability_zone.should.equal('us-west-1b')

tests/test_ec2/test_subnets.py

@@ -4,6 +4,7 @@ import tests.backport_assert_raises
 from nose.tools import assert_raises
 
 import boto
+import boto.vpc
 from boto.exception import EC2ResponseError
 import sure  # noqa
 
@@ -61,13 +62,21 @@ def test_subnet_tagging():
 
 
 @mock_ec2
-def test_get_subnets_filtering():
+def test_subnet_should_have_proper_availability_zone_set():
-    conn = boto.connect_vpc('the_key', 'the_secret')
+    conn = boto.vpc.connect_to_region('us-west-1')
     vpcA = conn.create_vpc("10.0.0.0/16")
-    subnetA = conn.create_subnet(vpcA.id, "10.0.0.0/24")
+    subnetA = conn.create_subnet(vpcA.id, "10.0.0.0/24", availability_zone='us-west-1b')
+    subnetA.availability_zone.should.equal('us-west-1b')
+
+
+@mock_ec2
+def test_get_subnets_filtering():
+    conn = boto.vpc.connect_to_region('us-west-1')
+    vpcA = conn.create_vpc("10.0.0.0/16")
+    subnetA = conn.create_subnet(vpcA.id, "10.0.0.0/24", availability_zone='us-west-1a')
     vpcB = conn.create_vpc("10.0.0.0/16")
-    subnetB1 = conn.create_subnet(vpcB.id, "10.0.0.0/24")
+    subnetB1 = conn.create_subnet(vpcB.id, "10.0.0.0/24", availability_zone='us-west-1a')
-    subnetB2 = conn.create_subnet(vpcB.id, "10.0.1.0/24")
+    subnetB2 = conn.create_subnet(vpcB.id, "10.0.1.0/24", availability_zone='us-west-1b')
 
     all_subnets = conn.get_all_subnets()
     all_subnets.should.have.length_of(3)
@@ -100,5 +109,10 @@ def test_get_subnets_filtering():
     subnets_by_id.should.have.length_of(1)
     set([subnet.id for subnet in subnets_by_id]).should.equal(set([subnetA.id]))
 
+    # Filter by availabilityZone
+    subnets_by_az = conn.get_all_subnets(filters={'availabilityZone': 'us-west-1a', 'vpc-id': vpcB.id})
+    subnets_by_az.should.have.length_of(1)
+    set([subnet.id for subnet in subnets_by_az]).should.equal(set([subnetB1.id]))
+
     # Unsupported filter
     conn.get_all_subnets.when.called_with(filters={'not-implemented-filter': 'foobar'}).should.throw(NotImplementedError)