sources/azure: move get_ip_from_lease_value out of shim (#1324)

Just a minor refactoring to cleanup the shim.

Update tests to use pytest parametrization.

Signed-off-by: Chris Patterson <cpatterson@microsoft.com>

3 files changed, 33 insertions, 50 deletions

diff --git a/cloudinit/sources/DataSourceAzure.py b/cloudinit/sources/DataSourceAzure.py
index ef5e0725..85f0211b 100755
--- a/cloudinit/sources/DataSourceAzure.py
+++ b/cloudinit/sources/DataSourceAzure.py
@@ -34,12 +34,12 @@ from cloudinit.sources.helpers import netlink
 from cloudinit.sources.helpers.azure import (
     DEFAULT_REPORT_FAILURE_USER_VISIBLE_MESSAGE,
     DEFAULT_WIRESERVER_ENDPOINT,
-    WALinuxAgentShim,
     azure_ds_reporter,
     azure_ds_telemetry_reporter,
     build_minimal_ovf,
     dhcp_log_cb,
     get_boot_telemetry,
+    get_ip_from_lease_value,
     get_metadata_from_fabric,
     get_system_info,
     is_byte_swapped,
@@ -429,10 +429,8 @@ class DataSourceAzure(sources.DataSource):
 
                 # Update wireserver IP from DHCP options.
                 if "unknown-245" in lease:
-                    self._wireserver_endpoint = (
-                        WALinuxAgentShim.get_ip_from_lease_value(
-                            lease["unknown-245"]
-                        )
+                    self._wireserver_endpoint = get_ip_from_lease_value(
+                        lease["unknown-245"]
                     )
 
     @azure_ds_telemetry_reporter
diff --git a/cloudinit/sources/helpers/azure.py b/cloudinit/sources/helpers/azure.py
index a6edc816..196ee031 100755
--- a/cloudinit/sources/helpers/azure.py
+++ b/cloudinit/sources/helpers/azure.py
@@ -315,6 +315,20 @@ def cd(newdir):
         os.chdir(prevdir)
 
 
+def get_ip_from_lease_value(fallback_lease_value):
+    unescaped_value = fallback_lease_value.replace("\\", "")
+    if len(unescaped_value) > 4:
+        hex_string = ""
+        for hex_pair in unescaped_value.split(":"):
+            if len(hex_pair) == 1:
+                hex_pair = "0" + hex_pair
+            hex_string += hex_pair
+        packed_bytes = struct.pack(">L", int(hex_string.replace(":", ""), 16))
+    else:
+        packed_bytes = unescaped_value.encode("utf-8")
+    return socket.inet_ntoa(packed_bytes)
+
+
 @azure_ds_telemetry_reporter
 def http_with_retries(url, **kwargs) -> url_helper.UrlResponse:
     """Wrapper around url_helper.readurl() with custom telemetry logging
@@ -799,22 +813,6 @@ class WALinuxAgentShim:
         if self.openssl_manager is not None:
             self.openssl_manager.clean_up()
 
-    @staticmethod
-    def get_ip_from_lease_value(fallback_lease_value):
-        unescaped_value = fallback_lease_value.replace("\\", "")
-        if len(unescaped_value) > 4:
-            hex_string = ""
-            for hex_pair in unescaped_value.split(":"):
-                if len(hex_pair) == 1:
-                    hex_pair = "0" + hex_pair
-                hex_string += hex_pair
-            packed_bytes = struct.pack(
-                ">L", int(hex_string.replace(":", ""), 16)
-            )
-        else:
-            packed_bytes = unescaped_value.encode("utf-8")
-        return socket.inet_ntoa(packed_bytes)
-
     @azure_ds_telemetry_reporter
     def eject_iso(self, iso_dev) -> None:
         try:
diff --git a/tests/unittests/sources/test_azure_helper.py b/tests/unittests/sources/test_azure_helper.py
index bcdccd83..326b41cc 100644
--- a/tests/unittests/sources/test_azure_helper.py
+++ b/tests/unittests/sources/test_azure_helper.py
@@ -8,6 +8,8 @@ from textwrap import dedent
 from xml.etree import ElementTree
 from xml.sax.saxutils import escape, unescape
 
+import pytest
+
 from cloudinit.sources.helpers import azure as azure_helper
 from cloudinit.sources.helpers.azure import WALinuxAgentShim as wa_shim
 from cloudinit.util import load_file
@@ -87,35 +89,20 @@ class SentinelException(Exception):
     pass
 
 
-class TestExtractIpAddressFromLeaseValue(CiTestCase):
-    def test_hex_string(self):
-        ip_address, encoded_address = "98.76.54.32", "62:4c:36:20"
-        self.assertEqual(
-            ip_address, wa_shim.get_ip_from_lease_value(encoded_address)
-        )
-
-    def test_hex_string_with_single_character_part(self):
-        ip_address, encoded_address = "4.3.2.1", "4:3:2:1"
-        self.assertEqual(
-            ip_address, wa_shim.get_ip_from_lease_value(encoded_address)
-        )
-
-    def test_packed_string(self):
-        ip_address, encoded_address = "98.76.54.32", "bL6 "
-        self.assertEqual(
-            ip_address, wa_shim.get_ip_from_lease_value(encoded_address)
-        )
-
-    def test_packed_string_with_escaped_quote(self):
-        ip_address, encoded_address = "100.72.34.108", 'dH\\"l'
-        self.assertEqual(
-            ip_address, wa_shim.get_ip_from_lease_value(encoded_address)
-        )
-
-    def test_packed_string_containing_a_colon(self):
-        ip_address, encoded_address = "100.72.58.108", "dH:l"
-        self.assertEqual(
-            ip_address, wa_shim.get_ip_from_lease_value(encoded_address)
+class TestGetIpFromLeaseValue:
+    @pytest.mark.parametrize(
+        "encoded_address,ip_address",
+        [
+            ("62:4c:36:20", "98.76.54.32"),
+            ("4:3:2:1", "4.3.2.1"),
+            ("bL6 ", "98.76.54.32"),
+            ('dH\\"l', "100.72.34.108"),
+            ("dH:l", "100.72.58.108"),
+        ],
+    )
+    def test_get_ip_from_lease_value(self, encoded_address, ip_address):
+        assert (
+            azure_helper.get_ip_from_lease_value(encoded_address) == ip_address
         )