From 725189af49791686542c5f580b2c30660ea653bb Mon Sep 17 00:00:00 2001
From: Viranch Mehta <email@viranch.me>
Date: Mon, 14 Jun 2021 13:58:20 -0700
Subject: [PATCH] Handle re-used pools in Azure DNS dynamic records

---
 octodns/provider/azuredns.py            | 28 ++++++++++-
 tests/test_octodns_provider_azuredns.py | 67 +++++++++++++++++++++++++
 2 files changed, 93 insertions(+), 2 deletions(-)

diff --git a/octodns/provider/azuredns.py b/octodns/provider/azuredns.py
index 0fd3dae..93478b8 100644
--- a/octodns/provider/azuredns.py
+++ b/octodns/provider/azuredns.py
@@ -717,6 +717,8 @@ class AzureProvider(BaseProvider):
                         country, province = code.split('-', 1)
                         country = GeoCodes.country_to_code(country)
                         geos.append('{}-{}'.format(country, province))
+                    elif code == 'WORLD':
+                        geos.append(code)
                     else:
                         # country
                         geos.append(GeoCodes.country_to_code(code))
@@ -788,6 +790,13 @@ class AzureProvider(BaseProvider):
 
             rules.append(rule)
 
+        # add separate rule for re-used world pool
+        for rule in list(rules):
+            geos = rule.get('geos', [])
+            if len(geos) > 1 and 'WORLD' in geos:
+                geos.remove('WORLD')
+                rules.append({'pool': rule['pool']})
+
         # Order and convert to a list
         default = sorted(default)
 
@@ -904,14 +913,29 @@ class AzureProvider(BaseProvider):
     def _generate_traffic_managers(self, record):
         traffic_managers = []
         pools = record.dynamic.pools
+        rules = record.dynamic.rules
 
         default = record.value[:-1]
         profile = self._generate_tm_profile
 
+        # a pool can be re-used only with a world pool, record the pool
+        # to later consolidate it with a geo pool if one exists since we
+        # can't have multiple endpoints with the same target in ATM
+        world_pool = None
+        for rule in rules:
+            if not rule.data.get('geos', []):
+                world_pool = rule.data['pool']
+        world_seen = False
+
         geo_endpoints = []
         pool_profiles = {}
 
         for rule in record.dynamic.rules:
+            pool_name = rule.data['pool']
+            if pool_name == world_pool and world_seen:
+                # this world pool is already mentioned in another geo rule
+                continue
+
             # Prepare the list of Traffic manager geos
             rule_geos = rule.data.get('geos', [])
             geos = []
@@ -933,10 +957,10 @@ class AzureProvider(BaseProvider):
                         geo = 'AP'
 
                     geos.append('GEO-{}'.format(geo))
-            if not geos:
+            if not geos or pool_name == world_pool:
                 geos.append('WORLD')
+                world_seen = True
 
-            pool_name = rule.data['pool']
             rule_endpoints = []
             priority = 1
             default_seen = False
diff --git a/tests/test_octodns_provider_azuredns.py b/tests/test_octodns_provider_azuredns.py
index 81a3084..56a783d 100644
--- a/tests/test_octodns_provider_azuredns.py
+++ b/tests/test_octodns_provider_azuredns.py
@@ -1506,6 +1506,73 @@ class TestAzureDnsProvider(TestCase):
                 self.assertNotIn(tm.name, seen)
                 seen.add(tm.name)
 
+    def test_dynamic_reused_pool(self):
+        # test that traffic managers are generated as expected
+        provider = self._get_provider()
+        nested = 'Microsoft.Network/trafficManagerProfiles/nestedEndpoints'
+
+        record = Record.new(zone, 'foo', data={
+            'type': 'CNAME',
+            'ttl': 60,
+            'value': 'default.unit.tests.',
+            'dynamic': {
+                'pools': {
+                    'iad': {
+                        'values': [
+                            {'value': 'iad.unit.tests.'},
+                        ],
+                        'fallback': 'lhr',
+                    },
+                    'lhr': {
+                        'values': [
+                            {'value': 'lhr.unit.tests.'},
+                        ],
+                    },
+                },
+                'rules': [
+                    {'geos': ['EU'], 'pool': 'iad'},
+                    {'geos': ['EU-GB'], 'pool': 'lhr'},
+                    {'pool': 'lhr'},
+                ],
+            }
+        })
+        profiles = provider._generate_traffic_managers(record)
+
+        self.assertEqual(len(profiles), 3)
+        self.assertTrue(_profile_is_match(profiles[-1], Profile(
+            name='foo--unit--tests',
+            traffic_routing_method='Geographic',
+            dns_config=DnsConfig(
+                relative_name='foo--unit--tests', ttl=record.ttl),
+            monitor_config=_get_monitor(record),
+            endpoints=[
+                Endpoint(
+                    name='rule-iad',
+                    type=nested,
+                    target_resource_id=profiles[0].id,
+                    geo_mapping=['GEO-EU'],
+                ),
+                Endpoint(
+                    name='rule-lhr',
+                    type=nested,
+                    target_resource_id=profiles[1].id,
+                    geo_mapping=['GB', 'WORLD'],
+                ),
+            ],
+        )))
+
+        # test that same record gets populated back from traffic managers
+        tm_list = provider._tm_client.profiles.list_by_resource_group
+        tm_list.return_value = profiles
+        azrecord = RecordSet(
+            ttl=60,
+            target_resource=SubResource(id=profiles[-1].id),
+        )
+        azrecord.name = record.name or '@'
+        azrecord.type = 'Microsoft.Network/dnszones/{}'.format(record._type)
+        record2 = provider._populate_record(zone, azrecord)
+        self.assertEqual(record2.dynamic._data(), record.dynamic._data())
+
     def test_sync_traffic_managers(self):
         provider, zone, record = self._get_dynamic_package()
         provider._populate_traffic_managers()