def test_s3_data_write_key(mocker):
mocker.patch.object(utils.config, "S3_IMPORT_DATA_BUCKET", "test-bucket")
mock_load = mocker.patch('dataflow.utils.S3Hook.load_string')
s3 = utils.S3Data('table', '20010101')
s3.write_key("key.json", {"data": "value"})
mock_load.assert_called_once_with(
'{"data": "value"}',
'import-data/table/20010101/key.json',
bucket_name="test-bucket",
encrypt=True,
replace=True,
)
def test_s3_data_list_keys_retries_requests(mocker):
mocker.patch("time.sleep") # skip backoff retry delay
mocker.patch.object(utils.config, "S3_IMPORT_DATA_BUCKET", "test-bucket")
mocker.patch(
'dataflow.utils.S3Hook.list_keys',
side_effect=[
botocore.exceptions.EndpointConnectionError(endpoint_url='aws'),
[],
],
)
s3 = utils.S3Data('table', '20010101')
assert s3.list_keys() == []
def test_s3_data_write_key_handles_decimals(mocker):
mocker.patch.object(utils.config, "S3_IMPORT_DATA_BUCKET", "test-bucket")
mock_load = mocker.patch('dataflow.utils.S3Hook.load_string')
s3 = utils.S3Data('table', '20020202')
with freezegun.freeze_time('2020-01-01T12:00:00Z'):
s3.write_key("key.json", {'a_decimal': decimal.Decimal('1.11')})
mock_load.assert_called_once_with(
'{"a_decimal": "1.11"}',
'import-data/table/20020202/key.json',
bucket_name="test-bucket",
encrypt=True,
replace=True,
)
def test_s3_data_read_key_retries_requests(mocker):
mocker.patch("time.sleep") # skip backoff retry delay
mocker.patch.object(utils.config, "S3_IMPORT_DATA_BUCKET", "test-bucket")
mocker.patch(
'dataflow.utils.S3Hook.get_key',
side_effect=[
botocore.exceptions.EndpointConnectionError(endpoint_url='aws'),
mocker.Mock(get=mocker.Mock(
side_effect=lambda:
{'Body': BytesIO(b'[{"key":"val_1"},{"key":"val_2"}]')})),
],
)
s3 = utils.S3Data('table', '20010101')
assert list(s3.read_key("key")) == [{"key": "val_1"}, {"key": "val_2"}]
def test_s3_data_write_key_handles_dates_and_times(mocker):
mocker.patch.object(utils.config, "S3_IMPORT_DATA_BUCKET", "test-bucket")
mock_load = mocker.patch('dataflow.utils.S3Hook.load_string')
s3 = utils.S3Data('table', '20010101')
with freezegun.freeze_time('2020-01-01T12:00:00Z'):
s3.write_key("key.json", {
"date": date.today(),
"datetime": datetime.now()
})
mock_load.assert_called_once_with(
'{"date": "2020-01-01", "datetime": "2020-01-01T12:00:00"}',
'import-data/table/20010101/key.json',
bucket_name="test-bucket",
encrypt=True,
replace=True,
)
def test_s3_data_read_keys(mocker):
mocker.patch.object(utils.config, "S3_IMPORT_DATA_BUCKET", "test-bucket")
mock_read = mocker.patch(
'dataflow.utils.S3Hook.get_key',
return_value=mocker.Mock(get=mocker.Mock(
return_value={
'Body': BytesIO(b'[{"key":"val_1"},{"key":"val_2"}]')
})),
)
s3 = utils.S3Data('table', '20010101')
assert list(s3.read_key("prefix/key.json")) == [{
"key": "val_1"
}, {
"key": "val_2"
}]
mock_read.assert_called_once_with("prefix/key.json",
bucket_name="test-bucket")
def test_s3_data_iter_keys(mocker):
mocker.patch.object(utils.config, "S3_IMPORT_DATA_BUCKET", "test-bucket")
mock_list = mocker.patch('dataflow.utils.S3Hook.list_keys',
return_value=iter(["key1", "key2"]))
mocker.patch(
'dataflow.utils.S3Hook.get_key',
return_value=mocker.Mock(get=mocker.Mock(
side_effect=lambda:
{'Body': BytesIO(b'[{"key":"val_1"},{"key":"val_2"}]')})),
)
s3 = utils.S3Data('table', '20010101')
records = [
key_record for (key, key_records) in s3.iter_keys()
for key_record in key_records
]
assert records == [{"key": "val_1"}, {"key": "val_2"}] * 2
mock_list.assert_called_once_with(bucket_name="test-bucket",
prefix='import-data/table/20010101/')
def test_s3_data(mocker):
s3 = utils.S3Data('table', '20010101')
assert s3.prefix == 'import-data/table/20010101/'