TheRock/build_tools/tests/artifact_backend_test.py

#!/usr/bin/env python
# Copyright Advanced Micro Devices, Inc.
# SPDX-License-Identifier: MIT

"""Unit tests for artifact_backend.py."""

import os
import socket
import sys
import tempfile
import unittest
from botocore import UNSIGNED
from pathlib import Path
from unittest import mock

sys.path.insert(0, os.fspath(Path(__file__).parent.parent))

from _therock_utils.artifact_backend import (
    ArtifactBackend,
    LocalDirectoryBackend,
    S3Backend,
    create_backend_from_env,
)
from _therock_utils.workflow_outputs import WorkflowOutputRoot


def _make_local_root(run_id="test-run-123", platform="linux"):
    return WorkflowOutputRoot.for_local(run_id=run_id, platform=platform)


def _make_s3_root(
    bucket="test-bucket",
    run_id="test-run-456",
    platform="linux",
    external_repo="external/",
):
    return WorkflowOutputRoot(
        bucket=bucket,
        external_repo=external_repo,
        run_id=run_id,
        platform=platform,
    )


def _is_public_s3_available():
    try:
        with socket.create_connection(
            ("therock-ci-artifacts.s3.amazonaws.com", 443), timeout=5
        ):
            return True
    except OSError:
        return False


