mirror of
https://github.com/wassname/optuna-dashboard.git
synced 2026-08-21 11:18:45 +08:00
64 lines
2.3 KiB
Python
64 lines
2.3 KiB
Python
from __future__ import annotations
|
|
|
|
import io
|
|
from unittest import TestCase
|
|
|
|
import boto3
|
|
from optuna_dashboard.artifact.boto3 import Boto3Backend
|
|
from optuna_dashboard.artifact.exceptions import ArtifactNotFound
|
|
|
|
|
|
try:
|
|
from moto import mock_aws
|
|
except ImportError:
|
|
from moto import mock_s3 as mock_aws
|
|
|
|
|
|
@mock_aws
|
|
class Boto3BackendTestCase(TestCase):
|
|
def setUp(self) -> None:
|
|
self.s3_client = boto3.client("s3")
|
|
self.bucket_name = "moto-bucket"
|
|
self.s3_client.create_bucket(Bucket=self.bucket_name)
|
|
|
|
def tearDown(self) -> None:
|
|
objects = self.s3_client.list_objects(Bucket=self.bucket_name).get("Contents", [])
|
|
if objects:
|
|
self.s3_client.delete_objects(
|
|
Bucket=self.bucket_name,
|
|
Delete={"Objects": [{"Key": obj["Key"] for obj in objects}], "Quiet": True},
|
|
)
|
|
self.s3_client.delete_bucket(Bucket=self.bucket_name)
|
|
|
|
def test_upload_download(self) -> None:
|
|
artifact_id = "dummy-uuid"
|
|
dummy_content = b"Hello World"
|
|
buf = io.BytesIO(dummy_content)
|
|
|
|
backend = Boto3Backend(self.bucket_name)
|
|
backend.write(artifact_id, buf)
|
|
self.assertEqual(len(self.s3_client.list_objects(Bucket=self.bucket_name)["Contents"]), 1)
|
|
obj = self.s3_client.get_object(Bucket=self.bucket_name, Key=artifact_id)
|
|
self.assertEqual(obj["Body"].read(), dummy_content)
|
|
|
|
with backend.open(artifact_id) as f:
|
|
actual = f.read()
|
|
self.assertEqual(actual, dummy_content)
|
|
self.assertFalse(buf.closed)
|
|
|
|
def test_remove(self) -> None:
|
|
artifact_id = "dummy-uuid"
|
|
backend = Boto3Backend(self.bucket_name)
|
|
backend.write(artifact_id, io.BytesIO(b"Hello"))
|
|
objects = self.s3_client.list_objects(Bucket=self.bucket_name)["Contents"]
|
|
self.assertEqual(len([obj for obj in objects if obj["Key"] == artifact_id]), 1)
|
|
|
|
backend.remove(artifact_id)
|
|
objects = self.s3_client.list_objects(Bucket=self.bucket_name).get("Contents", [])
|
|
self.assertEqual(len([obj for obj in objects if obj["Key"] == artifact_id]), 0)
|
|
|
|
def test_file_not_found_exception(self) -> None:
|
|
backend = Boto3Backend(self.bucket_name)
|
|
with self.assertRaises(ArtifactNotFound):
|
|
backend.open("not-found-id")
|