Azure: add image-id support

This allows users to specify a private image by ID.

Change-Id: Iccb4ec714bcd6c897b4b4d4eb61fd08484c7dc95

doc/source/azure.rst

@@ -293,9 +293,19 @@ section of the configuration.
            long-standing issue with ``ansible_shell_type`` in combination
            with ``become``
 
+      .. attr:: image-id
+         :type: str
+
+         Specifies a private image to use.  Either this field or
+         :attr:`providers.[azure].cloud-images.image-reference` must be
+         provided.
+
       .. attr:: image-reference
          :type: dict
-         :required:
+
+         Specifies a public image to use.  Either this field or
+         :attr:`providers.[azure].cloud-images.image-id` must be
+         provided.
 
          .. attr:: sku
             :type: str

nodepool/driver/azure/adapter.py

@@ -641,7 +641,10 @@ class AzureAdapter(statemachine.Adapter):
             image_reference = {'id': remote_image['id']}
         else:
             image = label.cloud_image
+            if label.cloud_image.image_reference:
                 image_reference = label.cloud_image.image_reference
+            else:
+                image_reference = {'id': label.cloud_image.image_id}
         os_profile = {'computerName': hostname}
         if image.username and image.key:
             linux_config = {

nodepool/driver/azure/config.py

@@ -34,7 +34,8 @@ class AzureProviderCloudImage(ConfigValue):
         self.username = image['username']
         # TODO(corvus): remove zuul_public_key
         self.key = image.get('key', zuul_public_key)
-        self.image_reference = image['image-reference']
+        self.image_reference = image.get('image-reference')
+        self.image_id = image.get('image-id')
         self.python_path = image.get('python-path')
         self.shell_type = image.get('shell-type')
         self.connection_type = image.get('connection-type', 'ssh')
@@ -45,7 +46,7 @@ class AzureProviderCloudImage(ConfigValue):
     @property
     def external_name(self):
         '''Human readable version of external.'''
-        return self.image_id or self.name
+        return self.image_id or self.image_reference or self.name
 
     @staticmethod
     def getSchema():
@@ -56,17 +57,24 @@ class AzureProviderCloudImage(ConfigValue):
             v.Required('offer'): str,
         }
 
-        return {
+        return v.All({
             v.Required('name'): str,
             v.Required('username'): str,
             # TODO(corvus): make required when zuul_public_key removed
             'key': str,
-            v.Required('image-reference'): azure_image_reference,
+            v.Exclusive('image-reference', 'spec'): azure_image_reference,
+            v.Exclusive('image-id', 'spec'): str,
             'connection-type': str,
             'connection-port': int,
             'python-path': str,
             'shell-type': str,
-        }
+        }, {
+            v.Required(
+                v.Any('image-reference', 'image-id'),
+                msg='Provide either "image-reference" or "image-id" keys'
+            ): object,
+            object: object,
+        })
 
 
 class AzureProviderDiskImage(ConfigValue):

nodepool/tests/fixtures/azure-external-image.yaml

