#!/usr/bin/env python """Unit tests for storage_backend.py.""" import os import sys import tempfile import unittest from pathlib import Path from unittest import mock sys.path.insert(0, os.fspath(Path(__file__).parent.parent)) from _therock_utils.workflow_outputs import WorkflowOutputRoot from _therock_utils.storage_location import StorageLocation from _therock_utils.storage_backend import ( LocalStorageBackend, S3StorageBackend, StorageBackend, create_storage_backend, infer_content_type, ) # --------------------------------------------------------------------------- # Content-type inference # --------------------------------------------------------------------------- class TestInferContentType(unittest.TestCase): def test_html(self): self.assertEqual(infer_content_type(Path("index.html")), "text/html") def test_json(self): self.assertEqual(infer_content_type(Path("manifest.json")), "application/json") def test_log(self): self.assertEqual(infer_content_type(Path("build.log")), "text/plain") def test_markdown(self): self.assertEqual(infer_content_type(Path("summary.md")), "text/plain") def test_gzip(self): self.assertEqual(infer_content_type(Path("ninja.tar.gz")), "application/gzip") def test_xz(self): self.assertEqual(infer_content_type(Path("core.tar.xz")), "application/x-xz") def test_zstd(self): self.assertEqual(infer_content_type(Path("core.tar.zst")), "application/zstd") def test_whl(self): ct = infer_content_type(Path("rocm-1.0-py3-none-any.whl")) self.assertIn(ct, ("application/zip", "application/octet-stream")) def test_unknown_extension(self): self.assertEqual( infer_content_type(Path("data.xyz123")), "application/octet-stream" ) def test_sha256sum(self): # .sha256sum is not a known extension - falls to default. ct = infer_content_type(Path("core.tar.xz.sha256sum")) # The suffix is ".sha256sum", not ".xz". self.assertEqual(ct, "application/octet-stream") def test_case_insensitive(self): self.assertEqual(infer_content_type(Path("page.HTML")), "text/html") # --------------------------------------------------------------------------- # WorkflowOutputRoot.root() # --------------------------------------------------------------------------- class TestWorkflowOutputRootRoot(unittest.TestCase): def test_root_returns_output_location(self): rr = WorkflowOutputRoot( bucket="my-bucket", external_repo="", run_id="123", platform="linux" ) loc = rr.root() self.assertIsInstance(loc, StorageLocation) self.assertEqual(loc.bucket, "my-bucket") self.assertEqual(loc.relative_path, "123-linux") def test_root_with_external_repo(self): rr = WorkflowOutputRoot( bucket="b", external_repo="owner-repo/", run_id="99", platform="windows" ) loc = rr.root() self.assertEqual(loc.relative_path, "owner-repo/99-windows") # --------------------------------------------------------------------------- # LocalStorageBackend # --------------------------------------------------------------------------- class TestLocalStorageBackendUploadFile(unittest.TestCase): def test_copies_file(self): with tempfile.TemporaryDirectory() as staging, tempfile.TemporaryDirectory() as src: staging_dir = Path(staging) src_dir = Path(src) source = src_dir / "hello.txt" source.write_text("content") dest = StorageLocation("bucket", "run-1/hello.txt") backend = LocalStorageBackend(staging_dir) backend.upload_file(source, dest) target = staging_dir / "run-1" / "hello.txt" self.assertTrue(target.is_file()) self.assertEqual(target.read_text(), "content") def test_creates_parent_dirs(self): with tempfile.TemporaryDirectory() as staging, tempfile.TemporaryDirectory() as src: staging_dir = Path(staging) src_dir = Path(src) source = src_dir / "data.json" source.write_text("{}") dest = StorageLocation("bucket", "run-1/deep/nested/data.json") backend = LocalStorageBackend(staging_dir) backend.upload_file(source, dest) target = staging_dir / "run-1" / "deep" / "nested" / "data.json" self.assertTrue(target.is_file()) def test_dry_run_does_not_copy(self): with tempfile.TemporaryDirectory() as staging, tempfile.TemporaryDirectory() as src: staging_dir = Path(staging) src_dir = Path(src) source = src_dir / "hello.txt" source.write_text("content") dest = StorageLocation("bucket", "run-1/hello.txt") backend = LocalStorageBackend(staging_dir, dry_run=True) backend.upload_file(source, dest) target = staging_dir / "run-1" / "hello.txt" self.assertFalse(target.exists()) class TestLocalStorageBackendUploadDirectory(unittest.TestCase): def _make_tree(self, base: Path): """Create a test directory tree: base/ file1.tar.xz file1.tar.xz.sha256sum file2.log sub/ nested.html deep/ deep.txt """ base.mkdir(parents=True, exist_ok=True) (base / "file1.tar.xz").write_bytes(b"xz-data") (base / "file1.tar.xz.sha256sum").write_text("abc123") (base / "file2.log").write_text("log line") sub = base / "sub" sub.mkdir() (sub / "nested.html").write_text("") deep = sub / "deep" deep.mkdir() (deep / "deep.txt").write_text("deep content") def test_upload_all_files(self): with tempfile.TemporaryDirectory() as staging, tempfile.TemporaryDirectory() as src: staging_dir = Path(staging) source_dir = Path(src) / "artifacts" self._make_tree(source_dir) dest = StorageLocation("bucket", "run-1") backend = LocalStorageBackend(staging_dir) count = backend.upload_directory(source_dir, dest) self.assertEqual(count, 5) self.assertTrue((staging_dir / "run-1" / "file1.tar.xz").is_file()) self.assertTrue( (staging_dir / "run-1" / "file1.tar.xz.sha256sum").is_file() ) self.assertTrue((staging_dir / "run-1" / "file2.log").is_file()) self.assertTrue((staging_dir / "run-1" / "sub" / "nested.html").is_file()) def test_upload_with_include_filter(self): with tempfile.TemporaryDirectory() as staging, tempfile.TemporaryDirectory() as src: staging_dir = Path(staging) source_dir = Path(src) / "artifacts" self._make_tree(source_dir) dest = StorageLocation("bucket", "run-1") backend = LocalStorageBackend(staging_dir) count = backend.upload_directory(source_dir, dest, include=["*.tar.xz*"]) # Only .tar.xz and .tar.xz.sha256sum should match self.assertEqual(count, 2) self.assertTrue((staging_dir / "run-1" / "file1.tar.xz").is_file()) self.assertTrue( (staging_dir / "run-1" / "file1.tar.xz.sha256sum").is_file() ) self.assertFalse((staging_dir / "run-1" / "file2.log").exists()) self.assertFalse((staging_dir / "run-1" / "sub" / "nested.html").exists()) def test_preserves_subdirectory_structure(self): with tempfile.TemporaryDirectory() as staging, tempfile.TemporaryDirectory() as src: staging_dir = Path(staging) source_dir = Path(src) / "logs" self._make_tree(source_dir) dest = StorageLocation("bucket", "run-1/logs/gfx94X") backend = LocalStorageBackend(staging_dir) backend.upload_directory(source_dir, dest) self.assertTrue( ( staging_dir / "run-1" / "logs" / "gfx94X" / "sub" / "nested.html" ).is_file() ) def test_skips_symlinks(self): with tempfile.TemporaryDirectory() as staging, tempfile.TemporaryDirectory() as src: staging_dir = Path(staging) source_dir = Path(src) / "artifacts" source_dir.mkdir() (source_dir / "real.tar.xz").write_bytes(b"data") try: (source_dir / "link.tar.xz").symlink_to(source_dir / "real.tar.xz") except OSError: self.skipTest("Cannot create symlinks on this platform") dest = StorageLocation("bucket", "run-1") backend = LocalStorageBackend(staging_dir) count = backend.upload_directory(source_dir, dest) self.assertEqual(count, 1) self.assertTrue((staging_dir / "run-1" / "real.tar.xz").is_file()) self.assertFalse((staging_dir / "run-1" / "link.tar.xz").exists()) def test_nonexistent_source_raises(self): with tempfile.TemporaryDirectory() as staging: staging_dir = Path(staging) dest = StorageLocation("bucket", "run-1") backend = LocalStorageBackend(staging_dir) with self.assertRaises(FileNotFoundError): backend.upload_directory(Path("/nonexistent"), dest) def test_dry_run_does_not_copy(self): with tempfile.TemporaryDirectory() as staging, tempfile.TemporaryDirectory() as src: staging_dir = Path(staging) source_dir = Path(src) / "artifacts" self._make_tree(source_dir) dest = StorageLocation("bucket", "run-1") backend = LocalStorageBackend(staging_dir, dry_run=True) count = backend.upload_directory(source_dir, dest) # Count should reflect files that would be uploaded. self.assertEqual(count, 5) # But nothing should actually be written. self.assertFalse((staging_dir / "run-1").exists()) def test_exclude_direct_children_only(self): """sub/* excludes direct children but not deeply nested files.""" with tempfile.TemporaryDirectory() as staging, tempfile.TemporaryDirectory() as src: staging_dir = Path(staging) source_dir = Path(src) / "artifacts" self._make_tree(source_dir) dest = StorageLocation("bucket", "run-1") backend = LocalStorageBackend(staging_dir) count = backend.upload_directory(source_dir, dest, exclude=["sub/*"]) # sub/nested.html excluded, but sub/deep/deep.txt still uploaded. self.assertEqual(count, 4) self.assertTrue((staging_dir / "run-1" / "file1.tar.xz").is_file()) self.assertFalse((staging_dir / "run-1" / "sub" / "nested.html").exists()) self.assertTrue( (staging_dir / "run-1" / "sub" / "deep" / "deep.txt").is_file() ) def test_exclude_recursive(self): """sub/**/* excludes files at all depths.""" with tempfile.TemporaryDirectory() as staging, tempfile.TemporaryDirectory() as src: staging_dir = Path(staging) source_dir = Path(src) / "artifacts" self._make_tree(source_dir) dest = StorageLocation("bucket", "run-1") backend = LocalStorageBackend(staging_dir) count = backend.upload_directory(source_dir, dest, exclude=["sub/**/*"]) # sub/nested.html and sub/deep/deep.txt both excluded. self.assertEqual(count, 3) self.assertTrue((staging_dir / "run-1" / "file1.tar.xz").is_file()) self.assertTrue((staging_dir / "run-1" / "file2.log").is_file()) self.assertFalse((staging_dir / "run-1" / "sub" / "nested.html").exists()) self.assertFalse( (staging_dir / "run-1" / "sub" / "deep" / "deep.txt").exists() ) def test_exclude_with_include(self): with tempfile.TemporaryDirectory() as staging, tempfile.TemporaryDirectory() as src: staging_dir = Path(staging) source_dir = Path(src) / "artifacts" self._make_tree(source_dir) dest = StorageLocation("bucket", "run-1") backend = LocalStorageBackend(staging_dir) # Include all files, but exclude .log files. count = backend.upload_directory(source_dir, dest, exclude=["*.log"]) self.assertEqual(count, 4) self.assertTrue((staging_dir / "run-1" / "file1.tar.xz").is_file()) self.assertTrue((staging_dir / "run-1" / "sub" / "nested.html").is_file()) self.assertFalse((staging_dir / "run-1" / "file2.log").exists()) def test_empty_directory_returns_zero(self): with tempfile.TemporaryDirectory() as staging, tempfile.TemporaryDirectory() as src: staging_dir = Path(staging) source_dir = Path(src) / "empty" source_dir.mkdir() dest = StorageLocation("bucket", "run-1") backend = LocalStorageBackend(staging_dir) count = backend.upload_directory(source_dir, dest) self.assertEqual(count, 0) class TestLocalStorageBackendCopyFile(unittest.TestCase): def test_copies_between_locations(self): with tempfile.TemporaryDirectory() as staging: staging_dir = Path(staging) # Create source file src_path = staging_dir / "run-1" / "file.whl" src_path.parent.mkdir(parents=True) src_path.write_text("wheel-data") source = StorageLocation("bucket", "run-1/file.whl") dest = StorageLocation("bucket", "release/file.whl") backend = LocalStorageBackend(staging_dir) backend.copy_file(source, dest) dst_path = staging_dir / "release" / "file.whl" self.assertTrue(dst_path.is_file()) self.assertEqual(dst_path.read_text(), "wheel-data") def test_creates_parent_dirs(self): with tempfile.TemporaryDirectory() as staging: staging_dir = Path(staging) src_path = staging_dir / "a" / "b.txt" src_path.parent.mkdir(parents=True) src_path.write_text("data") source = StorageLocation("bucket", "a/b.txt") dest = StorageLocation("bucket", "x/y/z/b.txt") backend = LocalStorageBackend(staging_dir) backend.copy_file(source, dest) self.assertTrue((staging_dir / "x" / "y" / "z" / "b.txt").is_file()) def test_dry_run_does_not_copy(self): with tempfile.TemporaryDirectory() as staging: staging_dir = Path(staging) src_path = staging_dir / "run-1" / "file.whl" src_path.parent.mkdir(parents=True) src_path.write_text("wheel-data") source = StorageLocation("bucket", "run-1/file.whl") dest = StorageLocation("bucket", "release/file.whl") backend = LocalStorageBackend(staging_dir, dry_run=True) backend.copy_file(source, dest) self.assertFalse((staging_dir / "release" / "file.whl").exists()) # --------------------------------------------------------------------------- # S3StorageBackend # --------------------------------------------------------------------------- class TestS3StorageBackendUploadFile(unittest.TestCase): def test_calls_boto3_upload_file(self): backend = S3StorageBackend() mock_client = mock.MagicMock() backend._s3_client = mock_client source = Path("/tmp/build.log") dest = StorageLocation("my-bucket", "run-1/logs/build.log") backend.upload_file(source, dest) mock_client.upload_file.assert_called_once_with( str(source), "my-bucket", "run-1/logs/build.log", ExtraArgs={"ContentType": "text/plain"}, ) def test_content_type_for_html(self): backend = S3StorageBackend() mock_client = mock.MagicMock() backend._s3_client = mock_client source = Path("/tmp/index.html") dest = StorageLocation("my-bucket", "run-1/index.html") backend.upload_file(source, dest) mock_client.upload_file.assert_called_once_with( str(source), "my-bucket", "run-1/index.html", ExtraArgs={"ContentType": "text/html"}, ) def test_retries_on_failure(self): backend = S3StorageBackend() mock_client = mock.MagicMock() mock_client.upload_file.side_effect = [ Exception("transient"), None, # succeeds on second attempt ] backend._s3_client = mock_client source = Path("/tmp/data.json") dest = StorageLocation("bucket", "run-1/data.json") with mock.patch("_therock_utils.storage_backend.time.sleep"): backend.upload_file(source, dest) self.assertEqual(mock_client.upload_file.call_count, 2) def test_raises_after_max_retries(self): backend = S3StorageBackend() mock_client = mock.MagicMock() mock_client.upload_file.side_effect = Exception("persistent") backend._s3_client = mock_client source = Path("/tmp/data.json") dest = StorageLocation("bucket", "run-1/data.json") with mock.patch("_therock_utils.storage_backend.time.sleep"): with self.assertRaises(RuntimeError) as ctx: backend.upload_file(source, dest) self.assertIn("3 attempts", str(ctx.exception)) self.assertEqual(mock_client.upload_file.call_count, 3) def test_dry_run_does_not_call_boto3(self): backend = S3StorageBackend(dry_run=True) mock_client = mock.MagicMock() backend._s3_client = mock_client source = Path("/tmp/build.log") dest = StorageLocation("bucket", "run-1/build.log") backend.upload_file(source, dest) mock_client.upload_file.assert_not_called() class TestS3StorageBackendCopyFile(unittest.TestCase): def test_calls_copy(self): backend = S3StorageBackend() mock_client = mock.MagicMock() backend._s3_client = mock_client source = StorageLocation("src-bucket", "run-1/file.whl") dest = StorageLocation("dest-bucket", "release/file.whl") backend.copy_file(source, dest) mock_client.copy.assert_called_once_with( {"Bucket": "src-bucket", "Key": "run-1/file.whl"}, "dest-bucket", "release/file.whl", ) def test_same_bucket_copy(self): backend = S3StorageBackend() mock_client = mock.MagicMock() backend._s3_client = mock_client source = StorageLocation("bucket", "staging/file.whl") dest = StorageLocation("bucket", "release/file.whl") backend.copy_file(source, dest) mock_client.copy.assert_called_once_with( {"Bucket": "bucket", "Key": "staging/file.whl"}, "bucket", "release/file.whl", ) def test_retries_on_failure(self): backend = S3StorageBackend() mock_client = mock.MagicMock() mock_client.copy.side_effect = [ Exception("transient"), None, ] backend._s3_client = mock_client source = StorageLocation("src", "a.whl") dest = StorageLocation("dst", "b.whl") with mock.patch("_therock_utils.storage_backend.time.sleep"): backend.copy_file(source, dest) self.assertEqual(mock_client.copy.call_count, 2) def test_dry_run_does_not_call_boto3(self): backend = S3StorageBackend(dry_run=True) mock_client = mock.MagicMock() backend._s3_client = mock_client source = StorageLocation("src", "a.whl") dest = StorageLocation("dst", "b.whl") backend.copy_file(source, dest) mock_client.copy.assert_not_called() class TestLocalStorageBackendUploadFiles(unittest.TestCase): def test_uploads_all_files(self): with tempfile.TemporaryDirectory() as staging, tempfile.TemporaryDirectory() as src: staging_dir = Path(staging) src_dir = Path(src) (src_dir / "a.txt").write_text("aaa") (src_dir / "b.txt").write_text("bbb") files = [ (src_dir / "a.txt", StorageLocation("bucket", "run-1/a.txt")), (src_dir / "b.txt", StorageLocation("bucket", "run-1/b.txt")), ] backend = LocalStorageBackend(staging_dir) count = backend.upload_files(files) self.assertEqual(count, 2) self.assertEqual((staging_dir / "run-1" / "a.txt").read_text(), "aaa") self.assertEqual((staging_dir / "run-1" / "b.txt").read_text(), "bbb") def test_empty_list_returns_zero(self): with tempfile.TemporaryDirectory() as staging: backend = LocalStorageBackend(Path(staging)) self.assertEqual(backend.upload_files([]), 0) class TestS3StorageBackendUploadFiles(unittest.TestCase): def test_uploads_all_files_in_parallel(self): backend = S3StorageBackend() mock_client = mock.MagicMock() backend._s3_client = mock_client files = [ (Path("/tmp/a.log"), StorageLocation("bucket", "run-1/a.log")), (Path("/tmp/b.log"), StorageLocation("bucket", "run-1/b.log")), (Path("/tmp/c.log"), StorageLocation("bucket", "run-1/c.log")), ] count = backend.upload_files(files) self.assertEqual(count, 3) self.assertEqual(mock_client.upload_file.call_count, 3) def test_empty_list_returns_zero(self): backend = S3StorageBackend() mock_client = mock.MagicMock() backend._s3_client = mock_client count = backend.upload_files([]) self.assertEqual(count, 0) mock_client.upload_file.assert_not_called() def test_single_file_skips_thread_pool(self): """A single file should not use a thread pool (no overhead).""" backend = S3StorageBackend() mock_client = mock.MagicMock() backend._s3_client = mock_client files = [ (Path("/tmp/only.log"), StorageLocation("bucket", "run-1/only.log")), ] with mock.patch( "_therock_utils.storage_backend.concurrent.futures.ThreadPoolExecutor" ) as mock_pool: count = backend.upload_files(files) self.assertEqual(count, 1) mock_pool.assert_not_called() mock_client.upload_file.assert_called_once() def test_dry_run_skips_thread_pool(self): backend = S3StorageBackend(dry_run=True) mock_client = mock.MagicMock() backend._s3_client = mock_client files = [ (Path("/tmp/a.log"), StorageLocation("bucket", "run-1/a.log")), (Path("/tmp/b.log"), StorageLocation("bucket", "run-1/b.log")), ] with mock.patch( "_therock_utils.storage_backend.concurrent.futures.ThreadPoolExecutor" ) as mock_pool: count = backend.upload_files(files) self.assertEqual(count, 2) mock_pool.assert_not_called() # Dry run: no actual boto3 calls mock_client.upload_file.assert_not_called() def test_error_propagation(self): """If a file fails after retries, upload_files raises RuntimeError.""" backend = S3StorageBackend() mock_client = mock.MagicMock() # Second file always fails (after _s3_retry exhausts retries) def upload_side_effect(filename, bucket, key, **kwargs): if "bad" in key: raise Exception("persistent failure") mock_client.upload_file.side_effect = upload_side_effect backend._s3_client = mock_client files = [ (Path("/tmp/good.log"), StorageLocation("bucket", "run-1/good.log")), (Path("/tmp/bad.log"), StorageLocation("bucket", "run-1/bad.log")), ] with mock.patch("_therock_utils.storage_backend.time.sleep"): with self.assertRaises(RuntimeError) as ctx: backend.upload_files(files) self.assertIn("bad.log", str(ctx.exception)) def test_concurrency_param_wired_through(self): backend = S3StorageBackend(upload_concurrency=5) self.assertEqual(backend._upload_concurrency, 5) class TestS3StorageBackendMaxPoolConnections(unittest.TestCase): def test_default_pool_connections_match_concurrency(self): with mock.patch("boto3.client") as mock_boto3: backend = S3StorageBackend() _ = backend.s3_client mock_boto3.assert_called_once() config = mock_boto3.call_args.kwargs.get("config") self.assertIsNotNone(config) self.assertEqual(config.max_pool_connections, 10) def test_custom_concurrency_sets_pool_connections(self): with mock.patch("boto3.client") as mock_boto3: backend = S3StorageBackend(upload_concurrency=20) _ = backend.s3_client config = mock_boto3.call_args.kwargs.get("config") self.assertEqual(config.max_pool_connections, 20) # --------------------------------------------------------------------------- # LocalStorageBackend.list_files # --------------------------------------------------------------------------- class TestLocalStorageBackendListFiles(unittest.TestCase): def test_lists_all_files(self): with tempfile.TemporaryDirectory() as staging: staging_dir = Path(staging) d = staging_dir / "run-1" / "tarballs" d.mkdir(parents=True) (d / "a.tar.gz").write_bytes(b"a") (d / "b.tar.gz").write_bytes(b"b") (d / "c.log").write_text("log") loc = StorageLocation("bucket", "run-1/tarballs") backend = LocalStorageBackend(staging_dir) files = backend.list_files(loc) names = [f.relative_path.rsplit("/", 1)[-1] for f in files] self.assertEqual(names, ["a.tar.gz", "b.tar.gz", "c.log"]) def test_include_filter(self): with tempfile.TemporaryDirectory() as staging: staging_dir = Path(staging) d = staging_dir / "run-1" / "tarballs" d.mkdir(parents=True) (d / "a.tar.gz").write_bytes(b"a") (d / "b.tar.gz").write_bytes(b"b") (d / "c.log").write_text("log") loc = StorageLocation("bucket", "run-1/tarballs") backend = LocalStorageBackend(staging_dir) files = backend.list_files(loc, include=["*.tar.gz"]) names = [f.relative_path.rsplit("/", 1)[-1] for f in files] self.assertEqual(names, ["a.tar.gz", "b.tar.gz"]) def test_empty_directory(self): with tempfile.TemporaryDirectory() as staging: staging_dir = Path(staging) d = staging_dir / "run-1" / "tarballs" d.mkdir(parents=True) loc = StorageLocation("bucket", "run-1/tarballs") backend = LocalStorageBackend(staging_dir) files = backend.list_files(loc) self.assertEqual(files, []) def test_nonexistent_directory(self): with tempfile.TemporaryDirectory() as staging: loc = StorageLocation("bucket", "run-1/tarballs") backend = LocalStorageBackend(Path(staging)) files = backend.list_files(loc) self.assertEqual(files, []) def test_includes_nested_files(self): with tempfile.TemporaryDirectory() as staging: staging_dir = Path(staging) d = staging_dir / "run-1" / "tarballs" d.mkdir(parents=True) (d / "a.tar.gz").write_bytes(b"a") sub = d / "subdir" sub.mkdir() (sub / "nested.tar.gz").write_bytes(b"nested") loc = StorageLocation("bucket", "run-1/tarballs") backend = LocalStorageBackend(staging_dir) files = backend.list_files(loc) paths = [f.relative_path for f in files] self.assertEqual( paths, [ "run-1/tarballs/a.tar.gz", "run-1/tarballs/subdir/nested.tar.gz", ], ) # --------------------------------------------------------------------------- # S3StorageBackend.list_files # --------------------------------------------------------------------------- class TestS3StorageBackendListFiles(unittest.TestCase): def _make_backend_with_pages(self, pages): """Create an S3StorageBackend with a mock paginator returning *pages*.""" backend = S3StorageBackend() mock_client = mock.MagicMock() mock_paginator = mock.MagicMock() mock_paginator.paginate.return_value = pages mock_client.get_paginator.return_value = mock_paginator backend._s3_client = mock_client return backend def test_lists_all_files(self): pages = [ { "Contents": [ {"Key": "run-1/tarballs/a.tar.gz"}, {"Key": "run-1/tarballs/b.tar.gz"}, ] } ] backend = self._make_backend_with_pages(pages) loc = StorageLocation("bucket", "run-1/tarballs") files = backend.list_files(loc) self.assertEqual(len(files), 2) self.assertEqual(files[0].relative_path, "run-1/tarballs/a.tar.gz") self.assertEqual(files[1].relative_path, "run-1/tarballs/b.tar.gz") def test_include_filter(self): pages = [ { "Contents": [ {"Key": "run-1/tarballs/a.tar.gz"}, {"Key": "run-1/tarballs/index.html"}, ] } ] backend = self._make_backend_with_pages(pages) loc = StorageLocation("bucket", "run-1/tarballs") files = backend.list_files(loc, include=["*.tar.gz"]) self.assertEqual(len(files), 1) self.assertEqual(files[0].relative_path, "run-1/tarballs/a.tar.gz") def test_empty_contents(self): pages = [{"Contents": []}] backend = self._make_backend_with_pages(pages) loc = StorageLocation("bucket", "run-1/tarballs") files = backend.list_files(loc) self.assertEqual(files, []) def test_no_contents_key(self): pages = [{}] backend = self._make_backend_with_pages(pages) loc = StorageLocation("bucket", "run-1/tarballs") files = backend.list_files(loc) self.assertEqual(files, []) def test_skips_prefix_directory_marker(self): pages = [ { "Contents": [ {"Key": "run-1/tarballs/"}, {"Key": "run-1/tarballs/a.tar.gz"}, ] } ] backend = self._make_backend_with_pages(pages) loc = StorageLocation("bucket", "run-1/tarballs") files = backend.list_files(loc) self.assertEqual(len(files), 1) self.assertEqual(files[0].relative_path, "run-1/tarballs/a.tar.gz") def test_multiple_pages(self): pages = [ {"Contents": [{"Key": "run-1/tarballs/a.tar.gz"}]}, {"Contents": [{"Key": "run-1/tarballs/b.tar.gz"}]}, ] backend = self._make_backend_with_pages(pages) loc = StorageLocation("bucket", "run-1/tarballs") files = backend.list_files(loc) self.assertEqual(len(files), 2) def test_paginate_uses_trailing_slash(self): """Prefix should end with / to list directory contents.""" backend = self._make_backend_with_pages([{}]) mock_paginator = backend._s3_client.get_paginator.return_value loc = StorageLocation("bucket", "run-1/tarballs") backend.list_files(loc) mock_paginator.paginate.assert_called_once_with( Bucket="bucket", Prefix="run-1/tarballs/" ) # --------------------------------------------------------------------------- # LocalStorageBackend.copy_directory # --------------------------------------------------------------------------- class TestLocalStorageBackendCopyDirectory(unittest.TestCase): def test_copies_filtered_files(self): with tempfile.TemporaryDirectory() as staging: staging_dir = Path(staging) src_dir = staging_dir / "run-1" / "tarballs" src_dir.mkdir(parents=True) (src_dir / "a.tar.gz").write_bytes(b"tarball-a") (src_dir / "b.tar.gz").write_bytes(b"tarball-b") (src_dir / "index.html").write_text("") source = StorageLocation("src-bucket", "run-1/tarballs") dest = StorageLocation("dst-bucket", "v3") backend = LocalStorageBackend(staging_dir) count = backend.copy_directory(source, dest, include=["*.tar.gz"]) self.assertEqual(count, 2) self.assertEqual( (staging_dir / "v3" / "a.tar.gz").read_bytes(), b"tarball-a" ) self.assertEqual( (staging_dir / "v3" / "b.tar.gz").read_bytes(), b"tarball-b" ) self.assertFalse((staging_dir / "v3" / "index.html").exists()) def test_copies_all_without_include(self): with tempfile.TemporaryDirectory() as staging: staging_dir = Path(staging) src_dir = staging_dir / "run-1" / "tarballs" src_dir.mkdir(parents=True) (src_dir / "a.tar.gz").write_bytes(b"data") (src_dir / "index.html").write_text("") source = StorageLocation("bucket", "run-1/tarballs") dest = StorageLocation("bucket", "v3") backend = LocalStorageBackend(staging_dir) count = backend.copy_directory(source, dest) self.assertEqual(count, 2) def test_empty_source_copies_nothing(self): with tempfile.TemporaryDirectory() as staging: staging_dir = Path(staging) src_dir = staging_dir / "run-1" / "tarballs" src_dir.mkdir(parents=True) source = StorageLocation("bucket", "run-1/tarballs") dest = StorageLocation("bucket", "v3") backend = LocalStorageBackend(staging_dir) count = backend.copy_directory(source, dest) self.assertEqual(count, 0) def test_cross_bucket_locations(self): """Source and dest can have different bucket names.""" with tempfile.TemporaryDirectory() as staging: staging_dir = Path(staging) src_dir = staging_dir / "run-1" / "tarballs" src_dir.mkdir(parents=True) (src_dir / "a.tar.gz").write_bytes(b"data") source = StorageLocation("therock-dev-artifacts", "run-1/tarballs") dest = StorageLocation("therock-dev-tarball", "v3") backend = LocalStorageBackend(staging_dir) count = backend.copy_directory(source, dest, include=["*.tar.gz"]) self.assertEqual(count, 1) # Local backend ignores bucket names for file paths self.assertTrue((staging_dir / "v3" / "a.tar.gz").is_file()) def test_preserves_subdirectory_structure(self): with tempfile.TemporaryDirectory() as staging: staging_dir = Path(staging) src_dir = staging_dir / "run-1" / "packages" src_dir.mkdir(parents=True) sub = src_dir / "gfx94X-dcgpu" sub.mkdir() (sub / "rocm_sdk-7.10.0.whl").write_bytes(b"wheel") (src_dir / "top.whl").write_bytes(b"top") source = StorageLocation("bucket", "run-1/packages") dest = StorageLocation("bucket", "v3") backend = LocalStorageBackend(staging_dir) count = backend.copy_directory(source, dest, include=["*.whl"]) self.assertEqual(count, 2) self.assertTrue( (staging_dir / "v3" / "gfx94X-dcgpu" / "rocm_sdk-7.10.0.whl").is_file() ) self.assertTrue((staging_dir / "v3" / "top.whl").is_file()) def test_dry_run_does_not_copy(self): with tempfile.TemporaryDirectory() as staging: staging_dir = Path(staging) src_dir = staging_dir / "run-1" / "tarballs" src_dir.mkdir(parents=True) (src_dir / "a.tar.gz").write_bytes(b"data") source = StorageLocation("bucket", "run-1/tarballs") dest = StorageLocation("bucket", "v3") backend = LocalStorageBackend(staging_dir, dry_run=True) count = backend.copy_directory(source, dest, include=["*.tar.gz"]) self.assertEqual(count, 1) self.assertFalse((staging_dir / "v3").exists()) # --------------------------------------------------------------------------- # S3StorageBackend.copy_files # --------------------------------------------------------------------------- class TestS3StorageBackendCopyFiles(unittest.TestCase): def test_copies_all_files_in_parallel(self): backend = S3StorageBackend() mock_client = mock.MagicMock() backend._s3_client = mock_client files = [ ( StorageLocation("src", "run-1/a.tar.gz"), StorageLocation("dst", "v3/a.tar.gz"), ), ( StorageLocation("src", "run-1/b.tar.gz"), StorageLocation("dst", "v3/b.tar.gz"), ), ( StorageLocation("src", "run-1/c.tar.gz"), StorageLocation("dst", "v3/c.tar.gz"), ), ] count = backend.copy_files(files) self.assertEqual(count, 3) self.assertEqual(mock_client.copy.call_count, 3) def test_empty_list_returns_zero(self): backend = S3StorageBackend() mock_client = mock.MagicMock() backend._s3_client = mock_client count = backend.copy_files([]) self.assertEqual(count, 0) mock_client.copy.assert_not_called() def test_single_file_skips_thread_pool(self): backend = S3StorageBackend() mock_client = mock.MagicMock() backend._s3_client = mock_client files = [ ( StorageLocation("src", "run-1/a.tar.gz"), StorageLocation("dst", "v3/a.tar.gz"), ), ] with mock.patch( "_therock_utils.storage_backend.concurrent.futures.ThreadPoolExecutor" ) as mock_pool: count = backend.copy_files(files) self.assertEqual(count, 1) mock_pool.assert_not_called() mock_client.copy.assert_called_once() def test_dry_run_does_not_call_boto3(self): backend = S3StorageBackend(dry_run=True) mock_client = mock.MagicMock() backend._s3_client = mock_client files = [ ( StorageLocation("src", "run-1/a.tar.gz"), StorageLocation("dst", "v3/a.tar.gz"), ), ( StorageLocation("src", "run-1/b.tar.gz"), StorageLocation("dst", "v3/b.tar.gz"), ), ] count = backend.copy_files(files) self.assertEqual(count, 2) mock_client.copy.assert_not_called() def test_error_propagation(self): backend = S3StorageBackend() mock_client = mock.MagicMock() def copy_side_effect(copy_source, bucket, key): if "bad" in key: raise Exception("persistent failure") mock_client.copy.side_effect = copy_side_effect backend._s3_client = mock_client files = [ ( StorageLocation("src", "run-1/good.tar.gz"), StorageLocation("dst", "v3/good.tar.gz"), ), ( StorageLocation("src", "run-1/bad.tar.gz"), StorageLocation("dst", "v3/bad.tar.gz"), ), ] with mock.patch("_therock_utils.storage_backend.time.sleep"): with self.assertRaises(RuntimeError) as ctx: backend.copy_files(files) self.assertIn("bad.tar.gz", str(ctx.exception)) # --------------------------------------------------------------------------- # Factory # --------------------------------------------------------------------------- class TestCreateStorageBackend(unittest.TestCase): def test_returns_s3_backend_by_default(self): backend = create_storage_backend() self.assertIsInstance(backend, S3StorageBackend) def test_returns_local_backend_with_staging_dir(self): backend = create_storage_backend(staging_dir=Path("/tmp/staging")) self.assertIsInstance(backend, LocalStorageBackend) def test_dry_run_passed_through(self): backend = create_storage_backend(dry_run=True) self.assertIsInstance(backend, S3StorageBackend) self.assertTrue(backend._dry_run) def test_local_dry_run_passed_through(self): backend = create_storage_backend(staging_dir=Path("/tmp/staging"), dry_run=True) self.assertIsInstance(backend, LocalStorageBackend) self.assertTrue(backend._dry_run) def test_upload_concurrency_passed_to_s3_backend(self): backend = create_storage_backend(upload_concurrency=25) self.assertIsInstance(backend, S3StorageBackend) self.assertEqual(backend._upload_concurrency, 25) def test_upload_concurrency_ignored_for_local(self): backend = create_storage_backend( staging_dir=Path("/tmp/staging"), upload_concurrency=25 ) self.assertIsInstance(backend, LocalStorageBackend) if __name__ == "__main__": unittest.main()