def test_list_prefixes(self):
store = S3Service()
b = store.get_bucket("bucket")
b.create()
b.put_object(Key="a", Body=b"a")
b.put_object(Key="dir/b", Body=b"b")
assert store.list_prefixes(bucket_name="bucket",
prefix="non-existent/") == []
assert store.list_prefixes(bucket_name="bucket",
delimiter="/") == ["dir"]
assert store.list_keys(bucket_name="bucket",
delimiter="/") == [("a", 1)]
assert store.list_keys(bucket_name="bucket",
prefix="dir/") == [("b", 1)]
def test_check_prefix_format(self):
assert S3Service.check_prefix_format(prefix="foo", delimiter="") == "foo"
assert S3Service.check_prefix_format(prefix="foo", delimiter="/") == "foo/"
assert S3Service.check_prefix_format(prefix="foo/", delimiter="/") == "foo/"
assert S3Service.check_prefix_format(prefix="/foo/", delimiter="/") == "/foo/"
assert (
S3Service.check_prefix_format(prefix="/foo/boo", delimiter="/")
== "/foo/boo/"
)
assert S3Service.check_prefix_format(prefix="", delimiter="/") == ""
def test_list_keys_paged(self):
store = S3Service()
b = store.get_bucket("bucket")
b.create()
keys = ["x", "y"]
for key in keys:
b.put_object(Key=key, Body=b"a")
self.assertListEqual(
sorted([(k, 1) for k in keys]),
sorted(
store.list_keys(bucket_name="bucket",
delimiter="/",
page_size=2)),
)
def test_list_prefixes_paged(self):
store = S3Service()
b = store.get_bucket("bucket")
b.create()
# Test only one page
keys = ["x/b", "y/b"]
dirs = ["x", "y"]
for key in keys:
b.put_object(Key=key, Body=b"a")
self.assertListEqual(
sorted(dirs),
sorted(
store.list_prefixes(bucket_name="bucket", delimiter="/", page_size=1)
),
)
def test_upload_download(self):
store = S3Service()
store.connection.create_bucket(Bucket="bucket")
dirname = tempfile.mkdtemp()
fpath1 = dirname + "/test1.txt"
with open(fpath1, "w") as f:
f.write("data1")
fpath2 = dirname + "/test2.txt"
with open(fpath2, "w") as f:
f.write("data2")
fpath3 = dirname + "/test3.txt"
with open(fpath3, "w") as f:
f.write("data3")
# Upload
store.upload_file(fpath1, "my_key1.txt", "bucket", use_basename=False)
assert store.check_key("my_key1.txt", "bucket") is True
store.upload_file(fpath2, "my_key2.txt", "bucket", use_basename=False)
assert store.check_key("my_key2.txt", "bucket") is True
store.upload_file(fpath3, "foo/", "bucket", use_basename=True)
assert store.check_key("foo/test3.txt", "bucket") is True
# Download
store.download_file(
"my_key1.txt",
local_path=dirname + "/foo1.txt",
bucket_name="bucket",
use_basename=False,
)
assert os.path.basename(dirname + "/foo1.txt") == "foo1.txt"
assert open(os.path.join(dirname + "/foo1.txt")).read() == "data1"
dirname2 = tempfile.mkdtemp()
store.download_file(
"foo/test3.txt",
local_path=dirname2,
bucket_name="bucket",
use_basename=True,
)
assert os.path.basename(dirname2 + "/test3.txt") == "test3.txt"
assert open(os.path.join(dirname2 + "/test3.txt")).read() == "data3"
def test_ls(self):
store = S3Service()
b = store.get_bucket("bucket")
b.create()
b.put_object(Key="a", Body=b"a")
b.put_object(Key="dir/b", Body=b"b")
full_response = {"files": [("a", 1)], "dirs": ["dir"]}
empty_response = {"dirs": [], "files": []}
dir_response = {"dirs": [], "files": [("b", 1)]}
assert store.ls("s3://bucket") == full_response
assert store.ls("s3://bucket/") == full_response
assert store.ls("s3://bucket/non-existent") == empty_response
assert store.ls("s3://bucket/non-existent/") == empty_response
assert store.ls("s3://bucket/dir") == dir_response
assert store.ls("s3://bucket/dir/") == dir_response
def get_connection_from_type(connection_type: V1ConnectionType, **kwargs):
# We assume that `None` refers to local store as well
if not connection_type or connection_type.kind in {
V1ConnectionKind.VOLUME_CLAIM,
V1ConnectionKind.HOST_PATH,
}:
from polyaxon.stores.local_store import LocalStore
return LocalStore()
if connection_type.kind == V1ConnectionKind.WASB:
from polyaxon.connections.azure.azure_blobstore import AzureBlobStoreService
return AzureBlobStoreService(connection_name=connection_type.name,
**kwargs)
if connection_type.kind == V1ConnectionKind.S3:
from polyaxon.connections.aws.s3 import S3Service
return S3Service(connection_name=connection_type.name, **kwargs)
if connection_type.kind == V1ConnectionKind.GCS:
from polyaxon.connections.gcp.gcs import GCSService
return GCSService(connection_name=connection_type.name, **kwargs)
def test_get_bucket(self):
store = S3Service()
bucket = store.get_bucket("bucket")
assert bucket is not None
def test_check_bucket_raises_with_invalid_client_resources(self):
store = S3Service(connection="foo", resource="bar")
with self.assertRaises(Exception):
store.check_bucket("bucket")
def test_parse_s3_url(self):
s3_url = "s3://test/this/is/valid/key.txt"
parsed_url = S3Service.parse_s3_url(s3_url)
assert parsed_url == ("test", "this/is/valid/key.txt")
def test_read_key(self):
store = S3Service()
store.connection.create_bucket(Bucket="bucket")
store.connection.put_object(Bucket="bucket", Key="my_key", Body=b"M\xC3\xA9nar")
self.assertEqual(store.read_key("my_key", "bucket"), u"Ménar")
