Rework configured_sub_zones and add tests specifically for it

octodns/manager.py

@@ -233,22 +233,27 @@ class Manager(object):
     def zone_tree(self):
         if self._zone_tree is None:
             zone_tree = {}
-            # Sort so we iterate on the deepest nodes first, ensuring if a parent
-            # zone exists it will be seen after the subzone, thus we can easily
-            # reparent children to their parent zone from the tree root.
-            for name in sorted(
-                self.config['zones'].keys(), key=lambda s: 0 - s.count('.')
-            ):
-                # Trim the trailing dot from FQDN
-                name = name[:-1]
-                this = {}
-                for sz in [k for k in zone_tree.keys() if k.endswith(name)]:
-                    # Found a zone in tree root that is our child, slice the
-                    # name and move its tree under ours.
-                    this[sz[: -(len(name) + 1)]] = zone_tree.pop(sz)
-                # Add to tree root where it will be reparented as we iterate up
-                # the tree.
-                zone_tree[name] = this
+
+            # Get a list of all of our zone names
+            zones = list(self.config['zones'].keys())
+            # Sort them from shortest to longest so that parents will always
+            # come before their subzones
+            zones.sort(key=lambda z: len(z))
+            # Until we're done processing zones
+            while zones:
+                # Grab the one we'lre going to work on now
+                zone = zones.pop(0)
+                trimmer = len(zone) + 1
+                subs = set()
+                # look at all the zone names that come after it
+                for candidate in zones:
+                    # If they end with this zone's name them they're a sub
+                    if candidate.endswith(zone):
+                        # We want subs to exclude the zone portion
+                        subs.add(candidate[:-trimmer])
+
+                zone_tree[zone] = subs
+
             self._zone_tree = zone_tree
 
         return self._zone_tree
@@ -323,23 +328,7 @@ class Manager(object):
         return kwargs
 
     def configured_sub_zones(self, zone_name):
-        name = zone_name[:-1]
-        where = self.zone_tree
-        while True:
-            # Find parent if it exists
-            parent = next((k for k in where if name.endswith(k)), None)
-            if not parent:
-                # The zone_name in the tree has been reached, stop searching.
-                break
-            # Move down the tree and slice name to get the remainder for the
-            # next round of the search.
-            where = where[parent]
-            name = name[: -(len(parent) + 1)]
-        # `where` is now pointing at the dictionary of children for zone_name
-        # in the tree
-        sub_zone_names = where.keys()
-        self.log.debug('configured_sub_zones: subs=%s', sub_zone_names)
-        return set(sub_zone_names)
+        return self.zone_tree.get(zone_name, set())
 
     def _populate_and_plan(
         self,
@@ -717,6 +706,7 @@ class Manager(object):
                 clz = SplitYamlProvider
             target = clz('dump', output_dir)
 
+        # TODO: use get_zone???
         zone = Zone(zone, self.configured_sub_zones(zone))
         for source in sources:
             source.populate(zone, lenient=lenient)

tests/test_octodns_manager.py

@@ -715,6 +715,7 @@ class TestManager(TestCase):
     def test_subzone_handling(self):
         manager = Manager(get_config_filename('simple.yaml'))
 
+        # tree with multiple branches, one that skips
         manager.config['zones'] = {
             'unit.tests.': {},
             'sub.unit.tests.': {},
@@ -723,23 +724,80 @@ class TestManager(TestCase):
         }
 
         self.assertEqual(
-            {'unit.tests': {'skipped.alevel': {}, 'sub': {'another': {}}}},
+            {
+                'unit.tests.': {'sub', 'another.sub', 'skipped.alevel'},
+                'sub.unit.tests.': {'another'},
+                'another.sub.unit.tests.': set(),
+                'skipped.alevel.unit.tests.': set(),
+            },
             manager.zone_tree,
         )
         self.assertEqual(
-            {'sub', 'skipped.alevel'},
+            {'another.sub', 'sub', 'skipped.alevel'},
             manager.configured_sub_zones('unit.tests.'),
         )
         self.assertEqual(
             {'another'}, manager.configured_sub_zones('sub.unit.tests.')
         )
         self.assertEqual(
-            set(), manager.configured_sub_zones('another.unit.tests.')
+            set(), manager.configured_sub_zones('another.sub.unit.tests.')
         )
         self.assertEqual(
             set(), manager.configured_sub_zones('skipped.alevel.unit.tests.')
         )
 
+        # two parallel trees, make sure they don't interfere
+        manager.config['zones'] = {
+            'unit.tests.': {},
+            'unit2.tests.': {},
+            'sub.unit.tests.': {},
+            'sub.unit2.tests.': {},
+            'another.sub.unit.tests.': {},
+            'another.sub.unit2.tests.': {},
+            'skipped.alevel.unit.tests.': {},
+            'skipped.alevel.unit2.tests.': {},
+        }
+        manager._zone_tree = None
+        self.assertEqual(
+            {
+                'unit.tests.': {'sub', 'another.sub', 'skipped.alevel'},
+                'sub.unit.tests.': {'another'},
+                'another.sub.unit.tests.': set(),
+                'skipped.alevel.unit.tests.': set(),
+                'unit2.tests.': {'sub', 'another.sub', 'skipped.alevel'},
+                'sub.unit2.tests.': {'another'},
+                'another.sub.unit2.tests.': set(),
+                'skipped.alevel.unit2.tests.': set(),
+            },
+            manager.zone_tree,
+        )
+        self.assertEqual(
+            {'another.sub', 'sub', 'skipped.alevel'},
+            manager.configured_sub_zones('unit.tests.'),
+        )
+        self.assertEqual(
+            {'another'}, manager.configured_sub_zones('sub.unit.tests.')
+        )
+        self.assertEqual(
+            set(), manager.configured_sub_zones('another.sub.unit.tests.')
+        )
+        self.assertEqual(
+            set(), manager.configured_sub_zones('skipped.alevel.unit.tests.')
+        )
+        self.assertEqual(
+            {'another.sub', 'sub', 'skipped.alevel'},
+            manager.configured_sub_zones('unit2.tests.'),
+        )
+        self.assertEqual(
+            {'another'}, manager.configured_sub_zones('sub.unit2.tests.')
+        )
+        self.assertEqual(
+            set(), manager.configured_sub_zones('another.sub.unit2.tests.')
+        )
+        self.assertEqual(
+            set(), manager.configured_sub_zones('skipped.alevel.unit2.tests.')
+        )
+
 
 class TestMainThreadExecutor(TestCase):
     def test_success(self):