Unit tests for list_files (#643)

Was trying to repro an issue where we saw `list_files` on a `Dir` was
returning the parent folder in addition to the files - was not able to
repro, but adding some tests to check. Also noticed we weren't using a
variable in `Dir.from_local`.

---------

Signed-off-by: Yee Hing Tong <wild-endeavor@users.noreply.github.com>

Author: wild-endeavor

src/flyte/io/_dir.py

@@ -662,7 +662,7 @@ async def _lazy_uploader() -> tuple[str | None, str]:
 
         # todo: in the future, mirror File and set the file to_path here
         output_path = await storage.put(
-            from_path=local_path_str, to_path=remote_destination, recursive=True, batch_size=batch_size
+            from_path=local_path_str, to_path=resolved_remote_path, recursive=True, batch_size=batch_size
         )
         return cls(path=output_path, name=dirname, hash=dir_cache_key)
 

tests/flyte/io_types/test_dirs.py

@@ -510,3 +510,46 @@ async def test_dir_lazy_uploader_preserves_hash(tmp_dir_structure):
         assert not lv.hash  # Empty string or None
     finally:
         _run_mode_var.set(None)
+
+
+@pytest.mark.sandbox
+@pytest.mark.asyncio
+async def test_list_dir_mock_s3(tmp_path, tmp_dir_structure, ctx_with_test_local_s3_stack_raw_data_path):
+    """
+    Test downloading a directory to a directory path.
+    When a directory path is provided, the directory contents should be downloaded directly into that path.
+    """
+    from flyte.storage import S3
+
+    await flyte.init.aio(storage=S3.for_sandbox())
+
+    # Upload to S3
+    uploaded_dir = await Dir.from_local(tmp_dir_structure)
+    uploaded_path = uploaded_dir.path + "/"
+
+    # List files with a trailing /
+    replica_dir = Dir.from_existing_remote(uploaded_path)
+    files = await replica_dir.list_files()
+    assert len(files) == 3
+
+    # List without the /
+    replica_dir = Dir.from_existing_remote(uploaded_dir.path)
+    files = await replica_dir.list_files()
+    assert len(files) == 3
+
+
+@pytest.mark.asyncio
+async def test_list_dir_local_fs(tmp_path, tmp_dir_structure, ctx_with_test_raw_data_path):
+    """
+    Test listing files in a directory using local filesystem storage (no S3).
+    Mirrors test_list_dir_mock_s3 but uses local raw data path instead.
+    """
+    flyte.init()
+
+    # Upload to local "remote"
+    uploaded_dir = await Dir.from_local(tmp_dir_structure)
+
+    # List files
+    replica_dir = Dir(path=uploaded_dir.path + "/")
+    files = await replica_dir.list_files()
+    assert len(files) == 3

tests/internal/storage/test_storage.py

@@ -139,6 +139,46 @@ async def test_storage_exists():
     assert not await storage.exists("/non-existent/test")
 
 
+@pytest.mark.sandbox
+@pytest.mark.asyncio
+async def test_get_underlying_filesystem_upload_download(tmp_path, ctx_with_test_local_s3_stack_raw_data_path):
+    """
+    Sandbox integration test that uses get_underlying_filesystem with the sandbox S3
+    (LocalStack) to upload and download a file.
+    """
+    from flyte.storage import S3
+
+    await flyte.init.aio(storage=S3.for_sandbox())
+
+    # Create a local file with known content
+    original_content = b"hello from sandbox integration test"
+    local_file = tmp_path / "upload_me.txt"
+    local_file.write_bytes(original_content)
+
+    # Upload the file to sandbox S3 via storage.put
+    s3_path = "s3://bucket/tests/default_upload/upload_me.txt"
+    await storage.put(str(local_file), s3_path)
+
+    # Use get_underlying_filesystem to verify the file exists on S3
+    fs = storage.get_underlying_filesystem(path=s3_path)
+    assert fs.exists(s3_path)
+
+    # Download the file back using get_underlying_filesystem
+    downloaded_file = tmp_path / "downloaded.txt"
+    fs.get(s3_path, str(downloaded_file))
+
+    # Verify the downloaded content matches the original
+    assert downloaded_file.read_bytes() == original_content
+
+    # Also upload via the filesystem directly and read back with storage.get
+    s3_path_2 = "s3://bucket/tests/default_upload/fs_uploaded.txt"
+    fs.put(str(local_file), s3_path_2)
+
+    downloaded_file_2 = tmp_path / "downloaded_2.txt"
+    await storage.get(s3_path_2, str(downloaded_file_2))
+    assert downloaded_file_2.exists()
+
+
 @pytest.mark.parametrize(
     "path,expected",
     [