class TestLocalDirectoryBackend(unittest.TestCase):
    """Tests for LocalDirectoryBackend."""

    def setUp(self):
        """Create a temporary directory for testing."""
        self.temp_dir = tempfile.mkdtemp()
        self.output_root = _make_local_root()
        self.backend = LocalDirectoryBackend(
            staging_dir=Path(self.temp_dir),
            output_root=self.output_root,
        )

    def tearDown(self):
        """Clean up temporary directory."""
        import shutil

        shutil.rmtree(self.temp_dir, ignore_errors=True)

    def test_base_uri(self):
        """Test that base_uri returns the correct path."""
        expected = str(Path(self.temp_dir) / "test-run-123-linux")
        self.assertEqual(self.backend.base_uri, expected)

    def test_base_path_created(self):
        """Test that the base path directory is created."""
        self.assertTrue(self.backend.base_path.exists())
        self.assertTrue(self.backend.base_path.is_dir())

    def test_list_artifacts_empty(self):
        """Test listing artifacts when none exist."""
        artifacts = self.backend.list_artifacts()
        self.assertEqual(artifacts, [])

    def test_list_artifacts_with_files(self):
        """Test listing artifacts with tar.zst and tar.xz files present."""
        # Create some test artifact files (both zstd and xz)
        (self.backend.base_path / "blas_lib_gfx94X.tar.zst").touch()
        (self.backend.base_path / "blas_dev_gfx94X.tar.xz").touch()
        (self.backend.base_path / "fft_lib_generic.tar.zst").touch()
        # Create a sha256sum file (should be excluded)
        (self.backend.base_path / "blas_lib_gfx94X.tar.zst.sha256sum").touch()
        # Create a non-artifact file (should be excluded)
        (self.backend.base_path / "README.txt").touch()

        artifacts = self.backend.list_artifacts()
        self.assertEqual(len(artifacts), 3)
        self.assertIn("blas_lib_gfx94X.tar.zst", artifacts)
        self.assertIn("blas_dev_gfx94X.tar.xz", artifacts)
        self.assertIn("fft_lib_generic.tar.zst", artifacts)

    def test_list_artifacts_with_name_filter(self):
        """Test filtering artifacts by name prefix."""
        (self.backend.base_path / "blas_lib_gfx94X.tar.xz").touch()
        (self.backend.base_path / "blas_dev_gfx94X.tar.xz").touch()
        (self.backend.base_path / "fft_lib_generic.tar.xz").touch()

        # Filter by "blas"
        artifacts = self.backend.list_artifacts(name_filter="blas")
        self.assertEqual(len(artifacts), 2)
        self.assertIn("blas_lib_gfx94X.tar.xz", artifacts)
        self.assertIn("blas_dev_gfx94X.tar.xz", artifacts)

        # Filter by "fft"
        artifacts = self.backend.list_artifacts(name_filter="fft")
        self.assertEqual(len(artifacts), 1)
        self.assertIn("fft_lib_generic.tar.xz", artifacts)

        # Filter with no matches
        artifacts = self.backend.list_artifacts(name_filter="nonexistent")
        self.assertEqual(len(artifacts), 0)

    def test_upload_and_download_artifact_xz(self):
        """Test uploading and downloading a .tar.xz artifact."""
        # Create a source file to upload
        source_dir = Path(self.temp_dir) / "source"
        source_dir.mkdir()
        source_file = source_dir / "test_artifact.tar.xz"
        source_file.write_text("test artifact content xz")

        # Also create a sha256sum file
        sha_file = source_dir / "test_artifact.tar.xz.sha256sum"
        sha_file.write_text("abc123  test_artifact.tar.xz\n")

        # Upload
        self.backend.upload_artifact(source_file, "test_artifact.tar.xz")

        # Verify it exists in the backend
        self.assertTrue(self.backend.artifact_exists("test_artifact.tar.xz"))
        self.assertTrue(
            (self.backend.base_path / "test_artifact.tar.xz.sha256sum").exists()
        )

        # Download to a new location
        dest_dir = Path(self.temp_dir) / "dest"
        dest_dir.mkdir()
        dest_file = dest_dir / "downloaded.tar.xz"

        self.backend.download_artifact("test_artifact.tar.xz", dest_file)

        # Verify content
        self.assertTrue(dest_file.exists())
        self.assertEqual(dest_file.read_text(), "test artifact content xz")
        # Verify sha256sum was also copied
        self.assertTrue((dest_dir / "test_artifact.tar.xz.sha256sum").exists())

    def test_upload_and_download_artifact_zst(self):
        """Test uploading and downloading a .tar.zst artifact."""
        # Create a source file to upload
        source_dir = Path(self.temp_dir) / "source_zst"
        source_dir.mkdir()
        source_file = source_dir / "test_artifact.tar.zst"
        source_file.write_text("test artifact content zst")

        # Also create a sha256sum file
        sha_file = source_dir / "test_artifact.tar.zst.sha256sum"
        sha_file.write_text("def456  test_artifact.tar.zst\n")

        # Upload
        self.backend.upload_artifact(source_file, "test_artifact.tar.zst")

        # Verify it exists in the backend
        self.assertTrue(self.backend.artifact_exists("test_artifact.tar.zst"))
        self.assertTrue(
            (self.backend.base_path / "test_artifact.tar.zst.sha256sum").exists()
        )

        # Download to a new location
        dest_dir = Path(self.temp_dir) / "dest_zst"
        dest_dir.mkdir()
        dest_file = dest_dir / "downloaded.tar.zst"

        self.backend.download_artifact("test_artifact.tar.zst", dest_file)

        # Verify content
        self.assertTrue(dest_file.exists())
        self.assertEqual(dest_file.read_text(), "test artifact content zst")
        # Verify sha256sum was also copied
        self.assertTrue((dest_dir / "test_artifact.tar.zst.sha256sum").exists())

    def test_download_nonexistent_artifact(self):
        """Test that downloading a nonexistent artifact raises FileNotFoundError."""
        dest_file = Path(self.temp_dir) / "dest" / "nonexistent.tar.xz"
        with self.assertRaises(FileNotFoundError):
            self.backend.download_artifact("nonexistent.tar.xz", dest_file)

    def test_upload_nonexistent_source(self):
        """Test that uploading a nonexistent source raises FileNotFoundError."""
        nonexistent = Path(self.temp_dir) / "nonexistent.tar.xz"
        with self.assertRaises(FileNotFoundError):
            self.backend.upload_artifact(nonexistent, "test.tar.xz")

    def test_artifact_exists_xz(self):
        """Test artifact_exists method with .tar.xz."""
        self.assertFalse(self.backend.artifact_exists("nonexistent.tar.xz"))

        (self.backend.base_path / "exists.tar.xz").touch()
        self.assertTrue(self.backend.artifact_exists("exists.tar.xz"))

    def test_artifact_exists_zst(self):
        """Test artifact_exists method with .tar.zst."""
        self.assertFalse(self.backend.artifact_exists("nonexistent.tar.zst"))

        (self.backend.base_path / "exists.tar.zst").touch()
        self.assertTrue(self.backend.artifact_exists("exists.tar.zst"))

    def test_copy_artifact_between_local_backends(self):
        """Test copy_artifact copies files between two local backends."""
        source = LocalDirectoryBackend(
            staging_dir=Path(self.temp_dir),
            output_root=_make_local_root(run_id="source-run"),
        )
        dest = LocalDirectoryBackend(
            staging_dir=Path(self.temp_dir),
            output_root=_make_local_root(run_id="dest-run"),
        )

        # Create artifact and sha256sum in source
        artifact_key = "test_lib_generic.tar.zst"
        (source.base_path / artifact_key).write_bytes(b"test content")
        (source.base_path / f"{artifact_key}.sha256sum").write_text(
            "abc123  test_lib_generic.tar.zst\n"
        )

        # Copy to dest
        dest.copy_artifact(artifact_key, source)

        # Verify artifact and sha256sum were both copied
        self.assertTrue(dest.artifact_exists(artifact_key))
        self.assertEqual((dest.base_path / artifact_key).read_bytes(), b"test content")
        self.assertTrue((dest.base_path / f"{artifact_key}.sha256sum").exists())

    def test_copy_artifact_without_sha256sum(self):
        """Test copy_artifact works when no sha256sum file exists."""
        source = LocalDirectoryBackend(
            staging_dir=Path(self.temp_dir),
            output_root=_make_local_root(run_id="source-run"),
        )
        dest = LocalDirectoryBackend(
            staging_dir=Path(self.temp_dir),
            output_root=_make_local_root(run_id="dest-run"),
        )

        # Create artifact without sha256sum
        artifact_key = "test_lib_generic.tar.zst"
        (source.base_path / artifact_key).write_bytes(b"test content")

        # Copy should succeed without error
        dest.copy_artifact(artifact_key, source)

        self.assertTrue(dest.artifact_exists(artifact_key))
        self.assertFalse((dest.base_path / f"{artifact_key}.sha256sum").exists())

    def test_copy_artifact_nonexistent_raises(self):
        """Test copy_artifact raises FileNotFoundError for missing source artifact."""
        source = LocalDirectoryBackend(
            staging_dir=Path(self.temp_dir),
            output_root=_make_local_root(run_id="source-run"),
        )
        dest = LocalDirectoryBackend(
            staging_dir=Path(self.temp_dir),
            output_root=_make_local_root(run_id="dest-run"),
        )

        with self.assertRaises(FileNotFoundError):
            dest.copy_artifact("nonexistent.tar.zst", source)

    def test_copy_artifact_wrong_backend_type_raises(self):
        """Test copy_artifact raises TypeError when source is a different backend type."""
        s3_source = S3Backend(
            output_root=_make_s3_root(run_id="source-run"),
        )

        with self.assertRaises(TypeError):
            self.backend.copy_artifact("test.tar.zst", s3_source)


