11308: Merge branch 'master' into 11308-python3
[arvados.git] / sdk / python / tests / test_stream.py
index 3970d672a6a46ac1ada71ca553bdbdd94834cd5e..063950bf684bd7f607d165bfd4b75f94ce735861 100644 (file)
-#!/usr/bin/env python
-
+from __future__ import absolute_import
+from builtins import object
+import bz2
+import gzip
+import io
 import mock
+import os
 import unittest
+import hashlib
 
 import arvados
 from arvados import StreamReader, StreamFileReader
+from arvados._ranges import Range
+
+from . import arvados_testutil as tutil
+from . import run_test_server
+
+class StreamFileReaderTestCase(unittest.TestCase):
+    def make_count_reader(self):
+        stream = tutil.MockStreamReader('.', '01234', '34567', '67890')
+        return StreamFileReader(stream, [Range(1, 0, 3), Range(6, 3, 3), Range(11, 6, 3)],
+                                'count.txt')
+
+    def test_read_block_crossing_behavior(self):
+        # read() calls will be aligned on block boundaries - see #3663.
+        sfile = self.make_count_reader()
+        self.assertEqual(b'123', sfile.read(10))
+
+    def test_small_read(self):
+        sfile = self.make_count_reader()
+        self.assertEqual(b'12', sfile.read(2))
+
+    def test_successive_reads(self):
+        sfile = self.make_count_reader()
+        for expect in [b'123', b'456', b'789', b'']:
+            self.assertEqual(expect, sfile.read(10))
+
+    def test_readfrom_spans_blocks(self):
+        sfile = self.make_count_reader()
+        self.assertEqual(b'6789', sfile.readfrom(5, 12))
+
+    def test_small_readfrom_spanning_blocks(self):
+        sfile = self.make_count_reader()
+        self.assertEqual(b'2345', sfile.readfrom(1, 4))
+
+    def test_readall(self):
+        sfile = self.make_count_reader()
+        self.assertEqual(b'123456789', b''.join(sfile.readall()))
+
+    def test_one_arg_seek(self):
+        self.test_absolute_seek([])
+
+    def test_absolute_seek(self, args=[os.SEEK_SET]):
+        sfile = self.make_count_reader()
+        sfile.seek(6, *args)
+        self.assertEqual(b'78', sfile.read(2))
+        sfile.seek(4, *args)
+        self.assertEqual(b'56', sfile.read(2))
+
+    def test_relative_seek(self, args=[os.SEEK_CUR]):
+        sfile = self.make_count_reader()
+        self.assertEqual(b'12', sfile.read(2))
+        sfile.seek(2, *args)
+        self.assertEqual(b'56', sfile.read(2))
+
+    def test_end_seek(self):
+        sfile = self.make_count_reader()
+        sfile.seek(-6, os.SEEK_END)
+        self.assertEqual(b'45', sfile.read(2))
+
+    def test_seek_min_zero(self):
+        sfile = self.make_count_reader()
+        self.assertEqual(0, sfile.tell())
+        with self.assertRaises(IOError):
+            sfile.seek(-2, os.SEEK_SET)
+        self.assertEqual(0, sfile.tell())
+
+    def test_seek_max_size(self):
+        sfile = self.make_count_reader()
+        sfile.seek(2, os.SEEK_END)
+        # POSIX permits seeking past end of file.
+        self.assertEqual(11, sfile.tell())
+
+    def test_size(self):
+        self.assertEqual(9, self.make_count_reader().size())
+
+    def test_tell_after_block_read(self):
+        sfile = self.make_count_reader()
+        sfile.read(5)
+        self.assertEqual(3, sfile.tell())
+
+    def test_tell_after_small_read(self):
+        sfile = self.make_count_reader()
+        sfile.read(1)
+        self.assertEqual(1, sfile.tell())
+
+    def test_no_read_after_close(self):
+        sfile = self.make_count_reader()
+        sfile.close()
+        self.assertRaises(ValueError, sfile.read, 2)
+
+    def test_context(self):
+        with self.make_count_reader() as sfile:
+            self.assertFalse(sfile.closed, "reader is closed inside context")
+            self.assertEqual(b'12', sfile.read(2))
+        self.assertTrue(sfile.closed, "reader is open after context")
+
+    def make_newlines_reader(self):
+        stream = tutil.MockStreamReader('.', 'one\ntwo\n\nth', 'ree\nfour\n\n')
+        return StreamFileReader(stream, [Range(0, 0, 11), Range(11, 11, 10)], 'count.txt')
+
+    def check_lines(self, actual):
+        self.assertEqual(['one\n', 'two\n', '\n', 'three\n', 'four\n', '\n'],
+                         actual)
+
+    def test_readline(self):
+        reader = self.make_newlines_reader()
+        actual = []
+        while True:
+            data = reader.readline()
+            if not data:
+                break
+            actual.append(data)
+        self.check_lines(actual)
+
+    def test_readlines(self):
+        self.check_lines(self.make_newlines_reader().readlines())
+
+    def test_iteration(self):
+        self.check_lines(list(iter(self.make_newlines_reader())))
+
+    def test_readline_size(self):
+        reader = self.make_newlines_reader()
+        self.assertEqual('on', reader.readline(2))
+        self.assertEqual('e\n', reader.readline(4))
+        self.assertEqual('two\n', reader.readline(6))
+        self.assertEqual('\n', reader.readline(8))
+        self.assertEqual('thre', reader.readline(4))
+
+    def test_readlines_sizehint(self):
+        result = self.make_newlines_reader().readlines(8)
+        self.assertEqual(['one\n', 'two\n'], result[:2])
+        self.assertNotIn('three\n', result)
+
+    def test_name_attribute(self):
+        # Test both .name and .name() (for backward compatibility)
+        stream = tutil.MockStreamReader()
+        sfile = StreamFileReader(stream, [Range(0, 0, 0)], 'nametest')
+        self.assertEqual('nametest', sfile.name)
+        self.assertEqual('nametest', sfile.name())
+
+    def check_decompressed_name(self, filename, expect):
+        stream = tutil.MockStreamReader('.', '')
+        reader = StreamFileReader(stream, [Range(0, 0, 0)], filename)
+        self.assertEqual(expect, reader.decompressed_name())
+
+    def test_decompressed_name_uncompressed_file(self):
+        self.check_decompressed_name('test.log', 'test.log')
+
+    def test_decompressed_name_gzip_file(self):
+        self.check_decompressed_name('test.log.gz', 'test.log')
+
+    def test_decompressed_name_bz2_file(self):
+        self.check_decompressed_name('test.log.bz2', 'test.log')
+
+    def check_decompression(self, compress_ext, compress_func):
+        test_text = b'decompression\ntest\n'
+        test_data = compress_func(test_text)
+        stream = tutil.MockStreamReader('.', test_data)
+        reader = StreamFileReader(stream, [Range(0, 0, len(test_data))],
+                                  'test.' + compress_ext)
+        self.assertEqual(test_text, b''.join(reader.readall_decompressed()))
+
+    @staticmethod
+    def gzip_compress(data):
+        compressed_data = io.BytesIO()
+        with gzip.GzipFile(fileobj=compressed_data, mode='wb') as gzip_file:
+            gzip_file.write(data)
+        return compressed_data.getvalue()
+
+    def test_no_decompression(self):
+        self.check_decompression('log', lambda s: s)
+
+    def test_gzip_decompression(self):
+        self.check_decompression('gz', self.gzip_compress)
+
+    def test_bz2_decompression(self):
+        self.check_decompression('bz2', bz2.compress)
+
+    def test_readline_then_readlines(self):
+        reader = self.make_newlines_reader()
+        data = reader.readline()
+        self.assertEqual('one\n', data)
+        data = reader.readlines()
+        self.assertEqual(['two\n', '\n', 'three\n', 'four\n', '\n'], data)
+
+    def test_readline_then_readall(self):
+        reader = self.make_newlines_reader()
+        data = reader.readline()
+        self.assertEqual('one\n', data)
+        self.assertEqual(b''.join([b'two\n', b'\n', b'three\n', b'four\n', b'\n']), b''.join(reader.readall()))
 
