optuna

1
from __future__ import annotations
2

3
import io
4
from typing import TYPE_CHECKING
5

6
import boto3
7
import pytest
8

9
from optuna.artifacts import Boto3ArtifactStore
10
from optuna.artifacts.exceptions import ArtifactNotFound
11

12

13
try:
14
    # TODO(nabenabe0928): Replace it with `from moto import mock_aws` after dropping Python3.7.
15
    from moto import mock_aws
16
except ImportError:
17
    from moto import mock_s3 as mock_aws  # type: ignore[attr-defined,no-redef]
18

19
if TYPE_CHECKING:
20
    from collections.abc import Iterator
21

22
    from mypy_boto3_s3 import S3Client
23
    from typing_extensions import Annotated
24

25
    # TODO(Shinichi) import Annotated from typing after python 3.8 support is dropped.
26

27

28
@pytest.fixture()
29
def init_mock_client() -> Iterator[tuple[str, S3Client]]:
30
    with mock_aws():
31
        # Runs before each test
32
        bucket_name = "moto-bucket"
33
        s3_client = boto3.client("s3")
34
        s3_client.create_bucket(Bucket=bucket_name)
35

36
        yield bucket_name, s3_client
37

38
        # Runs after each test
39
        objects = s3_client.list_objects(Bucket=bucket_name).get("Contents", [])
40
        if objects:
41
            s3_client.delete_objects(
42
                Bucket=bucket_name,
43
                Delete={"Objects": [{"Key": obj["Key"] for obj in objects}], "Quiet": True},
44
            )
45
        s3_client.delete_bucket(Bucket=bucket_name)
46

47

48
@pytest.mark.parametrize("avoid_buf_copy", [True, False])
49
def test_upload_download(
50
    init_mock_client: Annotated[tuple[str, S3Client], pytest.fixture],
51
    avoid_buf_copy: bool,
52
) -> None:
53
    bucket_name, s3_client = init_mock_client
54
    backend = Boto3ArtifactStore(bucket_name, avoid_buf_copy=avoid_buf_copy)
55

56
    artifact_id = "dummy-uuid"
57
    dummy_content = b"Hello World"
58
    buf = io.BytesIO(dummy_content)
59

60
    backend.write(artifact_id, buf)
61
    assert len(s3_client.list_objects(Bucket=bucket_name)["Contents"]) == 1
62

63
    obj = s3_client.get_object(Bucket=bucket_name, Key=artifact_id)
64
    assert obj["Body"].read() == dummy_content
65

66
    with backend.open_reader(artifact_id) as f:
67
        actual = f.read()
68
    assert actual == dummy_content
69
    if avoid_buf_copy is False:
70
        assert buf.closed is False
71

72

73
def test_remove(init_mock_client: Annotated[tuple[str, S3Client], pytest.fixture]) -> None:
74
    bucket_name, s3_client = init_mock_client
75
    backend = Boto3ArtifactStore(bucket_name)
76

77
    artifact_id = "dummy-uuid"
78
    backend.write(artifact_id, io.BytesIO(b"Hello"))
79
    objects = s3_client.list_objects(Bucket=bucket_name)["Contents"]
80
    assert len([obj for obj in objects if obj["Key"] == artifact_id]) == 1
81

82
    backend.remove(artifact_id)
83
    objects = s3_client.list_objects(Bucket=bucket_name).get("Contents", [])
84
    assert len([obj for obj in objects if obj["Key"] == artifact_id]) == 0
85

86

87
def test_file_not_found_exception(
88
    init_mock_client: Annotated[tuple[str, S3Client], pytest.fixture]
89
) -> None:
90
    bucket_name, _ = init_mock_client
91
    backend = Boto3ArtifactStore(bucket_name)
92

93
    with pytest.raises(ArtifactNotFound):
94
        backend.open_reader("not-found-id")
95