@@ -0,0 +1,49 @@
+webapp:
+  port: 8005
+  listen_address: '0.0.0.0'
+
+zookeeper-servers:
+  - host: {zookeeper_host}
+    port: {zookeeper_port}
+    chroot: {zookeeper_chroot}
+
+zookeeper-tls:
+  ca: {zookeeper_ca}
+  cert: {zookeeper_cert}
+  key: {zookeeper_key}
+
+labels:
+  - name: bionic
+    min-ready: 0
+
+providers:
+  - name: azure
+    driver: azure
+    zuul-public-key: ssh-rsa AAAAB3NzaC1yc2EAAAADAQABAAABAQC+mplenM+m6pNY9Un3fpO9eqf808Jrfb3d1gXg7BZVawCvtEZ/cDYvLQ3OF1AeL2kcIC0UAIglM5JXae7yO5CJbJRdkbXvv0u1LvpLxYSPM4ATR0r4IseC5YVxkfJQNi4ixSwTqD4ScEkuCXcSqSU9M+hB+KlnwXoR4IcYHf7vD2Z0Mdwm2ikk3SeERmspmMxx/uz0SPn58QxONuoTlNWQKqDWsV6bRyoPa6HWccMrIH1/e7E69Nw/30oioOQpKBgaDCauh+QkDtSkjRpRMOV47ZFh16Q9DqMgLx+FD8z6++9rsHlB65Zas1xyQsiRCFG09s00b7OR7Xz9ukQ5+vXV
+    resource-group-location: centralus
+    location: centralus
+    resource-group: nodepool
+    auth-path: {auth_path}
+    subnet-id: /subscriptions/c35cf7df-ed75-4c85-be00-535409a85120/resourceGroups/nodepool/providers/Microsoft.Network/virtualNetworks/NodePool/subnets/default
+    cloud-images:
+      - name: bionic
+        username: zuul
+        shell-type: sh
+        image-id: /subscriptions/c35cf7df-ed75-4c85-be00-535409a85120/resourceGroups/nodepool/providers/Microsoft.Compute/images/test-image-1234
+    pools:
+      - name: main
+        max-servers: 10
+        node-attributes:
+          key1: value1
+          key2: value2
+        labels:
+          - name: bionic
+            cloud-image: bionic
+            hardware-profile:
+              vm-size: Standard_B1ls
+            tags:
+              department: R&D
+              team: DevOps
+              systemPurpose: CI
+            user-data: "This is the user data"
+            custom-data: "This is the custom data"

nodepool/tests/unit/test_driver_azure.py

@@ -98,3 +98,44 @@ class TestDriverAzure(tests.DBTestCase):
         self.assertEqual(node.shell_type, None)
         self.assertEqual(node.attributes,
                          {'key1': 'value1', 'key2': 'value2'})
+
+    def test_azure_external_image(self):
+        configfile = self.setup_config(
+            'azure-external-image.yaml',
+            auth_path=self.fake_azure.auth_file.name)
+        pool = self.useNodepool(configfile, watermark_sleep=1)
+        pool.start()
+        req = zk.NodeRequest()
+        req.state = zk.REQUESTED
+        req.node_types.append('bionic')
+
+        self.zk.storeNodeRequest(req)
+        req = self.waitForNodeRequest(req)
+
+        self.assertEqual(req.state, zk.FULFILLED)
+        self.assertNotEqual(req.nodes, [])
+        node = self.zk.getNode(req.nodes[0])
+        self.assertEqual(node.allocated_to, req.id)
+        self.assertEqual(node.state, zk.READY)
+        self.assertIsNotNone(node.launcher)
+        self.assertEqual(node.connection_type, 'ssh')
+        self.assertEqual(node.shell_type, 'sh')
+        self.assertEqual(node.attributes,
+                         {'key1': 'value1', 'key2': 'value2'})
+        self.assertEqual(node.host_keys, ['ssh-rsa FAKEKEY'])
+        self.assertEqual(
+            self.fake_azure.crud['Microsoft.Compute/virtualMachines'].
+            items[0]['properties']['osProfile']['customData'],
+            'VGhpcyBpcyB0aGUgY3VzdG9tIGRhdGE=')  # This is the custom data
+        self.assertEqual(
+            self.fake_azure.crud['Microsoft.Compute/virtualMachines'].
+            requests[0]['properties']['userData'],
+            'VGhpcyBpcyB0aGUgdXNlciBkYXRh')  # This is the user data
+
+        self.assertEqual(
+            self.fake_azure.crud['Microsoft.Compute/virtualMachines'].
+            requests[0]['properties']['storageProfile']
+            ['imageReference']['id'],
+            "/subscriptions/c35cf7df-ed75-4c85-be00-535409a85120"
+            "/resourceGroups/nodepool/providers/Microsoft.Compute"
+            "/images/test-image-1234")