From 40820f351edc4de75da9b5b329cefb8e096b0a67 Mon Sep 17 00:00:00 2001
From: Ross McFarland <rwmcfa1@gmail.com>
Date: Wed, 16 Feb 2022 13:41:35 -0800
Subject: [PATCH] Implement and test Zone.root_ns helper property

---
 octodns/zone.py            | 14 +++++++++++
 tests/test_octodns_zone.py | 50 +++++++++++++++++++++++++++++++++++++-
 2 files changed, 63 insertions(+), 1 deletion(-)

diff --git a/octodns/zone.py b/octodns/zone.py
index c1929cf..4cd5e91 100644
--- a/octodns/zone.py
+++ b/octodns/zone.py
@@ -36,6 +36,7 @@ class Zone(object):
         # We're grouping by node, it allows us to efficiently search for
         # duplicates and detect when CNAMEs co-exist with other records
         self._records = defaultdict(set)
+        self._root_ns = None
         # optional leading . to match empty hostname
         # optional trailing . b/c some sources don't have it on their fqdn
         self._name_re = re.compile(fr'\.?{name}?$')
@@ -53,6 +54,12 @@ class Zone(object):
             return self._origin.records
         return set([r for _, node in self._records.items() for r in node])
 
+    @property
+    def root_ns(self):
+        if self._origin:
+            return self._origin.root_ns
+        return self._root_ns
+
     def hostname_from_fqdn(self, fqdn):
         return self._name_re.sub('', fqdn)
 
@@ -91,11 +98,18 @@ class Zone(object):
                                        f'{record.fqdn} cannot coexist with '
                                        'other records')
 
+        if record._type == 'NS' and record.name == '':
+            self._root_ns = record
+
         node.add(record)
 
     def remove_record(self, record):
         if self._origin:
             self.hydrate()
+
+        if record._type == 'NS' and record.name == '':
+            self._root_ns = None
+
         self._records[record.name].discard(record)
 
     # TODO: delete this
diff --git a/tests/test_octodns_zone.py b/tests/test_octodns_zone.py
index e032369..9385984 100644
--- a/tests/test_octodns_zone.py
+++ b/tests/test_octodns_zone.py
@@ -7,7 +7,8 @@ from __future__ import absolute_import, division, print_function, \
 
 from unittest import TestCase
 
-from octodns.record import ARecord, AaaaRecord, Create, Delete, Record, Update
+from octodns.record import ARecord, AaaaRecord, Create, Delete, NsRecord, \
+    Record, Update
 from octodns.zone import DuplicateRecordException, InvalidNodeException, \
     SubzoneRecordException, Zone
 
@@ -410,3 +411,50 @@ class TestZone(TestCase):
         self.assertTrue(copy.hydrate())
         # Doesn't the second
         self.assertFalse(copy.hydrate())
+
+    def test_root_ns(self):
+        zone = Zone('unit.tests.', [])
+
+        a = ARecord(zone, 'a', {'ttl': 42, 'value': '1.1.1.1'})
+        zone.add_record(a)
+        # No root NS yet
+        self.assertFalse(zone.root_ns)
+
+        non_root_ns = NsRecord(zone, 'sub', {'ttl': 42, 'values': (
+            'ns1.unit.tests.',
+            'ns2.unit.tests.',
+        )})
+        zone.add_record(non_root_ns)
+        # No root NS yet b/c this was a sub
+        self.assertFalse(zone.root_ns)
+
+        root_ns = NsRecord(zone, '', {'ttl': 42, 'values': (
+            'ns3.unit.tests.',
+            'ns4.unit.tests.',
+        )})
+        zone.add_record(root_ns)
+        # Now we have a root NS
+        self.assertEqual(root_ns, zone.root_ns)
+
+        # make a copy, it has a root_ns
+        copy = zone.copy()
+        self.assertEqual(root_ns, copy.root_ns)
+
+        # remove the root NS from it and we don't
+        copy.remove_record(root_ns)
+        self.assertFalse(copy.root_ns)
+
+        # original still does though
+        self.assertEqual(root_ns, zone.root_ns)
+
+        # remove the A, still has root NS
+        zone.remove_record(a)
+        self.assertEqual(root_ns, zone.root_ns)
+
+        # remove the sub NS, still has root NS
+        zone.remove_record(non_root_ns)
+        self.assertEqual(root_ns, zone.root_ns)
+
+        # finally remove the root NS, no more
+        zone.remove_record(root_ns)
+        self.assertFalse(zone.root_ns)