class TestS3Backend(unittest.TestCase):
    """Tests for S3Backend with mocked boto3 client."""

    def setUp(self):
        """Set up S3Backend with mocked client."""
        self.output_root = _make_s3_root()
        self.backend = S3Backend(output_root=self.output_root)

    def test_base_uri(self):
        """Test that base_uri returns the correct S3 URI."""
        expected = "s3://test-bucket/external/test-run-456-linux"
        self.assertEqual(self.backend.base_uri, expected)

    def test_s3_prefix(self):
        """Test that s3_prefix is constructed correctly."""
        self.assertEqual(self.backend.s3_prefix, "external/test-run-456-linux")

    @mock.patch.object(S3Backend, "s3_client", new_callable=mock.PropertyMock)
    def test_list_artifacts(self, mock_client_prop):
        """Test listing S3 artifacts (both zstd and xz)."""
        mock_client = mock.MagicMock()
        mock_client_prop.return_value = mock_client

        # Mock paginator with both zstd and xz artifacts
        mock_paginator = mock.MagicMock()
        mock_client.get_paginator.return_value = mock_paginator
        mock_paginator.paginate.return_value = [
            {
                "Contents": [
                    {"Key": "external/test-run-456-linux/blas_lib_gfx94X.tar.zst"},
                    {"Key": "external/test-run-456-linux/blas_dev_gfx94X.tar.xz"},
                    {
                        "Key": "external/test-run-456-linux/blas_lib_gfx94X.tar.zst.sha256sum"
                    },
                ]
            }
        ]

        artifacts = self.backend.list_artifacts()
        self.assertEqual(len(artifacts), 2)
        self.assertIn("blas_lib_gfx94X.tar.zst", artifacts)
        self.assertIn("blas_dev_gfx94X.tar.xz", artifacts)

    @mock.patch.object(S3Backend, "s3_client", new_callable=mock.PropertyMock)
    def test_list_artifacts_with_filter(self, mock_client_prop):
        """Test listing S3 artifacts with name filter."""
        mock_client = mock.MagicMock()
        mock_client_prop.return_value = mock_client

        mock_paginator = mock.MagicMock()
        mock_client.get_paginator.return_value = mock_paginator
        mock_paginator.paginate.return_value = [
            {
                "Contents": [
                    {"Key": "external/test-run-456-linux/blas_lib_gfx94X.tar.xz"},
                    {"Key": "external/test-run-456-linux/fft_lib_generic.tar.xz"},
                ]
            }
        ]

        artifacts = self.backend.list_artifacts(name_filter="blas")
        self.assertEqual(len(artifacts), 1)
        self.assertIn("blas_lib_gfx94X.tar.xz", artifacts)

    @mock.patch.object(S3Backend, "s3_client", new_callable=mock.PropertyMock)
    def test_download_artifact_xz(self, mock_client_prop):
        """Test downloading a .tar.xz artifact from S3."""
        mock_client = mock.MagicMock()
        mock_client_prop.return_value = mock_client

        with tempfile.TemporaryDirectory() as temp_dir:
            dest_path = Path(temp_dir) / "downloaded.tar.xz"
            self.backend.download_artifact("test.tar.xz", dest_path)

            mock_client.download_file.assert_called_once_with(
                "test-bucket",
                "external/test-run-456-linux/test.tar.xz",
                str(dest_path),
            )

    @mock.patch.object(S3Backend, "s3_client", new_callable=mock.PropertyMock)
    def test_download_artifact_zst(self, mock_client_prop):
        """Test downloading a .tar.zst artifact from S3."""
        mock_client = mock.MagicMock()
        mock_client_prop.return_value = mock_client

        with tempfile.TemporaryDirectory() as temp_dir:
            dest_path = Path(temp_dir) / "downloaded.tar.zst"
            self.backend.download_artifact("test.tar.zst", dest_path)

            mock_client.download_file.assert_called_once_with(
                "test-bucket",
                "external/test-run-456-linux/test.tar.zst",
                str(dest_path),
            )

    @mock.patch.object(S3Backend, "s3_client", new_callable=mock.PropertyMock)
    def test_upload_artifact_xz(self, mock_client_prop):
        """Test uploading a .tar.xz artifact to S3."""
        mock_client = mock.MagicMock()
        mock_client_prop.return_value = mock_client

        with tempfile.TemporaryDirectory() as temp_dir:
            source_path = Path(temp_dir) / "test.tar.xz"
            source_path.touch()

            self.backend.upload_artifact(source_path, "test.tar.xz")

            mock_client.upload_file.assert_called_once_with(
                str(source_path),
                "test-bucket",
                "external/test-run-456-linux/test.tar.xz",
            )

    @mock.patch.object(S3Backend, "s3_client", new_callable=mock.PropertyMock)
    def test_upload_artifact_zst(self, mock_client_prop):
        """Test uploading a .tar.zst artifact to S3."""
        mock_client = mock.MagicMock()
        mock_client_prop.return_value = mock_client

        with tempfile.TemporaryDirectory() as temp_dir:
            source_path = Path(temp_dir) / "test.tar.zst"
            source_path.touch()

            self.backend.upload_artifact(source_path, "test.tar.zst")

            mock_client.upload_file.assert_called_once_with(
                str(source_path),
                "test-bucket",
                "external/test-run-456-linux/test.tar.zst",
            )

    @mock.patch.object(S3Backend, "s3_client", new_callable=mock.PropertyMock)
    def test_artifact_exists_true(self, mock_client_prop):
        """Test artifact_exists when artifact exists."""
        mock_client = mock.MagicMock()
        mock_client_prop.return_value = mock_client

        self.assertTrue(self.backend.artifact_exists("exists.tar.xz"))
        mock_client.head_object.assert_called_once()

    @mock.patch.object(S3Backend, "s3_client", new_callable=mock.PropertyMock)
    def test_artifact_exists_false(self, mock_client_prop):
        """Test artifact_exists when artifact does not exist."""
        mock_client = mock.MagicMock()
        mock_client_prop.return_value = mock_client
        mock_client.head_object.side_effect = Exception("Not found")

        self.assertFalse(self.backend.artifact_exists("nonexistent.tar.xz"))

    @mock.patch.object(S3Backend, "s3_client", new_callable=mock.PropertyMock)
    def test_copy_artifact_same_bucket(self, mock_client_prop):
        """Test S3 server-side copy within the same bucket."""
        mock_client = mock.MagicMock()
        mock_client_prop.return_value = mock_client
        # sha256sum exists in source
        mock_client.head_object.return_value = {}

        source = S3Backend(
            output_root=_make_s3_root(
                bucket="test-bucket", run_id="source-run", external_repo=""
            )
        )
        dest = S3Backend(
            output_root=_make_s3_root(
                bucket="test-bucket", run_id="dest-run", external_repo=""
            )
        )
        # Share the mock client
        source._s3_client = mock_client

        dest.copy_artifact("artifact_lib_generic.tar.zst", source)

        # Verify both artifact and sha256sum were copied
        self.assertEqual(mock_client.copy.call_count, 2)
        mock_client.copy.assert_any_call(
            {
                "Bucket": "test-bucket",
                "Key": "source-run-linux/artifact_lib_generic.tar.zst",
            },
            "test-bucket",
            "dest-run-linux/artifact_lib_generic.tar.zst",
        )
        mock_client.copy.assert_any_call(
            {
                "Bucket": "test-bucket",
                "Key": "source-run-linux/artifact_lib_generic.tar.zst.sha256sum",
            },
            "test-bucket",
            "dest-run-linux/artifact_lib_generic.tar.zst.sha256sum",
        )

    @mock.patch.object(S3Backend, "s3_client", new_callable=mock.PropertyMock)
    def test_copy_artifact_cross_bucket(self, mock_client_prop):
        """Test S3 server-side copy across different buckets."""
        mock_client = mock.MagicMock()
        mock_client_prop.return_value = mock_client
        # sha256sum does not exist in source
        mock_client.head_object.side_effect = Exception("Not found")

        source = S3Backend(
            output_root=_make_s3_root(
                bucket="therock-ci-artifacts",
                run_id="source-run",
                external_repo="",
            )
        )
        dest = S3Backend(
            output_root=_make_s3_root(
                bucket="therock-ci-artifacts-external",
                run_id="dest-run",
                external_repo="ROCm-rocm-libraries/",
            )
        )
        source._s3_client = mock_client

        dest.copy_artifact("artifact_lib_generic.tar.zst", source)

        # Only artifact copied (no sha256sum in source)
        mock_client.copy.assert_called_once_with(
            {
                "Bucket": "therock-ci-artifacts",
                "Key": "source-run-linux/artifact_lib_generic.tar.zst",
            },
            "therock-ci-artifacts-external",
            "ROCm-rocm-libraries/dest-run-linux/artifact_lib_generic.tar.zst",
        )

    def test_copy_artifact_wrong_backend_type_raises(self):
        """Test copy_artifact raises TypeError when source is a different backend type."""
        import tempfile

        local_source = LocalDirectoryBackend(
            staging_dir=Path(tempfile.mkdtemp()),
            output_root=_make_local_root(run_id="source-run"),
        )

        with self.assertRaises(TypeError):
            self.backend.copy_artifact("test.tar.zst", local_source)