-import arvados_testutil as tutil
-import run_test_server
 
 class StreamRetryTestMixin(object):
     # Define reader_for(coll_name, **kwargs)
@@ -23,48 +215,48 @@ class StreamRetryTestMixin(object):
 
     @tutil.skip_sleep
     def test_success_without_retries(self):
-        reader = self.reader_for('bar_file')
-        with tutil.mock_responses('bar', 200):
-            self.assertEqual('bar', self.read_for_test(reader, 3))
+        with tutil.mock_keep_responses('bar', 200):
+            reader = self.reader_for('bar_file')
+            self.assertEqual(b'bar', self.read_for_test(reader, 3))
 
     @tutil.skip_sleep
     def test_read_no_default_retry(self):
-        reader = self.reader_for('user_agreement')
-        with tutil.mock_responses('', 500):
+        with tutil.mock_keep_responses('', 500):
+            reader = self.reader_for('user_agreement')
             with self.assertRaises(arvados.errors.KeepReadError):
                 self.read_for_test(reader, 10)
 
     @tutil.skip_sleep
     def test_read_with_instance_retries(self):
-        reader = self.reader_for('foo_file', num_retries=3)
-        with tutil.mock_responses('foo', 500, 200):
-            self.assertEqual('foo', self.read_for_test(reader, 3))
+        with tutil.mock_keep_responses('foo', 500, 200):
+            reader = self.reader_for('foo_file', num_retries=3)
+            self.assertEqual(b'foo', self.read_for_test(reader, 3))
 
     @tutil.skip_sleep
     def test_read_with_method_retries(self):
