sdk/python/tests/test_stream.py

   1 #!/usr/bin/env python
   2
   3 import bz2
   4 import gzip
   5 import io
   6 import mock
   7 import os
   8 import unittest
   9 import hashlib
  10
  11 import arvados
  12 from arvados import StreamReader, StreamFileReader
  13 from arvados._ranges import Range
  14
  15 import arvados_testutil as tutil
  16 import run_test_server
  17
  18 class StreamFileReaderTestCase(unittest.TestCase):
  19     def make_count_reader(self):
  20         stream = tutil.MockStreamReader('.', '01234', '34567', '67890')
  21         return StreamFileReader(stream, [Range(1, 0, 3), Range(6, 3, 3), Range(11, 6, 3)],
  22                                 'count.txt')
  23
  24     def test_read_returns_first_block(self):
  25         # read() calls will be aligned on block boundaries - see #3663.
  26         sfile = self.make_count_reader()
  27         self.assertEqual('123', sfile.read(10))
  28
  29     def test_small_read(self):
  30         sfile = self.make_count_reader()
  31         self.assertEqual('12', sfile.read(2))
  32
  33     def test_successive_reads(self):
  34         sfile = self.make_count_reader()
  35         for expect in ['123', '456', '789', '']:
  36             self.assertEqual(expect, sfile.read(10))
  37
  38     def test_readfrom_spans_blocks(self):
  39         sfile = self.make_count_reader()
  40         self.assertEqual('6789', sfile.readfrom(5, 12))
  41
  42     def test_small_readfrom_spanning_blocks(self):
  43         sfile = self.make_count_reader()
  44         self.assertEqual('2345', sfile.readfrom(1, 4))
  45
  46     def test_readall(self):
  47         sfile = self.make_count_reader()
  48         self.assertEqual('123456789', ''.join(sfile.readall()))
  49
  50     def test_one_arg_seek(self):
  51         self.test_absolute_seek([])
  52
  53     def test_absolute_seek(self, args=[os.SEEK_SET]):
  54         sfile = self.make_count_reader()
  55         sfile.seek(6, *args)
  56         self.assertEqual('78', sfile.read(2))
  57         sfile.seek(4, *args)
  58         self.assertEqual('56', sfile.read(2))
  59
  60     def test_relative_seek(self, args=[os.SEEK_CUR]):
  61         sfile = self.make_count_reader()
  62         self.assertEqual('12', sfile.read(2))
  63         sfile.seek(2, *args)
  64         self.assertEqual('56', sfile.read(2))
  65
  66     def test_end_seek(self):
  67         sfile = self.make_count_reader()
  68         sfile.seek(-6, os.SEEK_END)
  69         self.assertEqual('45', sfile.read(2))
  70
  71     def test_seek_min_zero(self):
  72         sfile = self.make_count_reader()
  73         sfile.seek(-2, os.SEEK_SET)
  74         self.assertEqual(0, sfile.tell())
  75
  76     def test_seek_max_size(self):
  77         sfile = self.make_count_reader()
  78         sfile.seek(2, os.SEEK_END)
  79         self.assertEqual(9, sfile.tell())
  80
  81     def test_size(self):
  82         self.assertEqual(9, self.make_count_reader().size())
  83
  84     def test_tell_after_block_read(self):
  85         sfile = self.make_count_reader()
  86         sfile.read(5)
  87         self.assertEqual(3, sfile.tell())
  88
  89     def test_tell_after_small_read(self):
  90         sfile = self.make_count_reader()
  91         sfile.read(1)
  92         self.assertEqual(1, sfile.tell())
  93
  94     def test_no_read_after_close(self):
  95         sfile = self.make_count_reader()
  96         sfile.close()
  97         self.assertRaises(ValueError, sfile.read, 2)
  98
  99     def test_context(self):
 100         with self.make_count_reader() as sfile:
 101             self.assertFalse(sfile.closed, "reader is closed inside context")
 102             self.assertEqual('12', sfile.read(2))
 103         self.assertTrue(sfile.closed, "reader is open after context")
 104
 105     def make_newlines_reader(self):
 106         stream = tutil.MockStreamReader('.', 'one\ntwo\n\nth', 'ree\nfour\n\n')
 107         return StreamFileReader(stream, [Range(0, 0, 11), Range(11, 11, 10)], 'count.txt')
 108
 109     def check_lines(self, actual):
 110         self.assertEqual(['one\n', 'two\n', '\n', 'three\n', 'four\n', '\n'],
 111                          actual)
 112
 113     def test_readline(self):
 114         reader = self.make_newlines_reader()
 115         actual = []
 116         while True:
 117             data = reader.readline()
 118             if not data:
 119                 break
 120             actual.append(data)
 121         self.check_lines(actual)
 122
 123     def test_readlines(self):
 124         self.check_lines(self.make_newlines_reader().readlines())
 125
 126     def test_iteration(self):
 127         self.check_lines(list(iter(self.make_newlines_reader())))
 128
 129     def test_readline_size(self):
 130         reader = self.make_newlines_reader()
 131         self.assertEqual('on', reader.readline(2))
 132         self.assertEqual('e\n', reader.readline(4))
 133         self.assertEqual('two\n', reader.readline(6))
 134         self.assertEqual('\n', reader.readline(8))
 135         self.assertEqual('thre', reader.readline(4))
 136
 137     def test_readlines_sizehint(self):
 138         result = self.make_newlines_reader().readlines(8)
 139         self.assertEqual(['one\n', 'two\n'], result[:2])
 140         self.assertNotIn('three\n', result)
 141
 142     def test_name_attribute(self):
 143         # Test both .name and .name() (for backward compatibility)
 144         stream = tutil.MockStreamReader()
 145         sfile = StreamFileReader(stream, [Range(0, 0, 0)], 'nametest')
 146         self.assertEqual('nametest', sfile.name)
 147         self.assertEqual('nametest', sfile.name())
 148
 149     def check_decompressed_name(self, filename, expect):
 150         stream = tutil.MockStreamReader('.', '')
 151         reader = StreamFileReader(stream, [Range(0, 0, 0)], filename)
 152         self.assertEqual(expect, reader.decompressed_name())
 153
 154     def test_decompressed_name_uncompressed_file(self):
 155         self.check_decompressed_name('test.log', 'test.log')
 156
 157     def test_decompressed_name_gzip_file(self):
 158         self.check_decompressed_name('test.log.gz', 'test.log')
 159
 160     def test_decompressed_name_bz2_file(self):
 161         self.check_decompressed_name('test.log.bz2', 'test.log')
 162
 163     def check_decompression(self, compress_ext, compress_func):
 164         test_text = 'decompression\ntest\n'
 165         test_data = compress_func(test_text)
 166         stream = tutil.MockStreamReader('.', test_data)
 167         reader = StreamFileReader(stream, [Range(0, 0, len(test_data))],
 168                                   'test.' + compress_ext)
 169         self.assertEqual(test_text, ''.join(reader.readall_decompressed()))
 170
 171     @staticmethod
 172     def gzip_compress(data):
 173         compressed_data = io.BytesIO()
 174         with gzip.GzipFile(fileobj=compressed_data, mode='wb') as gzip_file:
 175             gzip_file.write(data)
 176         return compressed_data.getvalue()
 177
 178     def test_no_decompression(self):
 179         self.check_decompression('log', lambda s: s)
 180
 181     def test_gzip_decompression(self):
 182         self.check_decompression('gz', self.gzip_compress)
 183
 184     def test_bz2_decompression(self):
 185         self.check_decompression('bz2', bz2.compress)
 186
 187
 188 class StreamRetryTestMixin(object):
 189     # Define reader_for(coll_name, **kwargs)
 190     # and read_for_test(reader, size, **kwargs).
 191     API_COLLECTIONS = run_test_server.fixture('collections')
 192
 193     def keep_client(self):
 194         return arvados.KeepClient(proxy='http://[%s]:1' % (tutil.TEST_HOST,),
 195                                   local_store='')
 196
 197     def manifest_for(self, coll_name):
 198         return self.API_COLLECTIONS[coll_name]['manifest_text']
 199
 200     @tutil.skip_sleep
 201     def test_success_without_retries(self):
 202         with tutil.mock_get_responses('bar', 200):
 203             reader = self.reader_for('bar_file')
 204             self.assertEqual('bar', self.read_for_test(reader, 3))
 205
 206     @tutil.skip_sleep
 207     def test_read_no_default_retry(self):
 208         with tutil.mock_get_responses('', 500):
 209             reader = self.reader_for('user_agreement')
 210             with self.assertRaises(arvados.errors.KeepReadError):
 211                 self.read_for_test(reader, 10)
 212
 213     @tutil.skip_sleep
 214     def test_read_with_instance_retries(self):
 215         with tutil.mock_get_responses('foo', 500, 200):
 216             reader = self.reader_for('foo_file', num_retries=3)
 217             self.assertEqual('foo', self.read_for_test(reader, 3))
 218
 219     @tutil.skip_sleep
 220     def test_read_with_method_retries(self):
 221         with tutil.mock_get_responses('foo', 500, 200):
 222             reader = self.reader_for('foo_file')
 223             self.assertEqual('foo',
 224                              self.read_for_test(reader, 3, num_retries=3))
 225
 226     @tutil.skip_sleep
 227     def test_read_instance_retries_exhausted(self):
 228         with tutil.mock_get_responses('bar', 500, 500, 500, 500, 200):
 229             reader = self.reader_for('bar_file', num_retries=3)
 230             with self.assertRaises(arvados.errors.KeepReadError):
 231                 self.read_for_test(reader, 3)
 232
 233     @tutil.skip_sleep
 234     def test_read_method_retries_exhausted(self):
 235         with tutil.mock_get_responses('bar', 500, 500, 500, 500, 200):
 236             reader = self.reader_for('bar_file')
 237             with self.assertRaises(arvados.errors.KeepReadError):
 238                 self.read_for_test(reader, 3, num_retries=3)
 239
 240     @tutil.skip_sleep
 241     def test_method_retries_take_precedence(self):
 242         with tutil.mock_get_responses('', 500, 500, 500, 200):
 243             reader = self.reader_for('user_agreement', num_retries=10)
 244             with self.assertRaises(arvados.errors.KeepReadError):
 245                 self.read_for_test(reader, 10, num_retries=1)
 246
 247
 248 class StreamReaderTestCase(unittest.TestCase, StreamRetryTestMixin):
 249     def reader_for(self, coll_name, **kwargs):
 250         return StreamReader(self.manifest_for(coll_name).split(),
 251                             self.keep_client(), **kwargs)
 252
 253     def read_for_test(self, reader, byte_count, **kwargs):
 254         return reader.readfrom(0, byte_count, **kwargs)
 255
 256     def test_manifest_text_without_keep_client(self):
 257         mtext = self.manifest_for('multilevel_collection_1')
 258         for line in mtext.rstrip('\n').split('\n'):
 259             reader = StreamReader(line.split())
 260             self.assertEqual(line + '\n', reader.manifest_text())
 261
 262
 263 class StreamFileReadTestCase(unittest.TestCase, StreamRetryTestMixin):
 264     def reader_for(self, coll_name, **kwargs):
 265         return StreamReader(self.manifest_for(coll_name).split(),
 266                             self.keep_client(), **kwargs).all_files()[0]
 267
 268     def read_for_test(self, reader, byte_count, **kwargs):
 269         return reader.read(byte_count, **kwargs)
 270
 271
 272 class StreamFileReadFromTestCase(StreamFileReadTestCase):
 273     def read_for_test(self, reader, byte_count, **kwargs):
 274         return reader.readfrom(0, byte_count, **kwargs)
 275
 276
 277 class StreamFileReadAllTestCase(StreamFileReadTestCase):
 278     def read_for_test(self, reader, byte_count, **kwargs):
 279         return ''.join(reader.readall(**kwargs))
 280
 281
 282 class StreamFileReadAllDecompressedTestCase(StreamFileReadTestCase):
 283     def read_for_test(self, reader, byte_count, **kwargs):
 284         return ''.join(reader.readall_decompressed(**kwargs))
 285
 286
 287 class StreamFileReadlinesTestCase(StreamFileReadTestCase):
 288     def read_for_test(self, reader, byte_count, **kwargs):
 289         return ''.join(reader.readlines(**kwargs))
 290
 291 if __name__ == '__main__':
 292     unittest.main()