handle IPv6 in test extra_df generation

asmap · Nov 25, 2024 · 3113502 · 3113502
1 parent bcc25eb
commit 3113502
Showing 1 changed file with 7 additions and 2 deletions.
diff --git a/tests/merge_base_class_test.py b/tests/merge_base_class_test.py
@@ -7,8 +7,13 @@ def _df_from_network(network, asn=123):
     '''
     Create a one-row dataframe that holds the extra file rows in the expected format for contains_row().
     '''
-    root_net = int(network.split(".", maxsplit=1)[0])
-    network_int = int(ipaddress.ip_network(network).network_address)
+    ipn = ipaddress.ip_network(network)
+    v = ipn.version
+    if v == 4:
+        root_net = int(str(pfx).split(".", maxsplit=1)[0])
+    else:
+        root_net = int(str(pfx).split(":", maxsplit=1)[0])
+    network_int = int(ipn.network_address)
     df_extra = pd.DataFrame(
         data={"INETS": network_int, "ASNS": asn, "PFXS": network, "PFXS_LEADING": root_net},
         index=[0],