def test_write(self):
"""
Test that we can write to the local file system.
:return:
"""
driver = FileStorageDriver('')
filepath = 'tests/test-output/FileStorageDriverWriteTest.txt'
# Make sure it doesn't exist already
if os.path.exists(filepath):
os.remove(filepath)
driver.get_filename = MagicMock(return_value=filepath)
data = 'this is the data to write.'
driver.write('the_dag', 'the_task', datetime(1983, 9, 5), data)
self.assertTrue(os.path.exists(filepath))
# Make sure we wrote the correct thing
with open(filepath, 'r') as f:
read_data = f.readline()
self.assertEqual(read_data, data)
read_data = open(filepath, 'r').readline()
# Clean up.
os.remove(filepath)
def test_get_read_stream_simple(self):
"""
Test reading a stream from the file system
"""
fake_read_stream_args = ['1', '2', datetime(2016, 1, 1)]
driver = FileStorageDriver('.')
driver.get_filename = MagicMock(return_value='tests/fixtures/SampleUniformData.json')
string_value = open('tests/fixtures/SampleUniformData.json', 'r').read()
actual_result = driver.get_read_stream(*fake_read_stream_args)
self.assertEqual(actual_result.read().decode('utf-8'), string_value)
def test_read(self):
"""
Test that we can read a fixture file on the local file system using an
encoding..
"""
filepath = 'tests/fixtures/FileStorageDriverReadTest.txt'
expected = 'this is a reading test.'
driver = FileStorageDriver('')
driver.get_filename = MagicMock(return_value=filepath)
actual = driver.read('the_dag', 'the_task', datetime(1983, 9, 5), 'utf-8')
self.assertEqual(actual, expected)
def test_get_filename(self):
"""
Test that get_filename returns the expected value.
"""
expected = '/this/is/a/test/the_dag/the_task/1983-09-05'
driver = FileStorageDriver('/this/is/a/test')
actual = driver.get_filename(
'the_dag',
'the_task',
datetime(1983, 9, 5)
)
self.assertEqual(actual, expected)
def test_write_from_stream(self):
"""
Test writing to the file system from a stream
"""
driver = FileStorageDriver('.')
filepath = 'tests/test-output/test_write_from_stream.txt'
# Make sure it doesn't exist already
if os.path.exists(filepath):
os.remove(filepath)
driver.get_filename = MagicMock(return_value=filepath)
f = open('tests/fixtures/SampleUniformData.json', 'r')
driver.write_from_stream('the_dag', 'the_task', datetime(1986, 4, 30), f)
# Go back to the start so we can read it in
f.seek(0)
file_data = f.read()
actual_data = open(filepath, 'r').read()
self.assertEqual(actual_data, file_data)