def test_empty_credentials_load(self, bad_credentials):
s3_data_set = PickleS3DataSet(
filepath=FILENAME, bucket_name=BUCKET_NAME, credentials=bad_credentials
)
pattern = r"Failed while loading data from data set PickleS3DataSet\(.+\)"
with pytest.raises(DataSetError, match=pattern):
s3_data_set.load()
def versioned_s3_data_set(load_version, save_version):
return PickleS3DataSet(
filepath=FILENAME,
bucket_name=BUCKET_NAME,
credentials=AWS_CREDENTIALS,
version=Version(load_version, save_version),
)
def test_version_str_repr(self, load_version, save_version):
"""Test that version is in string representation of the class instance
when applicable."""
ds = PickleS3DataSet(filepath=FILENAME, bucket_name=BUCKET_NAME)
ds_versioned = PickleS3DataSet(
filepath=FILENAME,
bucket_name=BUCKET_NAME,
version=Version(load_version, save_version),
)
assert FILENAME in str(ds)
assert "version" not in str(ds)
assert FILENAME in str(ds_versioned)
ver_str = "version=Version(load={}, save='{}')".format(
load_version, save_version)
assert ver_str in str(ds_versioned)
def test_s3fs_args_propagated(self, mocker):
mock = mocker.patch("kedro.io.pickle_s3.S3FileSystem")
PickleS3DataSet(FILENAME,
BUCKET_NAME,
AWS_CREDENTIALS,
s3fs_args=dict(custom=42))
mock.assert_called_with(client_kwargs=mocker.ANY, custom=42)
def s3_data_set_with_args():
return PickleS3DataSet(
filepath=FILENAME,
bucket_name=BUCKET_NAME,
credentials=AWS_CREDENTIALS,
load_args={"fix_imports": False},
save_args={"fix_imports": False},
)
def test_incorrect_credentials_load(self):
"""Test that incorrect credential keys won't instantiate dataset."""
pattern = "unexpected keyword argument"
with pytest.raises(TypeError, match=pattern):
PickleS3DataSet(
filepath=FILENAME,
bucket_name=BUCKET_NAME,
credentials={"access_token": "TOKEN", "access_key": "KEY"},
)
def test_load_with_protocol(self):
"""Test loading the data from S3."""
s3_data_set = PickleS3DataSet(
filepath="s3://{}/{}".format(BUCKET_NAME, FILENAME),
credentials={
"aws_access_key_id": "YOUR_KEY",
"aws_secret_access_key": "YOUR SECRET",
},
)
loaded_data = s3_data_set.load()
assert loaded_data == DUMMY_PICKABLE_OBJECT
def test_save_with_protocol(self):
"""Test saving the data to S3."""
s3_data_set = PickleS3DataSet(
filepath="s3://{}/{}".format(BUCKET_NAME, FILENAME),
credentials={
"aws_access_key_id": "YOUR_KEY",
"aws_secret_access_key": "YOUR SECRET",
},
)
new_data = {"x": "y"}
s3_data_set.save(new_data)
loaded_data = s3_data_set.load()
assert loaded_data == new_data
def test_pass_credentials(self, mocker):
"""Test that AWS credentials are passed successfully into boto3
client instantiation on creating S3 connection."""
mocker.patch("s3fs.core.boto3.Session.client")
s3_data_set = PickleS3DataSet(
filepath=FILENAME, bucket_name=BUCKET_NAME, credentials=AWS_CREDENTIALS
)
pattern = r"Failed while loading data from data set PickleS3DataSet\(.+\)"
with pytest.raises(DataSetError, match=pattern):
s3_data_set.load()
assert s3fs.core.boto3.Session.client.call_count == 1
args, kwargs = s3fs.core.boto3.Session.client.call_args_list[0]
assert args == ("s3",)
for k, v in AWS_CREDENTIALS.items():
assert kwargs[k] == v
def s3_data_set():
return PickleS3DataSet(
filepath=FILENAME, bucket_name=BUCKET_NAME, credentials=AWS_CREDENTIALS
)
def test_incomplete_credentials_load(self, bad_credentials):
"""Test that incomplete credentials passed in credentials.yml raises exception."""
with pytest.raises(PartialCredentialsError):
PickleS3DataSet(
filepath=FILENAME, bucket_name=BUCKET_NAME, credentials=bad_credentials
)