class TestS3BackendCredentials(unittest.TestCase):
    """Live tests for S3Backend credential resolution.

    These exercise the real boto3 credential chain (no mocks) to verify:
    - UNSIGNED fallback works for public bucket reads without credentials
    - AWS_SHARED_CREDENTIALS_FILE is used (so fork PRs can upload artifacts)
    """

    # Env vars that could provide credentials to boto3's default chain.
    _CREDENTIAL_ENV_VARS = [
        "AWS_ACCESS_KEY_ID",
        "AWS_SECRET_ACCESS_KEY",
        "AWS_SESSION_TOKEN",
        "AWS_SHARED_CREDENTIALS_FILE",
        "AWS_CONFIG_FILE",
        "AWS_PROFILE",
    ]

    def _make_backend(self):
        output_root = WorkflowOutputRoot(
            bucket="therock-ci-artifacts",
            external_repo="",
            run_id="placeholder",
            platform="linux",
        )
        return S3Backend(output_root=output_root)

    @unittest.skipUnless(
        _is_public_s3_available(),
        "Public S3 endpoint is not reachable (sandbox with no network access?)",
    )
    def test_list_objects_unsigned(self):
        """Listing objects in a public bucket must succeed without credentials."""
        # Point config/credentials files at nonexistent paths so that
        # ~/.aws/credentials on developer machines is not picked up.
        env_overrides = {
            "AWS_CONFIG_FILE": "/nonexistent/aws/config",
            "AWS_SHARED_CREDENTIALS_FILE": "/nonexistent/aws/credentials",
        }
        with mock.patch.dict(os.environ, env_overrides, clear=False):
            for var in self._CREDENTIAL_ENV_VARS:
                if var not in env_overrides:
                    os.environ.pop(var, None)

            backend = self._make_backend()

            config = backend.s3_client._client_config
            self.assertEqual(
                config.signature_version,
                UNSIGNED,
                "Client should be UNSIGNED when all credentials are cleared",
            )

            # A real list call against the public bucket should succeed
            response = backend.s3_client.list_objects_v2(
                Bucket="therock-ci-artifacts",
                MaxKeys=1,
            )
            self.assertIn("Contents", response)
            self.assertGreater(len(response["Contents"]), 0)

    def test_shared_credentials_file_produces_authenticated_client(self):
        """AWS_SHARED_CREDENTIALS_FILE must produce an authenticated client.

        Fork PRs provide credentials via a mounted credentials.ini file,
        not via AWS_ACCESS_KEY_ID/SECRET/TOKEN env vars. This test would
        fail with the old manual 3-env-var check that ignored
        AWS_SHARED_CREDENTIALS_FILE.
        """
        # Write a minimal credentials file. Values are intentionally not
        # real AWS keys — boto3 loads any syntactically valid values.
        with tempfile.NamedTemporaryFile(
            mode="w", suffix=".ini", delete=False
        ) as creds_file:
            creds_file.write(
                "[default]\n"
                "aws_access_key_id = test\n"
                "aws_secret_access_key = test\n"
            )
            creds_path = creds_file.name

        try:
            env_overrides = {
                "AWS_SHARED_CREDENTIALS_FILE": creds_path,
                "AWS_CONFIG_FILE": "/nonexistent/aws/config",
            }
            env_clear = [
                "AWS_ACCESS_KEY_ID",
                "AWS_SECRET_ACCESS_KEY",
                "AWS_SESSION_TOKEN",
                "AWS_PROFILE",
            ]
            with mock.patch.dict(os.environ, env_overrides, clear=False):
                for var in env_clear:
                    os.environ.pop(var, None)

                backend = self._make_backend()
                config = backend.s3_client._client_config

                self.assertNotEqual(
                    config.signature_version,
                    UNSIGNED,
                    "Client must NOT use UNSIGNED when credentials are "
                    "available via AWS_SHARED_CREDENTIALS_FILE. The old "
                    "manual env-var check ignored this credential source, "
                    "breaking uploads from fork PRs.",
                )
        finally:
            os.unlink(creds_path)


