ART/tests/unit/test_s3_utils.py at d27fcb3c5e13ebeccc1ed07cbb360b844a69973c · OpenPipe/ART · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
import asyncio
import importlib
from pathlib import Path
import zipfile

import pytest

s3_utils = importlib.import_module("art.utils.s3")


class _FakeProcess:
    def __init__(
        self,
        *,
        returncode: int = 0,
        stdout: bytes = b"",
        stderr: bytes = b"",
    ) -> None:
        self.returncode = returncode
        self._stdout = stdout
        self._stderr = stderr

    async def communicate(self) -> tuple[bytes, bytes]:
        return self._stdout, self._stderr


@pytest.mark.asyncio
async def test_archive_and_presign_step_url_uses_full_s3_uri(
    monkeypatch: pytest.MonkeyPatch, tmp_path: Path
) -> None:
    checkpoint_path = tmp_path / "checkpoint"
    checkpoint_path.mkdir()
    (checkpoint_path / "adapter_config.json").write_text('{"r": 8}', encoding="utf-8")

    upload_calls: list[tuple[str, str, bool, bool]] = []
    presign_calls: list[tuple[tuple[str, ...], object | None, object | None]] = []
    bucket_calls: list[str | None] = []

    async def fake_ensure_bucket_exists(
        s3_bucket: str | None = None, profile: str | None = None
    ) -> None:
        assert profile is None
        bucket_calls.append(s3_bucket)

    async def fake_s3_sync(
        source: str,
        destination: str,
        *,
        profile: str | None = None,
        verbose: bool = False,
        delete: bool = False,
        exclude: list[s3_utils.ExcludableOption] | None = None,
    ) -> None:
        assert profile is None
        assert exclude is None
        with zipfile.ZipFile(source) as archive:
            assert archive.namelist() == ["adapter_config.json"]
        upload_calls.append((source, destination, verbose, delete))

    async def fake_create_subprocess_exec(*cmd: str, stdout=None, stderr=None):
        presign_calls.append((cmd, stdout, stderr))
        return _FakeProcess(stdout=b"https://signed.example.com/model.zip\n")

    monkeypatch.setattr(s3_utils, "ensure_bucket_exists", fake_ensure_bucket_exists)
    monkeypatch.setattr(s3_utils, "s3_sync", fake_s3_sync)
    monkeypatch.setattr(
        s3_utils.asyncio, "create_subprocess_exec", fake_create_subprocess_exec
    )

    presigned_url = await s3_utils.archive_and_presign_step_url(
        model_name="demo-model",
        project="demo-project",
        step=7,
        s3_bucket="demo-bucket",
        prefix="exports",
        checkpoint_path=str(checkpoint_path),
    )

    expected_s3_uri = s3_utils.build_s3_zipped_step_path(
        model_name="demo-model",
        project="demo-project",
        step=7,
        s3_bucket="demo-bucket",
        prefix="exports",
    )

    assert presigned_url == "https://signed.example.com/model.zip"
    assert bucket_calls == ["demo-bucket"]
    assert len(upload_calls) == 1
    assert upload_calls[0][1:] == (expected_s3_uri, False, False)
    assert presign_calls == [
        (
            ("aws", "s3", "presign", expected_s3_uri, "--expires-in", "3600"),
            asyncio.subprocess.PIPE,
            asyncio.subprocess.PIPE,
        )
    ]


@pytest.mark.asyncio
async def test_archive_and_presign_step_url_surfaces_presign_failures(
    monkeypatch: pytest.MonkeyPatch, tmp_path: Path
) -> None:
    checkpoint_path = tmp_path / "checkpoint"
    checkpoint_path.mkdir()
    (checkpoint_path / "adapter_model.bin").write_text("weights", encoding="utf-8")

    async def fake_ensure_bucket_exists(
        s3_bucket: str | None = None, profile: str | None = None
    ) -> None:
        return None

    async def fake_s3_sync(
        source: str,
        destination: str,
        *,
        profile: str | None = None,
        verbose: bool = False,
        delete: bool = False,
        exclude: list[s3_utils.ExcludableOption] | None = None,
    ) -> None:
        return None

    async def fake_create_subprocess_exec(*cmd: str, stdout=None, stderr=None):
        return _FakeProcess(returncode=1, stderr=b"invalid S3 URI")

    monkeypatch.setattr(s3_utils, "ensure_bucket_exists", fake_ensure_bucket_exists)
    monkeypatch.setattr(s3_utils, "s3_sync", fake_s3_sync)
    monkeypatch.setattr(
        s3_utils.asyncio, "create_subprocess_exec", fake_create_subprocess_exec
    )

    with pytest.raises(RuntimeError, match="invalid S3 URI"):
        await s3_utils.archive_and_presign_step_url(
            model_name="demo-model",
            project="demo-project",
            step=3,
            s3_bucket="demo-bucket",
            checkpoint_path=str(checkpoint_path),
        )