From cfb2283b6dd6fd47451330a7eb5c315c84646ff3 Mon Sep 17 00:00:00 2001
From: "James E. Blair" <jim@acmegating.com>
Date: Mon, 23 Aug 2021 13:53:30 -0700
Subject: [PATCH] Support node-attributes in Azure driver

This is intended to be supported by all drivers, but we missed it
in the recent refactor.

Change-Id: Ib3f8c090b06c8458eb2f492db33e77a763aa64d0
---
 doc/source/azure.rst                                        | 6 ++++++
 nodepool/driver/azure/config.py                             | 1 +
 nodepool/tests/fixtures/azure-diskimage.yaml                | 3 +++
 nodepool/tests/fixtures/azure.yaml                          | 5 ++++-
 nodepool/tests/unit/test_driver_azure.py                    | 4 ++++
 .../notes/azure-node-attributes-8c3bb8bea3aea287.yaml       | 4 ++++
 6 files changed, 22 insertions(+), 1 deletion(-)
 create mode 100644 releasenotes/notes/azure-node-attributes-8c3bb8bea3aea287.yaml

diff --git a/doc/source/azure.rst b/doc/source/azure.rst
index 368d87f21..bda594655 100644
--- a/doc/source/azure.rst
+++ b/doc/source/azure.rst
@@ -392,6 +392,12 @@ section of the configuration.
 
           A unique name within the provider for this pool of resources.
 
+       .. attr:: node-attributes
+          :type: dict
+
+          A dictionary of key-value pairs that will be stored with the node data
+          in ZooKeeper. The keys and values can be any arbitrary string.
+
        .. attr:: ipv4
           :type: bool
 
diff --git a/nodepool/driver/azure/config.py b/nodepool/driver/azure/config.py
index 1b20c7f4b..a9bc9db99 100644
--- a/nodepool/driver/azure/config.py
+++ b/nodepool/driver/azure/config.py
@@ -166,6 +166,7 @@ class AzurePool(ConfigPool):
         self.load(pool_config)
 
     def load(self, pool_config):
+        super().load(pool_config)
         self.name = pool_config['name']
         self.max_servers = pool_config['max-servers']
         self.public_ipv4 = pool_config.get('public-ipv4',
diff --git a/nodepool/tests/fixtures/azure-diskimage.yaml b/nodepool/tests/fixtures/azure-diskimage.yaml
index 1999a60e0..825cc8955 100644
--- a/nodepool/tests/fixtures/azure-diskimage.yaml
+++ b/nodepool/tests/fixtures/azure-diskimage.yaml
@@ -35,6 +35,9 @@ providers:
     pools:
       - name: main
         max-servers: 10
+        node-attributes:
+          key1: value1
+          key2: value2
         labels:
           - name: bionic
             diskimage: fake-image
diff --git a/nodepool/tests/fixtures/azure.yaml b/nodepool/tests/fixtures/azure.yaml
index eab16defa..97e94bc6a 100644
--- a/nodepool/tests/fixtures/azure.yaml
+++ b/nodepool/tests/fixtures/azure.yaml
@@ -36,6 +36,9 @@ providers:
     pools:
       - name: main
         max-servers: 10
+        node-attributes:
+          key1: value1
+          key2: value2
         labels:
           - name: bionic
             cloud-image: bionic
@@ -44,4 +47,4 @@ providers:
             tags:
               department: R&D
               team: DevOps
-              systemPurpose: CI
\ No newline at end of file
+              systemPurpose: CI
diff --git a/nodepool/tests/unit/test_driver_azure.py b/nodepool/tests/unit/test_driver_azure.py
index af013efae..4e9544c79 100644
--- a/nodepool/tests/unit/test_driver_azure.py
+++ b/nodepool/tests/unit/test_driver_azure.py
@@ -54,6 +54,8 @@ class TestDriverAzure(tests.DBTestCase):
         self.assertEqual(node.state, zk.READY)
         self.assertIsNotNone(node.launcher)
         self.assertEqual(node.connection_type, 'ssh')
+        self.assertEqual(node.attributes,
+                         {'key1': 'value1', 'key2': 'value2'})
 
     def test_azure_diskimage(self):
         configfile = self.setup_config(
@@ -83,3 +85,5 @@ class TestDriverAzure(tests.DBTestCase):
         self.assertEqual(node.state, zk.READY)
         self.assertIsNotNone(node.launcher)
         self.assertEqual(node.connection_type, 'ssh')
+        self.assertEqual(node.attributes,
+                         {'key1': 'value1', 'key2': 'value2'})
diff --git a/releasenotes/notes/azure-node-attributes-8c3bb8bea3aea287.yaml b/releasenotes/notes/azure-node-attributes-8c3bb8bea3aea287.yaml
new file mode 100644
index 000000000..73b98a99a
--- /dev/null
+++ b/releasenotes/notes/azure-node-attributes-8c3bb8bea3aea287.yaml
@@ -0,0 +1,4 @@
+---
+features:
+  - |
+    The node-attributes setting has been added to the Azure driver.