ray-project--ray/python/ray/llm/tests/common/cloud/test_s3_filesystem.py

"""Tests for S3FileSystem class."""

import json
import os
import sys
import tempfile
from unittest.mock import MagicMock, patch

import pytest
from botocore.exceptions import ClientError

from ray.llm._internal.common.utils.cloud_filesystem.s3_filesystem import (
    S3FileSystem,
)


class TestS3FileSystem:
    """Tests for the S3FileSystem class."""

    @patch("boto3.client")
    @pytest.mark.parametrize(
        "decode_as_utf_8,file_content,expected_content",
        [
            (True, b"test file content", "test file content"),
            (False, b"test file content", b"test file content"),
        ],
    )
    def test_get_file(
        self,
        mock_boto_client,
        decode_as_utf_8,
        file_content,
        expected_content,
    ):
        """Test getting a file from S3 as string or bytes."""
        # Setup mock S3 client
        mock_s3_client = MagicMock()
        mock_boto_client.return_value = mock_s3_client

        # Mock get_object response
        mock_body = MagicMock()
        mock_body.read.return_value = file_content
        mock_s3_client.get_object.return_value = {
            "Body": mock_body,
            "ContentLength": len(file_content),
        }

        # Test getting file
        content = S3FileSystem.get_file(
            "s3://bucket/test.txt", decode_as_utf_8=decode_as_utf_8
        )

        assert content == expected_content
        mock_s3_client.get_object.assert_called_once_with(
            Bucket="bucket", Key="test.txt"
        )

    @patch("boto3.client")
    def test_get_file_not_found(self, mock_boto_client):
        """Test getting a non-existent file from S3."""
        # Setup mock S3 client
        mock_s3_client = MagicMock()
        mock_boto_client.return_value = mock_s3_client

        # Simulate NoSuchKey error
        mock_s3_client.get_object.side_effect = ClientError(
            {
                "Error": {
                    "Code": "NoSuchKey",
                    "Message": "The specified key does not exist.",
                }
            },
            "GetObject",
        )

        assert S3FileSystem.get_file("s3://bucket/nonexistent.txt") is None

    @patch("boto3.client")
    def test_get_file_anonymous(self, mock_boto_client):
        """Test getting a file from S3 with anonymous access."""
        # Setup mock S3 client
        mock_s3_client = MagicMock()
        mock_boto_client.return_value = mock_s3_client

        # Mock get_object response
        mock_body = MagicMock()
        mock_body.read.return_value = b"anonymous content"
        mock_s3_client.get_object.return_value = {
            "Body": mock_body,
        }

        # Test getting file with anonymous URI
        content = S3FileSystem.get_file("s3://anonymous@bucket/test.txt")

        assert content == "anonymous content"
        # Verify anonymous config was used (UNSIGNED signature)
        assert mock_boto_client.called

    @patch("boto3.client")
    @pytest.mark.parametrize(
        "uri,expected_prefix",
        [
            ("s3://bucket/parent", "parent/"),
            ("s3://bucket/parent/", "parent/"),
        ],
    )
    def test_list_subfolders(self, mock_boto_client, uri, expected_prefix):
        """Test listing subfolders in S3."""
        # Setup mock S3 client
        mock_s3_client = MagicMock()
        mock_boto_client.return_value = mock_s3_client

        # Mock list_objects_v2 response
        mock_s3_client.list_objects_v2.return_value = {
            "CommonPrefixes": [
                {"Prefix": f"{expected_prefix}folder1/"},
                {"Prefix": f"{expected_prefix}folder2/"},
            ]
        }

        folders = S3FileSystem.list_subfolders(uri)
        assert sorted(folders) == ["folder1", "folder2"]
        mock_s3_client.list_objects_v2.assert_called_once_with(
            Bucket="bucket", Prefix=expected_prefix, Delimiter="/"
        )

    @patch("boto3.client")
    def test_list_subfolders_exception(self, mock_boto_client):
        """Test listing subfolders when operation fails."""
        # Setup mock S3 client
        mock_s3_client = MagicMock()
        mock_boto_client.return_value = mock_s3_client

        mock_s3_client.list_objects_v2.side_effect = Exception("Network error")
        assert S3FileSystem.list_subfolders("s3://bucket/parent") == []

    def test_list_subfolders_invalid_uri(self):
        """Test listing subfolders with invalid URI."""
        # list_subfolders catches all exceptions and returns empty list
        result = S3FileSystem.list_subfolders("gs://bucket/parent")
        assert result == []

    @patch("boto3.client")
    @pytest.mark.parametrize(
        "uri,expected_prefix",
        [
            ("s3://bucket/dir", "dir/"),
            ("s3://bucket/dir/", "dir/"),
        ],
    )
    def test_download_files(self, mock_boto_client, uri, expected_prefix):
        """Test downloading files from S3."""
        # Setup mock S3 client
        mock_s3_client = MagicMock()
        mock_boto_client.return_value = mock_s3_client

        # Mock paginator
        mock_paginator = MagicMock()
        mock_s3_client.get_paginator.return_value = mock_paginator
        mock_paginator.paginate.return_value = [
            {
                "Contents": [
                    {"Key": f"{expected_prefix}file1.txt", "Size": 100},
                    {"Key": f"{expected_prefix}file2.txt", "Size": 200},
                ]
            }
        ]

        # Mock download_file to do nothing
        mock_s3_client.download_file = MagicMock()

        with tempfile.TemporaryDirectory() as tempdir:
            S3FileSystem.download_files(tempdir, uri, max_workers=2)

            # Verify paginator was called correctly
            mock_s3_client.get_paginator.assert_called_with("list_objects_v2")
            mock_paginator.paginate.assert_called_once_with(
                Bucket="bucket", Prefix=expected_prefix
            )

            # Verify files were downloaded
            assert mock_s3_client.download_file.call_count == 2

    @patch("boto3.client")
    def test_download_files_with_filters(self, mock_boto_client):
        """Test downloading files with inclusion and exclusion filters."""
        # Setup mock S3 client
        mock_s3_client = MagicMock()
        mock_boto_client.return_value = mock_s3_client

        # Mock paginator with various files
        mock_paginator = MagicMock()
        mock_s3_client.get_paginator.return_value = mock_paginator
        mock_paginator.paginate.return_value = [
            {
                "Contents": [
                    {"Key": "dir/config.json", "Size": 100},
                    {"Key": "dir/tokenizer.json", "Size": 200},
                    {"Key": "dir/data.tmp", "Size": 300},
                    {"Key": "dir/readme.txt", "Size": 400},
                    {"Key": "dir/other.bin", "Size": 500},
                ]
            }
        ]

        # Mock download_file to do nothing
        mock_s3_client.download_file = MagicMock()

        with tempfile.TemporaryDirectory() as tempdir:
            S3FileSystem.download_files(
                tempdir,
                "s3://bucket/dir",
                substrings_to_include=["config", "tokenizer"],
                suffixes_to_exclude=[".tmp", ".txt"],
                max_workers=2,
            )

            # Should only download config.json and tokenizer.json
            # (included by substring, not excluded by suffix)
            assert mock_s3_client.download_file.call_count == 2

            # Get the keys that were downloaded
            downloaded_keys = [
                call[0][1] for call in mock_s3_client.download_file.call_args_list
            ]
            assert "dir/config.json" in downloaded_keys
            assert "dir/tokenizer.json" in downloaded_keys

    @patch("boto3.client")
    def test_download_files_no_matching_files(self, mock_boto_client):
        """Test downloading when no files match filters."""
        # Setup mock S3 client
        mock_s3_client = MagicMock()
        mock_boto_client.return_value = mock_s3_client

        # Mock paginator with files that won't match
        mock_paginator = MagicMock()
        mock_s3_client.get_paginator.return_value = mock_paginator
        mock_paginator.paginate.return_value = [
            {
                "Contents": [
                    {"Key": "dir/file1.txt", "Size": 100},
                ]
            }
        ]

        with tempfile.TemporaryDirectory() as tempdir:
            # This should not raise, just return without downloading
            S3FileSystem.download_files(
                tempdir,
                "s3://bucket/dir",
                substrings_to_include=["nonexistent"],
                max_workers=2,
            )

            # Verify no files were downloaded
            mock_s3_client.download_file.assert_not_called()

    @patch("boto3.client")
    @patch("ray.llm._internal.common.utils.cloud_filesystem.s3_filesystem.logger")
    def test_download_files_concurrent_failure(self, mock_logger, mock_boto_client):
        """Test downloading files when some downloads fail."""
        # Setup mock S3 client
        mock_s3_client = MagicMock()
        mock_boto_client.return_value = mock_s3_client

        # Mock paginator
        mock_paginator = MagicMock()
        mock_s3_client.get_paginator.return_value = mock_paginator
        mock_paginator.paginate.return_value = [
            {
                "Contents": [
                    {"Key": "dir/file1.txt", "Size": 100},
                    {"Key": "dir/file2.txt", "Size": 200},
                ]
            }
        ]

        # Make download_file fail
        mock_s3_client.download_file.side_effect = Exception("Download failed")

        with tempfile.TemporaryDirectory() as tempdir:
            # Should complete without raising, but log errors
            S3FileSystem.download_files(tempdir, "s3://bucket/dir", max_workers=2)

            # Verify error was logged for failed downloads
            mock_logger.error.assert_called()
            error_call = mock_logger.error.call_args_list[0][0][0]
            assert "Failed to download" in error_call

    def test_download_files_invalid_uri(self):
        """Test downloading files with invalid URI."""
        with tempfile.TemporaryDirectory() as tempdir:
            with pytest.raises(ValueError, match="Invalid S3 URI"):
                S3FileSystem.download_files(tempdir, "gs://bucket/dir")

    @patch("boto3.client")
    @pytest.mark.parametrize(
        "uri,expected_prefix",
        [
            ("s3://bucket/dir", "dir/"),
            ("s3://bucket/dir/", "dir/"),
        ],
    )
    def test_upload_files_directory(self, mock_boto_client, uri, expected_prefix):
        """Test uploading a directory to S3."""
        # Setup mock S3 client
        mock_s3_client = MagicMock()
        mock_boto_client.return_value = mock_s3_client
        mock_s3_client.upload_file = MagicMock()

        with tempfile.TemporaryDirectory() as tempdir:
            # Create some test files
            test_file1 = os.path.join(tempdir, "file1.txt")
            test_file2 = os.path.join(tempdir, "subdir", "file2.txt")
            os.makedirs(os.path.dirname(test_file2), exist_ok=True)

            with open(test_file1, "w") as f:
                f.write("test1")
            with open(test_file2, "w") as f:
                f.write("test2")

            S3FileSystem.upload_files(tempdir, uri)

            # Verify files were uploaded
            assert mock_s3_client.upload_file.call_count == 2

            # Check the S3 keys that were used
            uploaded_keys = [
                call[0][2] for call in mock_s3_client.upload_file.call_args_list
            ]
            assert f"{expected_prefix}file1.txt" in uploaded_keys
            assert f"{expected_prefix}subdir/file2.txt" in uploaded_keys

    @patch("boto3.client")
    def test_upload_files_single_file(self, mock_boto_client):
        """Test uploading a single file to S3."""
        # Setup mock S3 client
        mock_s3_client = MagicMock()
        mock_boto_client.return_value = mock_s3_client
        mock_s3_client.upload_file = MagicMock()

        with tempfile.TemporaryDirectory() as tempdir:
            # Create a test file
            test_file = os.path.join(tempdir, "single.txt")
            with open(test_file, "w") as f:
                f.write("test content")

            S3FileSystem.upload_files(test_file, "s3://bucket/dir/")

            # Verify single file was uploaded
            mock_s3_client.upload_file.assert_called_once()
            call_args = mock_s3_client.upload_file.call_args[0]
            assert call_args[2] == "dir/single.txt"

    @patch("boto3.client")
    def test_upload_files_exception(self, mock_boto_client):
        """Test uploading files when operation fails."""
        # Setup mock S3 client
        mock_s3_client = MagicMock()
        mock_boto_client.return_value = mock_s3_client
        mock_s3_client.upload_file.side_effect = Exception("Network error")

        with tempfile.TemporaryDirectory() as tempdir:
            # Create a test file
            test_file = os.path.join(tempdir, "test.txt")
            with open(test_file, "w") as f:
                f.write("test")

            with pytest.raises(Exception, match="Network error"):
                S3FileSystem.upload_files(tempdir, "s3://bucket/dir")

    def test_upload_files_invalid_uri(self):
        """Test uploading files with invalid URI."""
        with tempfile.TemporaryDirectory() as tempdir:
            with pytest.raises(ValueError, match="Invalid S3 URI"):
                S3FileSystem.upload_files(tempdir, "gs://bucket/dir")

    def test_upload_files_nonexistent_path(self):
        """Test uploading from a path that doesn't exist."""
        with pytest.raises(ValueError, match="does not exist"):
            S3FileSystem.upload_files("/nonexistent/path", "s3://bucket/dir")

    def test_parse_s3_uri(self):
        """Test parsing S3 URIs."""
        # Standard URI
        bucket, key, is_anon = S3FileSystem._parse_s3_uri(
            "s3://bucket/path/to/file.txt"
        )
        assert bucket == "bucket"
        assert key == "path/to/file.txt"
        assert is_anon is False

        # Anonymous URI
        bucket, key, is_anon = S3FileSystem._parse_s3_uri(
            "s3://anonymous@bucket/file.txt"
        )
        assert bucket == "bucket"
        assert key == "file.txt"
        assert is_anon is True

        # Bucket only
        bucket, key, is_anon = S3FileSystem._parse_s3_uri("s3://bucket")
        assert bucket == "bucket"
        assert key == ""
        assert is_anon is False

    def test_calculate_optimal_workers(self):
        """Test worker calculation based on file characteristics."""
        # Many small files (< 1MB)
        workers = S3FileSystem._calculate_optimal_workers(
            num_files=50, total_size=50 * 500 * 1024  # 50 files * 500KB each
        )
        assert workers == 50  # Should use many workers for small files

        # Medium files (1-10MB)
        workers = S3FileSystem._calculate_optimal_workers(
            num_files=50, total_size=50 * 5 * 1024 * 1024  # 50 files * 5MB each
        )
        assert workers == 25  # Should use moderate workers

        # Large files (> 10MB)
        workers = S3FileSystem._calculate_optimal_workers(
            num_files=50, total_size=50 * 50 * 1024 * 1024  # 50 files * 50MB each
        )
        assert workers == 20  # Should cap at 20 for large files

        # Zero files
        workers = S3FileSystem._calculate_optimal_workers(num_files=0, total_size=0)
        assert workers == 10  # Should return default_min


