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 )