class TestCreateBackendFromEnv(unittest.TestCase):
    """Tests for create_backend_from_env factory function."""

    def test_local_backend_from_env(self):
        """Test that LocalDirectoryBackend is created when env var is set."""
        with tempfile.TemporaryDirectory() as temp_dir:
            with mock.patch.dict(
                os.environ,
                {
                    "THEROCK_LOCAL_STAGING_DIR": temp_dir,
                    "THEROCK_RUN_ID": "env-run-id",
                    "THEROCK_PLATFORM": "windows",
                },
            ):
                backend = create_backend_from_env()

                self.assertIsInstance(backend, LocalDirectoryBackend)
                self.assertIn("env-run-id", backend.base_uri)
                self.assertIn("windows", backend.base_uri)

    def test_local_backend_with_overrides(self):
        """Test LocalDirectoryBackend with explicit run_id and platform."""
        with tempfile.TemporaryDirectory() as temp_dir:
            with mock.patch.dict(
                os.environ,
                {
                    "THEROCK_LOCAL_STAGING_DIR": temp_dir,
                },
            ):
                backend = create_backend_from_env(
                    run_id="override-run", platform="override-platform"
                )

                self.assertIsInstance(backend, LocalDirectoryBackend)
                self.assertIn("override-run", backend.base_uri)
                self.assertIn("override-platform", backend.base_uri)

    @mock.patch("_therock_utils.workflow_outputs._retrieve_bucket_info")
    def test_s3_backend_when_no_local_dir(self, mock_retrieve):
        """Test that S3Backend is created when THEROCK_LOCAL_STAGING_DIR is not set."""
        mock_retrieve.return_value = ("", "test-bucket")
        with mock.patch.dict(
            os.environ,
            {"THEROCK_RUN_ID": "s3-run-id"},
            clear=True,
        ):
            backend = create_backend_from_env()

            self.assertIsInstance(backend, S3Backend)
            self.assertEqual(backend.bucket, "test-bucket")
            self.assertIn("s3-run-id", backend.s3_prefix)

    @mock.patch("_therock_utils.workflow_outputs._retrieve_bucket_info")
    def test_s3_backend_with_github_repository_override(self, mock_retrieve):
        """Test that explicit github_repository overrides GITHUB_REPOSITORY env var.

        When running on a fork (GITHUB_REPOSITORY=SomeUser/TheRock), passing
        github_repository="ROCm/TheRock" should use the main repo's bucket.
        """
        mock_retrieve.return_value = ("", "therock-ci-artifacts")
        with mock.patch.dict(
            os.environ,
            {
                "THEROCK_RUN_ID": "s3-run-id",
                "GITHUB_REPOSITORY": "SomeUser/TheRock",
            },
            clear=True,
        ):
            backend = create_backend_from_env(github_repository="ROCm/TheRock")

            self.assertIsInstance(backend, S3Backend)
            self.assertEqual(backend.bucket, "therock-ci-artifacts")
            # ROCm/TheRock has no external_repo prefix
            self.assertNotIn("SomeUser", backend.s3_prefix)


if __name__ == "__main__":
    unittest.main()