Updates test and set up.

plugins/inventory/gql_inventory.py

@@ -273,12 +273,12 @@ def add_variable(self, host: str, var: str, var_type: str):
         """
         self.inventory.set_variable(host, var_type, var)
 
-    def add_ip_address(self, default_ip_version, device):
+    def add_ip_address(self, device, default_ip_version="ipv4"):
         """Add primary IP address to host."""
         # Check to see what the primary IP host addition is, first case is IPv4, which is the default
         order_of_preference = ["primary_ip4"]
 
-        # if default_ip_version is IPv4, prepend, else postpend
+        # if default_ip_version is IPv4, prepend, else add to the end
         if default_ip_version.lower() == "ipv6":
             order_of_preference.insert(0, "primary_ip6")
         else:
@@ -424,7 +424,7 @@ def main(self):
         for device in json_data["data"].get("devices", []) + json_data["data"].get("virtual_machines", []):
             hostname = device["name"]
             self.inventory.add_host(host=hostname)
-            self.add_ip_address(self.default_ip_version, device)
+            self.add_ip_address(device, self.default_ip_version)
             self.add_ansible_platform(device)
             self.populate_variables(device)
             self.create_groups(device)

tests/unit/inventory/test_graphql.py

@@ -140,7 +140,34 @@ def test_group_by_empty_string(inventory_fixture, device_data):
 def test_add_ipv4(inventory_fixture, device_data):
     inventory_fixture.group_by = ["location"]
     inventory_fixture.create_groups(device_data)
-    inventory_fixture.add_ipv4_address(device_data)
+    inventory_fixture.add_ip_address(device_data, default_ip_version="ipv4")
+    mydevice_host = inventory_fixture.inventory.get_host("mydevice")
+    assert mydevice_host.vars.get("ansible_host") == "10.10.10.10"
+
+
+def test_add_ipv6(inventory_fixture, device_data):
+    inventory_fixture.group_by = ["location"]
+    inventory_fixture.create_groups(device_data)
+    inventory_fixture.add_ip_address(device_data, default_ip_version="ipv6")
+    mydevice_host = inventory_fixture.inventory.get_host("mydevice")
+    assert mydevice_host.vars.get("ansible_host") == "2001:db8::1"
+
+
+def test_add_ip_address_no_default(inventory_fixture, device_data):
+    inventory_fixture.group_by = ["location"]
+    inventory_fixture.create_groups(device_data)
+    inventory_fixture.add_ip_address(device_data)
+    mydevice_host = inventory_fixture.inventory.get_host("mydevice")
+    assert mydevice_host.vars.get("ansible_host") == "10.10.10.10"
+
+
+def test_add_ip_address_no_ipv6(inventory_fixture, device_data):
+    inventory_fixture.group_by = ["location"]
+
+    # Set the primary_ip6 to None as it would be if there was no ipv6 address assigned
+    device_data["primary_ip6"]["host"] = None
+    inventory_fixture.create_groups(device_data)
+    inventory_fixture.add_ip_address(device_data, default_ip_version="ipv6")
     mydevice_host = inventory_fixture.inventory.get_host("mydevice")
     assert mydevice_host.vars.get("ansible_host") == "10.10.10.10"