mirror of
https://github.com/Stone-Red-Code/dvc.git
synced 2026-11-07 23:09:21 +01:00
111 lines
2.6 KiB
Python
111 lines
2.6 KiB
Python
import locale
|
|
import os
|
|
import uuid
|
|
|
|
import pytest
|
|
from funcy import cached_property
|
|
from moto import mock_s3
|
|
|
|
from dvc.path_info import CloudURLInfo
|
|
from dvc.utils import env2bool
|
|
|
|
from .base import Base
|
|
|
|
TEST_AWS_REPO_BUCKET = os.environ.get("DVC_TEST_AWS_REPO_BUCKET", "dvc-temp")
|
|
|
|
|
|
class S3(Base, CloudURLInfo):
|
|
@staticmethod
|
|
def should_test():
|
|
do_test = env2bool("DVC_TEST_AWS", undefined=None)
|
|
if do_test is not None:
|
|
return do_test
|
|
|
|
if os.getenv("AWS_ACCESS_KEY_ID") and os.getenv(
|
|
"AWS_SECRET_ACCESS_KEY"
|
|
):
|
|
return True
|
|
|
|
return False
|
|
|
|
@staticmethod
|
|
def _get_storagepath():
|
|
return (
|
|
TEST_AWS_REPO_BUCKET
|
|
+ "/"
|
|
+ "dvc_test_caches"
|
|
+ "/"
|
|
+ str(uuid.uuid4())
|
|
)
|
|
|
|
@staticmethod
|
|
def get_url():
|
|
return "s3://" + S3._get_storagepath()
|
|
|
|
@cached_property
|
|
def _s3(self):
|
|
import boto3
|
|
|
|
return boto3.client("s3")
|
|
|
|
def is_file(self):
|
|
from botocore.exceptions import ClientError
|
|
|
|
if self.path.endswith("/"):
|
|
return False
|
|
|
|
try:
|
|
self._s3.head_object(Bucket=self.bucket, Key=self.path)
|
|
except ClientError as exc:
|
|
if exc.response["Error"]["Code"] != "404":
|
|
raise
|
|
return False
|
|
|
|
return True
|
|
|
|
def is_dir(self):
|
|
path = (self / "").path
|
|
resp = self._s3.list_objects(Bucket=self.bucket, Prefix=path)
|
|
return bool(resp.get("Contents"))
|
|
|
|
def exists(self):
|
|
return self.is_file() or self.is_dir()
|
|
|
|
def mkdir(self, mode=0o777, parents=False, exist_ok=False):
|
|
assert mode == 0o777
|
|
assert parents
|
|
|
|
def write_bytes(self, contents):
|
|
self._s3.put_object(
|
|
Bucket=self.bucket, Key=self.path, Body=contents,
|
|
)
|
|
|
|
def read_bytes(self):
|
|
data = self._s3.get_object(Bucket=self.bucket, Key=self.path)
|
|
return data["Body"].read()
|
|
|
|
def read_text(self, encoding=None, errors=None):
|
|
if not encoding:
|
|
encoding = locale.getpreferredencoding(False)
|
|
assert errors is None
|
|
return self.read_bytes().decode(encoding)
|
|
|
|
|
|
@pytest.fixture
|
|
def s3(test_config):
|
|
test_config.requires("s3")
|
|
with mock_s3():
|
|
import boto3
|
|
|
|
boto3.client("s3").create_bucket(Bucket=TEST_AWS_REPO_BUCKET)
|
|
|
|
yield S3(S3.get_url())
|
|
|
|
|
|
@pytest.fixture
|
|
def real_s3(test_config):
|
|
test_config.requires("s3")
|
|
if not S3.should_test():
|
|
pytest.skip("no real s3")
|
|
yield S3(S3.get_url())
|