From 93b319e8b90ab60d80f37e06802b876246a8e9fd Mon Sep 17 00:00:00 2001
From: Romil Bhardwaj <romil.bhardwaj@berkeley.edu>
Date: Thu, 25 Jul 2024 16:54:05 -0700
Subject: [PATCH] [Tests] Fix test_docker_storage_mounts for azure blob (#3784)

* fix azure smoke tests

* lint

* Add comment
---
 tests/test_smoke.py                           | 28 ++++++++++++++++---
 .../test_yamls/test_storage_mounting.yaml.j2  |  6 ++++
 2 files changed, 30 insertions(+), 4 deletions(-)

diff --git a/tests/test_smoke.py b/tests/test_smoke.py
index 32bf6c23582..0413447dcec 100644
--- a/tests/test_smoke.py
+++ b/tests/test_smoke.py
@@ -47,6 +47,7 @@
 from sky import global_user_state
 from sky import jobs
 from sky import serve
+from sky import skypilot_config
 from sky.adaptors import cloudflare
 from sky.adaptors import ibm
 from sky.clouds import AWS
@@ -1189,9 +1190,22 @@ def test_docker_storage_mounts(generic_cloud: str, image_id: str):
     template = jinja2.Template(template_str)
     # ubuntu 18.04 does not support fuse3, and blobfuse2 depends on fuse3.
     azure_mount_unsupported_ubuntu_version = '18.04'
+    # Commands to verify bucket upload. We need to check all three
+    # storage types because the optimizer may pick any of them.
+    s3_command = f'aws s3 ls {storage_name}/hello.txt'
+    gsutil_command = f'gsutil ls gs://{storage_name}/hello.txt'
+    azure_blob_command = TestStorageWithCredentials.cli_ls_cmd(
+        storage_lib.StoreType.AZURE, storage_name, suffix='hello.txt')
     if azure_mount_unsupported_ubuntu_version in image_id:
+        # The store for mount_private_mount is not specified in the template.
+        # If we're running on Azure, the private mount will be created on
+        # azure blob. That will not be supported on the ubuntu 18.04 image
+        # and thus fail. For other clouds, the private mount on other
+        # storage types (GCS/S3) should succeed.
+        include_private_mount = False if generic_cloud == 'azure' else True
         content = template.render(storage_name=storage_name,
-                                  include_azure_mount=False)
+                                  include_azure_mount=False,
+                                  include_private_mount=include_private_mount)
     else:
         content = template.render(storage_name=storage_name,)
     with tempfile.NamedTemporaryFile(suffix='.yaml', mode='w') as f:
@@ -1202,8 +1216,10 @@ def test_docker_storage_mounts(generic_cloud: str, image_id: str):
             *STORAGE_SETUP_COMMANDS,
             f'sky launch -y -c {name} --cloud {generic_cloud} --image-id {image_id} {file_path}',
             f'sky logs {name} 1 --status',  # Ensure job succeeded.
-            f'aws s3 ls {storage_name}/hello.txt || '
-            f'gsutil ls gs://{storage_name}/hello.txt',
+            # Check AWS, GCP, or Azure storage mount.
+            f'{s3_command} || '
+            f'{gsutil_command} || '
+            f'{azure_blob_command}',
         ]
         test = Test(
             'docker_storage_mounts',
@@ -4277,7 +4293,11 @@ def cli_ls_cmd(store_type, bucket_name, suffix=''):
             return f'gsutil ls {url}'
         if store_type == storage_lib.StoreType.AZURE:
             default_region = 'eastus'
-            storage_account_name = (
+            config_storage_account = skypilot_config.get_nested(
+                ('azure', 'storage_account'), None)
+            storage_account_name = config_storage_account if (
+                config_storage_account is not None
+            ) else (
                 storage_lib.AzureBlobStore.DEFAULT_STORAGE_ACCOUNT_NAME.format(
                     region=default_region,
                     user_hash=common_utils.get_user_hash()))
diff --git a/tests/test_yamls/test_storage_mounting.yaml.j2 b/tests/test_yamls/test_storage_mounting.yaml.j2
index c61250bae14..4241c63409e 100644
--- a/tests/test_yamls/test_storage_mounting.yaml.j2
+++ b/tests/test_yamls/test_storage_mounting.yaml.j2
@@ -28,11 +28,13 @@ file_mounts:
     source: ['~/tmp-workdir/tmp file', '~/tmp-workdir/tmp file2']
     mode: COPY
 
+  {% if include_private_mount | default(True) %}
   # Mounting private buckets in MOUNT mode
   /mount_private_mount:
     name: {{storage_name}}
     source: ~/tmp-workdir
     mode: MOUNT
+  {% endif %}
 
 run: |
   set -ex
@@ -49,12 +51,16 @@ run: |
   ls -ltr /mount_private_copy/tmp\ file
   ls -ltr /mount_private_copy_lof/tmp\ file
   ls -ltr /mount_private_copy_lof/tmp\ file2
+  {% if include_private_mount | default(True) %}
   ls -ltr /mount_private_mount/foo
   ls -ltr /mount_private_mount/tmp\ file
+  {% endif %}
   
   # Symlinks are not copied to buckets
   ! ls /mount_private_copy/circle-link
+  {% if include_private_mount | default(True) %}
   ! ls /mount_private_mount/circle-link
   
   # Write to private bucket in MOUNT mode should pass
   echo "hello" > /mount_private_mount/hello.txt
+  {% endif %}