class TestS3FileSystemIntegration:
    """Integration tests for S3FileSystem (requires actual S3 access)."""

    def test_list_subfolders_real_s3(self):
        """Test listing subfolders from real S3 bucket."""
        # Test listing subfolders in the parent directory which has actual subfolders
        folders = S3FileSystem.list_subfolders(
            "s3://anonymous@air-example-data/rayllm-ossci/"
        )
        # Verify we get expected subfolders
        assert isinstance(folders, list)
        assert "meta-Llama-3.2-1B-Instruct" in folders
        assert len(folders) > 0

    def test_get_file_real_s3(self):
        """Test getting a file from real S3 bucket."""
        # Test getting a small config file
        content = S3FileSystem.get_file(
            "s3://anonymous@air-example-data/rayllm-ossci/meta-Llama-3.2-1B-Instruct/config.json"
        )
        assert content is not None
        assert isinstance(content, str)
        # Verify it's valid JSON
        config = json.loads(content)
        assert "model_type" in config or "vocab_size" in config

    def test_download_files_with_exclusion(self):
        """Test downloading files with exclusion filter (exclude safetensors files)."""
        with tempfile.TemporaryDirectory() as tempdir:
            # Download files excluding safetensors
            S3FileSystem.download_files(
                tempdir,
                "s3://anonymous@air-example-data/rayllm-ossci/meta-Llama-3.2-1B-Instruct/",
                suffixes_to_exclude=[".safetensors"],
            )

            # Get list of downloaded files
            downloaded_files = set()
            for root, dirs, files in os.walk(tempdir):
                for file in files:
                    rel_path = os.path.relpath(os.path.join(root, file), tempdir)
                    downloaded_files.add(rel_path)

            # Verify safetensors file is excluded
            assert (
                "model.safetensors" not in downloaded_files
            ), "safetensors file should be excluded"

            # Verify other files are downloaded
            assert "config.json" in downloaded_files
            assert "tokenizer.json" in downloaded_files
            assert len(downloaded_files) > 0

    def test_download_files_with_inclusion(self):
        """Test downloading files with inclusion filter (include only .json files)."""
        with tempfile.TemporaryDirectory() as tempdir:
            # Download only .json files
            S3FileSystem.download_files(
                tempdir,
                "s3://anonymous@air-example-data/rayllm-ossci/meta-Llama-3.2-1B-Instruct/",
                substrings_to_include=[".json"],
            )

            # Get list of downloaded files
            downloaded_files = set()
            for root, dirs, files in os.walk(tempdir):
                for file in files:
                    rel_path = os.path.relpath(os.path.join(root, file), tempdir)
                    downloaded_files.add(rel_path)

            # Verify only .json files are downloaded
            expected_json_files = {
                "config.json",
                "generation_config.json",
                "special_tokens_map.json",
                "tokenizer.json",
                "tokenizer_config.json",
            }
            assert (
                downloaded_files == expected_json_files
            ), f"Expected {expected_json_files}, got {downloaded_files}"

            # Verify non-json files are excluded
            assert "model.safetensors" not in downloaded_files
            assert "README.md" not in downloaded_files
            assert "LICENSE.txt" not in downloaded_files


if __name__ == "__main__":
    sys.exit(pytest.main(["-v", __file__]))