Fix docker image passing through in AzureVMCluster (#350)

dask · May 16, 2022 · aee9a0b · aee9a0b
1 parent e5cbf93
commit aee9a0b
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 1 deletion.
diff --git a/dask_cloudprovider/azure/azurevm.py b/dask_cloudprovider/azure/azurevm.py
@@ -494,6 +494,9 @@ def __init__(
         self.security_group = self.config.get(
             "azurevm.security_group", override_with=security_group
         )
+        self.docker_image = self.config.get(
+            "azurevm.docker_image", override_with=docker_image
+        )
         if self.security_group is None:
             raise ConfigError(
                 "You must configure a security group which allows traffic on 8786 and 8787"

diff --git a/dask_cloudprovider/azure/tests/test_azurevm.py b/dask_cloudprovider/azure/tests/test_azurevm.py
@@ -127,7 +127,9 @@ async def test_render_cloud_init():
     assert " --privileged " in cloud_init
 
     cloud_init = AzureVMCluster.get_cloud_init(
-        extra_bootstrap=["echo 'hello world'", "echo 'foo bar'"]
+        docker_image="foo/bar:baz",
+        extra_bootstrap=["echo 'hello world'", "echo 'foo bar'"],
     )
+    assert "foo/bar:baz" in cloud_init
     assert "- echo 'hello world'" in cloud_init
     assert "- echo 'foo bar'" in cloud_init