Update filter for firewall modules, formatting

tests/infrastructure/hub-and-spoke-vpns/test_cloud_routers.py

@@ -18,21 +18,22 @@
 def test_hub_custom_routers(plan):
   "Custom routers should match input variables."
   hub_custom_router_tpl = ('google_compute_router.hub-to-spoke-%s-custom[0]')
-  count = range(1, 3)
-  for i in count:
+  for i in range(1, 3):
     name = hub_custom_router_tpl % i
     custom_router = plan.resource_changes[name]
+    adv_ranges = custom_router['change']['after']['bgp'][0]['advertised_ip_ranges']
+    spoke_subnets = plan.variables['spoke_%s_subnets' % (3 - i)]
     assert custom_router['change']['after']['bgp'][0]['advertise_mode'] == 'CUSTOM'
     assert custom_router['change']['after']['bgp'][0]['advertised_groups'] == ['ALL_SUBNETS']
     assert custom_router['change']['after']['bgp'][0]['asn'] == plan.variables['hub_bgp_asn']
-    assert [range['range'] for range in custom_router['change']['after']['bgp'][0]['advertised_ip_ranges']] == [subnet['subnet_ip'] for subnet in plan.variables['spoke_%s_subnets' % (3 - i)]]
+    assert [range['range'] for range in adv_ranges] == [subnet['subnet_ip']
+                                                            for subnet in spoke_subnets]
 
 def test_spoke_routers(plan):
   "Spoke routers should match input variables."
   spoke_router_tpl = ('google_compute_router.spoke-%s')
   spoke_bgp_asn_tpl = ('spoke_%s_bgp_asn')
-  count = range(1, 3)
-  for i in count:
+  for i in range(1, 3):
     spoke_router = plan.resource_changes[spoke_router_tpl % i]
     spoke_bgp_asn = plan.variables[spoke_bgp_asn_tpl % i]
     assert spoke_router['change']['after']['bgp'][0]['advertise_mode'] == 'DEFAULT'

tests/infrastructure/hub-and-spoke-vpns/test_firewall.py

@@ -19,8 +19,7 @@ import pytest
 
 @pytest.fixture(scope='module')
 def firewall_modules(plan):
-  names = ['module.firewall-%s' %
-           name for name in ('hub', 'spoke-1', 'spoke-2')]
+  names = [name for name in plan.modules.keys() if name.startswith('module.firewall-')]
   return dict((name, plan.modules[name]) for name in names)
 
 

tests/infrastructure/hub-and-spoke-vpns/test_outputs.py

@@ -27,8 +27,7 @@ def test_spokes_outputs(plan):
   "Spokes VPC ranges and regions should match input variables."
   spoke_output_tpl = ('spoke-%s')
   spole_subnets_tpl = ('spoke_%s_subnets')
-  count = range(1, 3)
-  for i in count:
+  for i in range(1, 3):
     spoke_output = plan.outputs[spoke_output_tpl % i]
     for spoke_subnet in plan.variables[spole_subnets_tpl % i]:
       assert spoke_output['subnets_ips'][spoke_subnet['subnet_name']] == spoke_subnet['subnet_ip']

tests/infrastructure/hub-and-spoke-vpns/test_vpns.py

@@ -43,7 +43,6 @@ def test_spokes_peer_asn(plan, spoke_to_hub_vpn):
 
 def test_hub_peer_asns(plan, hub_to_spoke_vpn):
   "Test that the hub-to-spoke VPNs mach input variables"
-  count = range(1,3)
   var_names, mods = hub_to_spoke_vpn
   for var_name, mod in zip(var_names, mods):
     hub_bgp_peer = mod.resources['google_compute_router_peer.bgp_peer[0]']