From a8505d66f13421f1c67242f4ecbf1fbeae85aef5 Mon Sep 17 00:00:00 2001
From: Robert Reichel <rreichel3@github.com>
Date: Fri, 29 Jan 2021 15:11:27 -0500
Subject: [PATCH] Improve checking and creating Azure DNS Zones

This change improves the process for checking for AzureDNS zones by using the known set and not relying upon custom error handling.

Since the provider already fetches the zones, octodns doesn't need to make a second call to check for the existence of the zone - _populate_zones already does that for us.
---
 octodns/provider/azuredns.py            | 36 ++++++++++---------------
 tests/test_octodns_provider_azuredns.py | 30 ++++++++++++++-------
 2 files changed, 34 insertions(+), 32 deletions(-)

diff --git a/octodns/provider/azuredns.py b/octodns/provider/azuredns.py
index a1ef6fe..0224f06 100644
--- a/octodns/provider/azuredns.py
+++ b/octodns/provider/azuredns.py
@@ -341,7 +341,8 @@ class AzureProvider(BaseProvider):
         self.log.debug('azure_zones: loading')
         list_zones = self._dns_client.zones.list_by_resource_group
         for zone in list_zones(self._resource_group):
-            self._azure_zones.add(zone.name)
+            if zone.name:
+                self._azure_zones.add(zone.name.rstrip('.'))
 
     def _check_zone(self, name, create=False):
         '''Checks whether a zone specified in a source exist in Azure server.
@@ -356,29 +357,20 @@ class AzureProvider(BaseProvider):
 
             :type return: str or None
         '''
-        self.log.debug('_check_zone: name=%s', name)
-        try:
-            if name in self._azure_zones:
-                return name
-            self._dns_client.zones.get(self._resource_group, name)
+        self.log.debug('_check_zone: name=%s create=%s', name, create)
+        # Check if the zone already exists in our set
+        if name in self._azure_zones:
+            return name
+        # If not, and its time to create, lets do it.
+        if create:
+            self.log.debug('_check_zone:no matching zone; creating %s', name)
+            create_zone = self._dns_client.zones.create_or_update
+            create_zone(self._resource_group, name, Zone(location='global'))
             self._azure_zones.add(name)
             return name
-        except CloudError as err:
-            msg = 'The Resource \'Microsoft.Network/dnszones/{}\''.format(name)
-            msg += ' under resource group \'{}\''.format(self._resource_group)
-            msg += ' was not found.'
-            if msg == err.message:
-                # Then the only error is that the zone doesn't currently exist
-                if create:
-                    self.log.debug('_check_zone:no matching zone; creating %s',
-                                   name)
-                    create_zone = self._dns_client.zones.create_or_update
-                    create_zone(self._resource_group, name,
-                                Zone(location='global'))
-                    return name
-                else:
-                    return
-            raise
+        else:
+            # Else return nothing (aka false)
+            return
 
     def populate(self, zone, target=False, lenient=False):
         '''Required function of manager.py to collect records from zone.
diff --git a/tests/test_octodns_provider_azuredns.py b/tests/test_octodns_provider_azuredns.py
index 2b2c1e7..3c93a27 100644
--- a/tests/test_octodns_provider_azuredns.py
+++ b/tests/test_octodns_provider_azuredns.py
@@ -498,32 +498,42 @@ class TestAzureDnsProvider(TestCase):
         record_list = provider._dns_client.record_sets.list_by_dns_zone
         record_list.return_value = rs
 
+        zone_list = provider._dns_client.zones.list_by_resource_group
+        zone_list.return_value = [zone]
+
         exists = provider.populate(zone)
-        self.assertTrue(exists)
+        
         self.assertEquals(len(zone.records), 17)
+        self.assertTrue(exists)
 
     def test_populate_zone(self):
         provider = self._get_provider()
 
         zone_list = provider._dns_client.zones.list_by_resource_group
-        zone_list.return_value = [AzureZone(location='global'),
-                                  AzureZone(location='global')]
+        zone_1 = AzureZone(location='global')
+        # This is far from ideal but the zone constructor doesn't let me set it on creation
+        zone_1.name = "zone-1"
+        zone_2 = AzureZone(location='global')
+        # This is far from ideal but the zone constructor doesn't let me set it on creation
+        zone_2.name = "zone-2"
+        zone_list.return_value = [zone_1,
+                                  zone_2,
+                                  zone_1]
 
         provider._populate_zones()
 
-        self.assertEquals(len(provider._azure_zones), 1)
+        # This should be returning two zones since two zones are the same
+        self.assertEquals(len(provider._azure_zones), 2)
 
     def test_bad_zone_response(self):
         provider = self._get_provider()
 
         _get = provider._dns_client.zones.get
         _get.side_effect = CloudError(Mock(status=404), 'Azure Error')
-        trip = False
-        try:
-            provider._check_zone('unit.test', create=False)
-        except CloudError:
-            trip = True
-        self.assertEquals(trip, True)
+        self.assertEquals(
+            provider._check_zone('unit.test', create=False), 
+            None
+        )
 
     def test_apply(self):
         provider = self._get_provider()