Repository: skypilot-org/skypilot
Pull request: [Tests] Fix test_docker_storage_mounts for azure blob
Closes #3781.

Tested: `pytest tests/test_smoke.py::test_docker_storage_mounts --azure`

Problem statement:
[Azure] `test_docker_storage_mounts` needs to be updated for azure blob storage
`pytest tests/test_smoke.py::test_docker_storage_mounts --azure` fails with:

```
+ aws s3 ls sky-test-17217151744714668/hello.txt || gsutil ls gs://sky-test-17217151744714668/hello.txt

An error occurred (NoSuchBucket) when calling the ListObjectsV2 operation: The specified bucket does not exist
BucketNotFoundException: 404 gs://sky-test-17217151744714668 bucket does not exist.
``` 

We probably need to update this test to check azure blob storage as well. 

Blocking 0.6.1.

Hints from issue discussion:

Commit message: lint

Patch to review:
diff --git a/tests/test_smoke.py b/tests/test_smoke.py
index 325a836cf4c..ec431895ca0 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,17 @@ 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:
+        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 +1211,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 +4288,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 %}


Code review comments:
diff_hunk: @@ -1189,9 +1190,17 @@ 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:
+        include_private_mount = False if generic_cloud == 'azure' else True
path: tests/test_smoke.py
text: @user1:
can we add a comment for why private bucket is not allowed for azure?

@author:
Added!

@user1:
Let's file an issue for this to check the dependencies earlier in the `setup.py`?