-        reader = self.reader_for('foo_file')
-        with tutil.mock_responses('foo', 500, 200):
-            self.assertEqual('foo',
+        with tutil.mock_keep_responses('foo', 500, 200):
+            reader = self.reader_for('foo_file')
+            self.assertEqual(b'foo',
                              self.read_for_test(reader, 3, num_retries=3))
 
     @tutil.skip_sleep
     def test_read_instance_retries_exhausted(self):
-        reader = self.reader_for('bar_file', num_retries=3)
-        with tutil.mock_responses('bar', 500, 500, 500, 500, 200):
+        with tutil.mock_keep_responses('bar', 500, 500, 500, 500, 200):
+            reader = self.reader_for('bar_file', num_retries=3)
             with self.assertRaises(arvados.errors.KeepReadError):
                 self.read_for_test(reader, 3)
 
     @tutil.skip_sleep
     def test_read_method_retries_exhausted(self):
-        reader = self.reader_for('bar_file')
-        with tutil.mock_responses('bar', 500, 500, 500, 500, 200):
+        with tutil.mock_keep_responses('bar', 500, 500, 500, 500, 200):
+            reader = self.reader_for('bar_file')
             with self.assertRaises(arvados.errors.KeepReadError):
                 self.read_for_test(reader, 3, num_retries=3)
 
     @tutil.skip_sleep
     def test_method_retries_take_precedence(self):
-        reader = self.reader_for('user_agreement', num_retries=10)
-        with tutil.mock_responses('', 500, 500, 500, 200):
+        with tutil.mock_keep_responses('', 500, 500, 500, 200):
+            reader = self.reader_for('user_agreement', num_retries=10)
             with self.assertRaises(arvados.errors.KeepReadError):
                 self.read_for_test(reader, 10, num_retries=1)
 
@@ -100,18 +292,17 @@ class StreamFileReadFromTestCase(StreamFileReadTestCase):
 
 class StreamFileReadAllTestCase(StreamFileReadTestCase):
     def read_for_test(self, reader, byte_count, **kwargs):
-        return ''.join(reader.readall(**kwargs))
+        return b''.join(reader.readall(**kwargs))
 
 
 class StreamFileReadAllDecompressedTestCase(StreamFileReadTestCase):
     def read_for_test(self, reader, byte_count, **kwargs):
-        return ''.join(reader.readall_decompressed(**kwargs))
+        return b''.join(reader.readall_decompressed(**kwargs))
 
 
 class StreamFileReadlinesTestCase(StreamFileReadTestCase):
     def read_for_test(self, reader, byte_count, **kwargs):
-        return ''.join(reader.readlines(**kwargs))
-
+        return ''.join(reader.readlines(**kwargs)).encode()
 
 if __name__ == '__main__':
     unittest.main()