add tests for util.py and update format_pfx logic

kartograf/util.py

@@ -138,12 +138,17 @@ def format_pfx(pfx):
     """
     try:
         if "/" in pfx:
+            pattern = r"^0+"
+            match = re.search(pattern, pfx)
+            if match:
+                pfx = re.sub(pattern, "", pfx)
             formatted_pfx = str(ipaddress.ip_network(pfx))
             return f"{formatted_pfx}"
         return str(ipaddress.ip_address(pfx))
     except ValueError:
         return pfx
 
+
 def get_root_network(pfx):
     """
     Extract the top-level network from an IPv4 or IPv6 address.

tests/merge_base_class_test.py

@@ -14,7 +14,6 @@ def _df_from_networks(networks, asn=123):
     for network in networks:
         ipn = ipaddress.ip_network(network)
         root_net = get_root_network(network)
-        print(root_net)
         network_int = int(ipn.network_address)
         df.loc[len(df)] = [network_int, asn, str(ipn), root_net]
     return df

tests/test_util.py

@@ -0,0 +1,54 @@
+import pytest
+
+from kartograf.util import format_pfx, get_root_network
+
+def test_valid_ipv4_network():
+    pfx = "192.144.11.0/21"
+    assert format_pfx(pfx) == pfx
+
+
+def test_valid_ipv4_addr():
+    pfx = "192.144.11.0"
+    assert format_pfx(pfx) == pfx
+
+
+def test_valid_ipv6_network():
+    pfx = "2001:db8::/64"
+    assert format_pfx(pfx) == pfx
+
+
+def test_valid_ipv6_addr():
+    pfx = "2001:db8::1"
+    assert format_pfx(pfx) == pfx
+
+
+def test_ipv4_prefix_with_leading_zeros():
+    pfx = "010.10.00.00/16"
+    expected_output = "10.10.00.00/16"
+    assert format_pfx(pfx) == expected_output
+
+
+def test_ipv6_prefix_with_leading_zeros():
+    pfx = "001:db8::0/24"
+    expected_output = "1:db8::0/24"
+    assert format_pfx(pfx) == expected_output
+
+
+def test_invalid_ip_network():
+    pfx = "192.1/asdf"
+    assert format_pfx(pfx) == pfx
+
+
+def test_invalid_input():
+    pfx = "no.slash"
+    assert format_pfx(pfx) == pfx
+
+
+def test_get_root_network():
+    ipv4 = "192.144.11.0/24"
+    assert get_root_network(ipv4) == 192
+    ipv6 = "2001:db8::/64"
+    assert get_root_network(ipv6) == int("2001", 16)
+    invalid = "not.a.network"
+    with pytest.raises(ValueError):
+        get_root_network(invalid)