From 0ea10ca797e7634c282c939038fce52cf44cee6a Mon Sep 17 00:00:00 2001
From: Ross McFarland <rwmcfa1@gmail.com>
Date: Tue, 18 Jun 2024 09:14:18 -0700
Subject: [PATCH] Add ICMP & UDP healthcheck protocols

---
 CHANGELOG.md                         |  1 +
 octodns/record/base.py               | 10 +++-
 tests/test_octodns_record_dynamic.py | 77 ++++++++++++++++++++++++++++
 3 files changed, 86 insertions(+), 2 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 3e959ce..892d31c 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -1,5 +1,6 @@
 ## v1.?.? - 2024-??-?? - ???
 
+* ICMP & UDP healthcheck protocol support added
 * Improved handling of present, but empty/None config file values.
 * Add PlanJson plan_output support
 * Include `record_type` in Change data
diff --git a/octodns/record/base.py b/octodns/record/base.py
index 456473c..2cc27f9 100644
--- a/octodns/record/base.py
+++ b/octodns/record/base.py
@@ -113,7 +113,9 @@ class Record(EqualityTupleMixin):
             if data['octodns']['healthcheck']['protocol'] not in (
                 'HTTP',
                 'HTTPS',
+                'ICMP',
                 'TCP',
+                'UDP',
             ):
                 reasons.append('invalid healthcheck protocol')
         except KeyError:
@@ -224,14 +226,16 @@ class Record(EqualityTupleMixin):
 
     def healthcheck_host(self, value=None):
         healthcheck = self.octodns.get('healthcheck', {})
-        if healthcheck.get('protocol', None) == 'TCP':
+        protocol = self.healthcheck_protocol
+        if protocol not in ('HTTP', 'HTTPS'):
             return None
         return healthcheck.get('host', self.fqdn[:-1]) or value
 
     @property
     def healthcheck_path(self):
         healthcheck = self.octodns.get('healthcheck', {})
-        if healthcheck.get('protocol', None) == 'TCP':
+        protocol = self.healthcheck_protocol
+        if protocol not in ('HTTP', 'HTTPS'):
             return None
         try:
             return healthcheck['path']
@@ -247,6 +251,8 @@ class Record(EqualityTupleMixin):
 
     @property
     def healthcheck_port(self):
+        if self.healthcheck_protocol == 'ICMP':
+            return None
         try:
             return int(self.octodns['healthcheck']['port'])
         except KeyError:
diff --git a/tests/test_octodns_record_dynamic.py b/tests/test_octodns_record_dynamic.py
index 8d7b402..471957b 100644
--- a/tests/test_octodns_record_dynamic.py
+++ b/tests/test_octodns_record_dynamic.py
@@ -82,6 +82,7 @@ class TestRecordDynamic(TestCase):
         )
         self.assertEqual('1.2.3.4', new.healthcheck_host(value="1.2.3.4"))
 
+        # defaults
         new = Record.new(
             self.zone, 'a', {'ttl': 44, 'type': 'A', 'value': '1.2.3.4'}
         )
@@ -90,6 +91,44 @@ class TestRecordDynamic(TestCase):
         self.assertEqual('HTTPS', new.healthcheck_protocol)
         self.assertEqual(443, new.healthcheck_port)
 
+    def test_healthcheck_icmp(self):
+        new = Record.new(
+            self.zone,
+            'a',
+            {
+                'ttl': 44,
+                'type': 'A',
+                'value': '1.2.3.4',
+                'octodns': {
+                    'healthcheck': {
+                        'path': '/ignored',
+                        'host': 'completely.ignored',
+                        'protocol': 'ICMP',
+                        'port': -99,
+                    }
+                },
+            },
+        )
+        self.assertIsNone(new.healthcheck_path)
+        self.assertIsNone(new.healthcheck_host())
+        self.assertEqual('ICMP', new.healthcheck_protocol)
+        self.assertIsNone(new.healthcheck_port)
+
+        new = Record.new(
+            self.zone,
+            'a',
+            {
+                'ttl': 44,
+                'type': 'A',
+                'value': '1.2.3.4',
+                'octodns': {'healthcheck': {'protocol': 'ICMP'}},
+            },
+        )
+        self.assertIsNone(new.healthcheck_path)
+        self.assertIsNone(new.healthcheck_host())
+        self.assertEqual('ICMP', new.healthcheck_protocol)
+        self.assertIsNone(new.healthcheck_port)
+
     def test_healthcheck_tcp(self):
         new = Record.new(
             self.zone,
@@ -128,6 +167,44 @@ class TestRecordDynamic(TestCase):
         self.assertEqual('TCP', new.healthcheck_protocol)
         self.assertEqual(443, new.healthcheck_port)
 
+    def test_healthcheck_udp(self):
+        new = Record.new(
+            self.zone,
+            'a',
+            {
+                'ttl': 44,
+                'type': 'A',
+                'value': '1.2.3.4',
+                'octodns': {
+                    'healthcheck': {
+                        'path': '/ignored',
+                        'host': 'completely.ignored',
+                        'protocol': 'UDP',
+                        'port': 8081,
+                    }
+                },
+            },
+        )
+        self.assertIsNone(new.healthcheck_path)
+        self.assertIsNone(new.healthcheck_host())
+        self.assertEqual('UDP', new.healthcheck_protocol)
+        self.assertEqual(8081, new.healthcheck_port)
+
+        new = Record.new(
+            self.zone,
+            'a',
+            {
+                'ttl': 44,
+                'type': 'A',
+                'value': '1.2.3.4',
+                'octodns': {'healthcheck': {'protocol': 'UDP'}},
+            },
+        )
+        self.assertIsNone(new.healthcheck_path)
+        self.assertIsNone(new.healthcheck_host())
+        self.assertEqual('UDP', new.healthcheck_protocol)
+        self.assertEqual(443, new.healthcheck_port)
+
     def test_simple_a_weighted(self):
         a_data = {
             'dynamic': {