diff --git a/python_tests/test_boto3_artifact.py b/python_tests/test_boto3_artifact.py index cafb9f91..0030c07a 100644 --- a/python_tests/test_boto3_artifact.py +++ b/python_tests/test_boto3_artifact.py @@ -9,23 +9,28 @@ from optuna_dashboard.artifact.boto3 import Boto3Backend @mock_s3 class Boto3BackendTestCase(TestCase): def setUp(self) -> None: - self.s3_client = boto3.resource("s3") - self.bucket = self.s3_client.create_bucket(Bucket="moto-bucket") + self.s3_client = boto3.client("s3") + self.bucket_name = "moto-bucket" + self.s3_client.create_bucket(Bucket=self.bucket_name) def tearDown(self) -> None: - self.bucket.objects.all().delete() - self.bucket.delete() + 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" - backend = Boto3Backend(self.bucket.name) + backend = Boto3Backend(self.bucket_name) backend.write(artifact_id, io.BytesIO(dummy_content)) - - objects = [obj for obj in self.bucket.objects.all() if obj.key == artifact_id] - assert len(objects) == 1 - assert objects[0].get()["Body"].read() == dummy_content + assert 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) + assert obj["Body"].read() == dummy_content with backend.open(artifact_id) as f: actual = f.read() @@ -33,9 +38,11 @@ class Boto3BackendTestCase(TestCase): def test_remove(self) -> None: artifact_id = "dummy-uuid" - backend = Boto3Backend(self.bucket.name) + backend = Boto3Backend(self.bucket_name) backend.write(artifact_id, io.BytesIO(b"Hello")) - assert len([obj for obj in self.bucket.objects.all() if obj.key == artifact_id]) == 1 + objects = self.s3_client.list_objects(Bucket=self.bucket_name)['Contents'] + assert len([obj for obj in objects if obj["Key"] == artifact_id]) == 1 backend.remove(artifact_id) - assert len([obj for obj in self.bucket.objects.all() if obj.key == artifact_id]) == 0 + objects = self.s3_client.list_objects(Bucket=self.bucket_name).get('Contents', []) + assert len([obj for obj in objects if obj["Key"] == artifact_id]) == 0