Allow mapping additional hostvars in write-inventory

In OpenDev, we use an inventory file with the public v4 and v6
addresses listed as hostvars.  We would like to use that feature in
testing as well.  In order to do that, we need to mutate the inventory
so that the public_ipv4 address supplied by nodepool is mapped to
the public_v4 address in the inventory.  This additional option allows
the user to configure such mappings.

Change-Id: I48f03cacdf4531c42d33f6e807845d5c2a1da1d5

roles/write-inventory/README.rst

@@ -24,3 +24,18 @@ with the inventory for the job.
    A list of facts about the host to exclude.  By default, all
    variables about a host will be included.  To exclude certain
    variables, list them here.
+
+.. zuul:rolevar:: write_inventory_additional_hostvars
+   :type: dict
+
+   Additional hostvars to include.  This can be used to map
+   information from nodepool into the inventory if used as follows:
+
+   .. code-block:: yaml
+
+      write_inventory_additional_hostvars:
+        public_v4: nodepool.public_ipv4
+        public_v6: nodepool.public_ipv6
+
+   This will map hostvars[hostname]['nodepool']['public_ipv4'] to
+   hostvars[hostname]['public_v4'].

roles/write-inventory/library/test_write_inventory.py

@@ -37,7 +37,6 @@ bionic:
     private_ipv4: 10.210.196.115
     provider: rax-ord
     public_ipv4: 104.130.217.77
-    public_ipv6: 2001:4801:7828:101:be76:4eff:fe10:14eb
     region: ORD
 xenial:
   ansible_connection: ssh
@@ -51,7 +50,6 @@ xenial:
     label: ubuntu-xenial
     private_ipv4: 149.202.170.85
     provider: ovh-gra1
-    public_ipv4: 149.202.170.85
     public_ipv6: 2001:41d0:302:1000::17:a32b
     region: GRA1
 """)
@@ -75,7 +73,7 @@ class TestWriteInventory(testtools.TestCase):
         '''Test passing all variables'''
         dest = self.useFixture(fixtures.TempDir()).path
         dest = os.path.join(dest, 'out.yaml')
-        run(dest, INPUT, GROUPS_INPUT, None, None)
+        run(dest, INPUT, GROUPS_INPUT, None, None, None)
 
         self.assertOutput(dest, {
             'all': {
@@ -108,7 +106,7 @@ class TestWriteInventory(testtools.TestCase):
         '''Test incuding vars'''
         dest = self.useFixture(fixtures.TempDir()).path
         dest = os.path.join(dest, 'out.yaml')
-        run(dest, INPUT, GROUPS_INPUT, ['ansible_host'], None)
+        run(dest, INPUT, GROUPS_INPUT, ['ansible_host'], None, None)
 
         self.assertOutput(dest, {
             'all': {
@@ -135,7 +133,7 @@ class TestWriteInventory(testtools.TestCase):
         '''Test passing all variables'''
         dest = self.useFixture(fixtures.TempDir()).path
         dest = os.path.join(dest, 'out.yaml')
-        run(dest, INPUT, GROUPS_INPUT, None, ['ansible_user'])
+        run(dest, INPUT, GROUPS_INPUT, None, ['ansible_user'], None)
 
         self.assertOutput(dest, {
             'all': {
@@ -161,3 +159,41 @@ class TestWriteInventory(testtools.TestCase):
                 }
             }
         })
+
+    def test_additional(self):
+        '''Test passing additional variables'''
+        dest = self.useFixture(fixtures.TempDir()).path
+        dest = os.path.join(dest, 'out.yaml')
+        run(dest, INPUT, GROUPS_INPUT, None, None,
+            {'public_v4': 'nodepool.public_ipv4',
+             'public_v6': 'nodepool.public_ipv6',
+        })
+
+        self.assertOutput(dest, {
+            'all': {
+                'children': {
+                    'puppet': {
+                        'hosts': {
+                            'bionic': None,
+                            'xenial': None,
+                        },
+                    },
+                },
+                'hosts': {
+                    'bionic': {
+                        "ansible_connection": "ssh",
+                        "ansible_user": "zuul",
+                        "ansible_host": "104.130.217.77",
+                        "ansible_port": 22,
+                        "public_v4": "104.130.217.77",
+                    },
+                    'xenial': {
+                        "ansible_connection": "ssh",
+                        "ansible_user": "zuul",
+                        "ansible_host": "149.202.170.85",
+                        "ansible_port": 22,
+                        "public_v6": "2001:41d0:302:1000::17:a32b",
+                    }
+                }
+            }
+        })

roles/write-inventory/library/write_inventory.py

@@ -28,7 +28,7 @@ VARS = [
 ]
 
 
-def run(dest, hostvars, groups, include, exclude):
+def run(dest, hostvars, groups, include, exclude, additional):
     children = {}
     for group, hostnames in groups.items():
         if group == 'all' or group == 'ungrouped':
@@ -58,6 +58,15 @@ def run(dest, hostvars, groups, include, exclude):
                     continue
             d[v] = hvars[v]
         out_all[host] = d
+        if additional:
+            for new_var_name, old_var_name in additional.items():
+                old_var = hvars
+                for part in old_var_name.split('.'):
+                    old_var = old_var.get(part)
+                    if old_var is None:
+                        break
+                else:
+                    d[new_var_name] = old_var
 
     with open(dest, 'w') as f:
         f.write(json.dumps(out))
@@ -71,6 +80,7 @@ def ansible_main():
             groups=dict(required=True, type='raw'),
             include_hostvars=dict(type='list'),
             exclude_hostvars=dict(type='list'),
+            additional_hostvars=dict(type='raw'),
         )
     )
 
@@ -81,8 +91,9 @@ def ansible_main():
     groups = p.get('groups')
     include = p.get('include_hostvars')
     exclude = p.get('exclude_hostvars')
+    additional = p.get('additional_hostvars', {})
 
-    run(dest, hostvars, groups, include, exclude)
+    run(dest, hostvars, groups, include, exclude, additional)
 
     module.exit_json(changed=True)
 

roles/write-inventory/tasks/main.yaml

@@ -5,3 +5,4 @@
     groups: "{{ groups }}"
     include_hostvars: "{{ write_inventory_include_hostvars | default(omit) }}"
     exclude_hostvars: "{{ write_inventory_exclude_hostvars | default(omit) }}"
+    additional_hostvars: "{{ write_inventory_additional_hostvars | default(omit) }}"