| # Licensed to the Apache Software Foundation (ASF) under one or more |
| # contributor license agreements. See the NOTICE file distributed with |
| # this work for additional information regarding copyright ownership. |
| # The ASF licenses this file to You under the Apache License, Version 2.0 |
| # (the "License"); you may not use this file except in compliance with |
| # the License. You may obtain a copy of the License at |
| # |
| # http://www.apache.org/licenses/LICENSE-2.0 |
| # |
| # Unless required by applicable law or agreed to in writing, software |
| # distributed under the License is distributed on an "AS IS" BASIS, |
| # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| # See the License for the specific language governing permissions and |
| # limitations under the License. |
| # |
| |
| # pytype: skip-file |
| |
| import bz2 |
| import gzip |
| import io |
| import logging |
| import math |
| import os |
| import random |
| import tempfile |
| import unittest |
| |
| import hamcrest as hc |
| |
| import apache_beam as beam |
| from apache_beam.io import filebasedsource |
| from apache_beam.io import iobase |
| from apache_beam.io import range_trackers |
| # importing following private classes for testing |
| from apache_beam.io.concat_source import ConcatSource |
| from apache_beam.io.filebasedsource import _SingleFileSource as SingleFileSource |
| from apache_beam.io.filebasedsource import FileBasedSource |
| from apache_beam.io.filesystem import CompressionTypes |
| from apache_beam.options.value_provider import RuntimeValueProvider |
| from apache_beam.options.value_provider import StaticValueProvider |
| from apache_beam.testing.test_pipeline import TestPipeline |
| from apache_beam.testing.util import assert_that |
| from apache_beam.testing.util import equal_to |
| from apache_beam.transforms.display import DisplayData |
| from apache_beam.transforms.display_test import DisplayDataItemMatcher |
| |
| |
| class LineSource(FileBasedSource): |
| def read_records(self, file_name, range_tracker): |
| f = self.open_file(file_name) |
| try: |
| start = range_tracker.start_position() |
| if start > 0: |
| # Any line that starts after 'start' does not belong to the current |
| # bundle. Seeking to (start - 1) and skipping a line moves the current |
| # position to the starting position of the first line that belongs to |
| # the current bundle. |
| start -= 1 |
| f.seek(start) |
| line = f.readline() |
| start += len(line) |
| current = start |
| line = f.readline() |
| while range_tracker.try_claim(current): |
| # When the source is unsplittable, try_claim is not enough to determine |
| # whether the file has reached to the end. |
| if not line: |
| return |
| yield line.rstrip(b'\n') |
| current += len(line) |
| line = f.readline() |
| finally: |
| f.close() |
| |
| |
| class EOL(object): |
| LF = 1 |
| CRLF = 2 |
| MIXED = 3 |
| LF_WITH_NOTHING_AT_LAST_LINE = 4 |
| |
| |
| def write_data( |
| num_lines, |
| no_data=False, |
| directory=None, |
| prefix=tempfile.template, |
| eol=EOL.LF): |
| """Writes test data to a temporary file. |
| |
| Args: |
| num_lines (int): The number of lines to write. |
| no_data (bool): If :data:`True`, empty lines will be written, otherwise |
| each line will contain a concatenation of b'line' and the line number. |
| directory (str): The name of the directory to create the temporary file in. |
| prefix (str): The prefix to use for the temporary file. |
| eol (int): The line ending to use when writing. |
| :class:`~apache_beam.io.filebasedsource_test.EOL` exposes attributes that |
| can be used here to define the eol. |
| |
| Returns: |
| Tuple[str, List[bytes]]: A tuple of the filename and a list of the written |
| data. |
| """ |
| all_data = [] |
| with tempfile.NamedTemporaryFile(delete=False, dir=directory, |
| prefix=prefix) as f: |
| sep_values = [b'\n', b'\r\n'] |
| for i in range(num_lines): |
| data = b'' if no_data else b'line' + str(i).encode() |
| all_data.append(data) |
| |
| if eol == EOL.LF: |
| sep = sep_values[0] |
| elif eol == EOL.CRLF: |
| sep = sep_values[1] |
| elif eol == EOL.MIXED: |
| sep = sep_values[i % len(sep_values)] |
| elif eol == EOL.LF_WITH_NOTHING_AT_LAST_LINE: |
| sep = b'' if i == (num_lines - 1) else sep_values[0] |
| else: |
| raise ValueError('Received unknown value %s for eol.' % eol) |
| |
| f.write(data + sep) |
| |
| return f.name, all_data |
| |
| |
| def _write_prepared_data( |
| data, directory=None, prefix=tempfile.template, suffix=''): |
| with tempfile.NamedTemporaryFile(delete=False, |
| dir=directory, |
| prefix=prefix, |
| suffix=suffix) as f: |
| f.write(data) |
| return f.name |
| |
| |
| def write_prepared_pattern(data, suffixes=None): |
| assert data, 'Data (%s) seems to be empty' % data |
| if suffixes is None: |
| suffixes = [''] * len(data) |
| temp_dir = tempfile.mkdtemp() |
| for i, d in enumerate(data): |
| file_name = _write_prepared_data( |
| d, temp_dir, prefix='mytemp', suffix=suffixes[i]) |
| return file_name[:file_name.rfind(os.path.sep)] + os.path.sep + 'mytemp*' |
| |
| |
| def write_pattern(lines_per_file, no_data=False): |
| """Writes a pattern of temporary files. |
| |
| Args: |
| lines_per_file (List[int]): The number of lines to write per file. |
| no_data (bool): If :data:`True`, empty lines will be written, otherwise |
| each line will contain a concatenation of b'line' and the line number. |
| |
| Returns: |
| Tuple[str, List[bytes]]: A tuple of the filename pattern and a list of the |
| written data. |
| """ |
| temp_dir = tempfile.mkdtemp() |
| |
| all_data = [] |
| file_name = None |
| start_index = 0 |
| for i in range(len(lines_per_file)): |
| file_name, data = write_data(lines_per_file[i], no_data=no_data, |
| directory=temp_dir, prefix='mytemp') |
| all_data.extend(data) |
| start_index += lines_per_file[i] |
| |
| assert file_name |
| return ( |
| file_name[:file_name.rfind(os.path.sep)] + os.path.sep + 'mytemp*', |
| all_data) |
| |
| |
| class TestConcatSource(unittest.TestCase): |
| class DummySource(iobase.BoundedSource): |
| def __init__(self, values): |
| self._values = values |
| |
| def split( |
| self, desired_bundle_size, start_position=None, stop_position=None): |
| # simply devides values into two bundles |
| middle = len(self._values) // 2 |
| yield iobase.SourceBundle( |
| 0.5, TestConcatSource.DummySource(self._values[:middle]), None, None) |
| yield iobase.SourceBundle( |
| 0.5, TestConcatSource.DummySource(self._values[middle:]), None, None) |
| |
| def get_range_tracker(self, start_position, stop_position): |
| if start_position is None: |
| start_position = 0 |
| if stop_position is None: |
| stop_position = len(self._values) |
| |
| return range_trackers.OffsetRangeTracker(start_position, stop_position) |
| |
| def read(self, range_tracker): |
| for index, value in enumerate(self._values): |
| if not range_tracker.try_claim(index): |
| return |
| |
| yield value |
| |
| def estimate_size(self): |
| return len(self._values) # Assuming each value to be 1 byte. |
| |
| def setUp(self): |
| # Reducing the size of thread pools. Without this test execution may fail in |
| # environments with limited amount of resources. |
| filebasedsource.MAX_NUM_THREADS_FOR_SIZE_ESTIMATION = 2 |
| |
| def test_read(self): |
| sources = [ |
| TestConcatSource.DummySource(range(start, start + 10)) |
| for start in [0, 10, 20] |
| ] |
| concat = ConcatSource(sources) |
| range_tracker = concat.get_range_tracker(None, None) |
| read_data = [value for value in concat.read(range_tracker)] |
| self.assertCountEqual(list(range(30)), read_data) |
| |
| def test_split(self): |
| sources = [ |
| TestConcatSource.DummySource(list(range(start, start + 10))) |
| for start in [0, 10, 20] |
| ] |
| concat = ConcatSource(sources) |
| splits = [split for split in concat.split()] |
| self.assertEqual(6, len(splits)) |
| |
| # Reading all splits |
| read_data = [] |
| for split in splits: |
| range_tracker_for_split = split.source.get_range_tracker( |
| split.start_position, split.stop_position) |
| read_data.extend( |
| [value for value in split.source.read(range_tracker_for_split)]) |
| self.assertCountEqual(list(range(30)), read_data) |
| |
| def test_estimate_size(self): |
| sources = [ |
| TestConcatSource.DummySource(range(start, start + 10)) |
| for start in [0, 10, 20] |
| ] |
| concat = ConcatSource(sources) |
| self.assertEqual(30, concat.estimate_size()) |
| |
| |
| class TestFileBasedSource(unittest.TestCase): |
| def setUp(self): |
| # Reducing the size of thread pools. Without this test execution may fail in |
| # environments with limited amount of resources. |
| filebasedsource.MAX_NUM_THREADS_FOR_SIZE_ESTIMATION = 2 |
| |
| def test_string_or_value_provider_only(self): |
| str_file_pattern = tempfile.NamedTemporaryFile(delete=False).name |
| self.assertEqual( |
| str_file_pattern, FileBasedSource(str_file_pattern)._pattern.value) |
| |
| static_vp_file_pattern = StaticValueProvider( |
| value_type=str, value=str_file_pattern) |
| self.assertEqual( |
| static_vp_file_pattern, |
| FileBasedSource(static_vp_file_pattern)._pattern) |
| |
| runtime_vp_file_pattern = RuntimeValueProvider( |
| option_name='arg', value_type=str, default_value=str_file_pattern) |
| self.assertEqual( |
| runtime_vp_file_pattern, |
| FileBasedSource(runtime_vp_file_pattern)._pattern) |
| # Reset runtime options to avoid side-effects in other tests. |
| RuntimeValueProvider.set_runtime_options(None) |
| |
| invalid_file_pattern = 123 |
| with self.assertRaises(TypeError): |
| FileBasedSource(invalid_file_pattern) |
| |
| def test_validation_file_exists(self): |
| file_name, _ = write_data(10) |
| LineSource(file_name) |
| |
| def test_validation_directory_non_empty(self): |
| temp_dir = tempfile.mkdtemp() |
| file_name, _ = write_data(10, directory=temp_dir) |
| LineSource(file_name) |
| |
| def test_validation_failing(self): |
| no_files_found_error = 'No files found based on the file pattern*' |
| with self.assertRaisesRegex(IOError, no_files_found_error): |
| LineSource('dummy_pattern') |
| with self.assertRaisesRegex(IOError, no_files_found_error): |
| temp_dir = tempfile.mkdtemp() |
| LineSource(os.path.join(temp_dir, '*')) |
| |
| def test_validation_file_missing_verification_disabled(self): |
| LineSource('dummy_pattern', validate=False) |
| |
| def test_fully_read_single_file(self): |
| file_name, expected_data = write_data(10) |
| assert len(expected_data) == 10 |
| fbs = LineSource(file_name) |
| range_tracker = fbs.get_range_tracker(None, None) |
| read_data = [record for record in fbs.read(range_tracker)] |
| self.assertCountEqual(expected_data, read_data) |
| |
| def test_single_file_display_data(self): |
| file_name, _ = write_data(10) |
| fbs = LineSource(file_name) |
| dd = DisplayData.create_from(fbs) |
| expected_items = [ |
| DisplayDataItemMatcher('file_pattern', file_name), |
| DisplayDataItemMatcher('compression', 'auto') |
| ] |
| hc.assert_that(dd.items, hc.contains_inanyorder(*expected_items)) |
| |
| def test_fully_read_file_pattern(self): |
| pattern, expected_data = write_pattern([5, 3, 12, 8, 8, 4]) |
| assert len(expected_data) == 40 |
| fbs = LineSource(pattern) |
| range_tracker = fbs.get_range_tracker(None, None) |
| read_data = [record for record in fbs.read(range_tracker)] |
| self.assertCountEqual(expected_data, read_data) |
| |
| def test_fully_read_file_pattern_with_empty_files(self): |
| pattern, expected_data = write_pattern([5, 0, 12, 0, 8, 0]) |
| assert len(expected_data) == 25 |
| fbs = LineSource(pattern) |
| range_tracker = fbs.get_range_tracker(None, None) |
| read_data = [record for record in fbs.read(range_tracker)] |
| self.assertCountEqual(expected_data, read_data) |
| |
| def test_estimate_size_of_file(self): |
| file_name, expected_data = write_data(10) |
| assert len(expected_data) == 10 |
| fbs = LineSource(file_name) |
| self.assertEqual(10 * 6, fbs.estimate_size()) |
| |
| def test_estimate_size_of_pattern(self): |
| pattern, expected_data = write_pattern([5, 3, 10, 8, 8, 4]) |
| assert len(expected_data) == 38 |
| fbs = LineSource(pattern) |
| self.assertEqual(38 * 6, fbs.estimate_size()) |
| |
| pattern, expected_data = write_pattern([5, 3, 9]) |
| assert len(expected_data) == 17 |
| fbs = LineSource(pattern) |
| self.assertEqual(17 * 6, fbs.estimate_size()) |
| |
| def test_estimate_size_with_sampling_same_size(self): |
| num_files = 2 * FileBasedSource.MIN_NUMBER_OF_FILES_TO_STAT |
| pattern, _ = write_pattern([10] * num_files) |
| # Each line will be of length 6 since write_pattern() uses |
| # ('line' + line number + '\n') as data. |
| self.assertEqual( |
| 6 * 10 * num_files, FileBasedSource(pattern).estimate_size()) |
| |
| def test_estimate_size_with_sampling_different_sizes(self): |
| num_files = 2 * FileBasedSource.MIN_NUMBER_OF_FILES_TO_STAT |
| |
| # Each line will be of length 8 since write_pattern() uses |
| # ('line' + line number + '\n') as data. |
| base_size = 500 |
| variance = 5 |
| |
| sizes = [] |
| for _ in range(num_files): |
| sizes.append( |
| int(random.uniform(base_size - variance, base_size + variance))) |
| pattern, _ = write_pattern(sizes) |
| tolerance = 0.05 |
| self.assertAlmostEqual( |
| base_size * 8 * num_files, |
| FileBasedSource(pattern).estimate_size(), |
| delta=base_size * 8 * num_files * tolerance) |
| |
| def test_splits_into_subranges(self): |
| pattern, expected_data = write_pattern([5, 9, 6]) |
| assert len(expected_data) == 20 |
| fbs = LineSource(pattern) |
| splits = [split for split in fbs.split(desired_bundle_size=15)] |
| expected_num_splits = ( |
| math.ceil(float(6 * 5) / 15) + math.ceil(float(6 * 9) / 15) + |
| math.ceil(float(6 * 6) / 15)) |
| assert len(splits) == expected_num_splits |
| |
| def test_read_splits_single_file(self): |
| file_name, expected_data = write_data(100) |
| assert len(expected_data) == 100 |
| fbs = LineSource(file_name) |
| splits = [split for split in fbs.split(desired_bundle_size=33)] |
| |
| # Reading all splits |
| read_data = [] |
| for split in splits: |
| source = split.source |
| range_tracker = source.get_range_tracker( |
| split.start_position, split.stop_position) |
| data_from_split = [data for data in source.read(range_tracker)] |
| read_data.extend(data_from_split) |
| |
| self.assertCountEqual(expected_data, read_data) |
| |
| def test_read_splits_file_pattern(self): |
| pattern, expected_data = write_pattern([34, 66, 40, 24, 24, 12]) |
| assert len(expected_data) == 200 |
| fbs = LineSource(pattern) |
| splits = [split for split in fbs.split(desired_bundle_size=50)] |
| |
| # Reading all splits |
| read_data = [] |
| for split in splits: |
| source = split.source |
| range_tracker = source.get_range_tracker( |
| split.start_position, split.stop_position) |
| data_from_split = [data for data in source.read(range_tracker)] |
| read_data.extend(data_from_split) |
| |
| self.assertCountEqual(expected_data, read_data) |
| |
| def _run_source_test(self, pattern, expected_data, splittable=True): |
| with TestPipeline() as pipeline: |
| pcoll = pipeline | 'Read' >> beam.io.Read( |
| LineSource(pattern, splittable=splittable)) |
| assert_that(pcoll, equal_to(expected_data)) |
| |
| def test_source_file(self): |
| file_name, expected_data = write_data(100) |
| assert len(expected_data) == 100 |
| self._run_source_test(file_name, expected_data) |
| |
| def test_source_pattern(self): |
| pattern, expected_data = write_pattern([34, 66, 40, 24, 24, 12]) |
| assert len(expected_data) == 200 |
| self._run_source_test(pattern, expected_data) |
| |
| def test_unsplittable_does_not_split(self): |
| pattern, expected_data = write_pattern([5, 9, 6]) |
| assert len(expected_data) == 20 |
| fbs = LineSource(pattern, splittable=False) |
| splits = [split for split in fbs.split(desired_bundle_size=15)] |
| self.assertEqual(3, len(splits)) |
| |
| def test_source_file_unsplittable(self): |
| file_name, expected_data = write_data(100) |
| assert len(expected_data) == 100 |
| self._run_source_test(file_name, expected_data, False) |
| |
| def test_source_pattern_unsplittable(self): |
| pattern, expected_data = write_pattern([34, 66, 40, 24, 24, 12]) |
| assert len(expected_data) == 200 |
| self._run_source_test(pattern, expected_data, False) |
| |
| def test_read_file_bzip2(self): |
| _, lines = write_data(10) |
| filename = tempfile.NamedTemporaryFile( |
| delete=False, prefix=tempfile.template).name |
| with bz2.BZ2File(filename, 'wb') as f: |
| f.write(b'\n'.join(lines)) |
| |
| with TestPipeline() as pipeline: |
| pcoll = pipeline | 'Read' >> beam.io.Read( |
| LineSource( |
| filename, |
| splittable=False, |
| compression_type=CompressionTypes.BZIP2)) |
| assert_that(pcoll, equal_to(lines)) |
| |
| def test_read_file_gzip(self): |
| _, lines = write_data(10) |
| filename = tempfile.NamedTemporaryFile( |
| delete=False, prefix=tempfile.template).name |
| with gzip.GzipFile(filename, 'wb') as f: |
| f.write(b'\n'.join(lines)) |
| |
| with TestPipeline() as pipeline: |
| pcoll = pipeline | 'Read' >> beam.io.Read( |
| LineSource( |
| filename, |
| splittable=False, |
| compression_type=CompressionTypes.GZIP)) |
| assert_that(pcoll, equal_to(lines)) |
| |
| def test_read_pattern_bzip2(self): |
| _, lines = write_data(200) |
| splits = [0, 34, 100, 140, 164, 188, 200] |
| chunks = [lines[splits[i - 1]:splits[i]] for i in range(1, len(splits))] |
| compressed_chunks = [] |
| for c in chunks: |
| compressobj = bz2.BZ2Compressor() |
| compressed_chunks.append( |
| compressobj.compress(b'\n'.join(c)) + compressobj.flush()) |
| file_pattern = write_prepared_pattern(compressed_chunks) |
| with TestPipeline() as pipeline: |
| pcoll = pipeline | 'Read' >> beam.io.Read( |
| LineSource( |
| file_pattern, |
| splittable=False, |
| compression_type=CompressionTypes.BZIP2)) |
| assert_that(pcoll, equal_to(lines)) |
| |
| def test_read_pattern_gzip(self): |
| _, lines = write_data(200) |
| splits = [0, 34, 100, 140, 164, 188, 200] |
| chunks = [lines[splits[i - 1]:splits[i]] for i in range(1, len(splits))] |
| compressed_chunks = [] |
| for c in chunks: |
| out = io.BytesIO() |
| with gzip.GzipFile(fileobj=out, mode="wb") as f: |
| f.write(b'\n'.join(c)) |
| compressed_chunks.append(out.getvalue()) |
| file_pattern = write_prepared_pattern(compressed_chunks) |
| with TestPipeline() as pipeline: |
| pcoll = pipeline | 'Read' >> beam.io.Read( |
| LineSource( |
| file_pattern, |
| splittable=False, |
| compression_type=CompressionTypes.GZIP)) |
| assert_that(pcoll, equal_to(lines)) |
| |
| def test_read_auto_single_file_bzip2(self): |
| _, lines = write_data(10) |
| filename = tempfile.NamedTemporaryFile( |
| delete=False, prefix=tempfile.template, suffix='.bz2').name |
| with bz2.BZ2File(filename, 'wb') as f: |
| f.write(b'\n'.join(lines)) |
| |
| with TestPipeline() as pipeline: |
| pcoll = pipeline | 'Read' >> beam.io.Read( |
| LineSource(filename, compression_type=CompressionTypes.AUTO)) |
| assert_that(pcoll, equal_to(lines)) |
| |
| def test_read_auto_single_file_gzip(self): |
| _, lines = write_data(10) |
| filename = tempfile.NamedTemporaryFile( |
| delete=False, prefix=tempfile.template, suffix='.gz').name |
| with gzip.GzipFile(filename, 'wb') as f: |
| f.write(b'\n'.join(lines)) |
| |
| with TestPipeline() as pipeline: |
| pcoll = pipeline | 'Read' >> beam.io.Read( |
| LineSource(filename, compression_type=CompressionTypes.AUTO)) |
| assert_that(pcoll, equal_to(lines)) |
| |
| def test_read_auto_pattern(self): |
| _, lines = write_data(200) |
| splits = [0, 34, 100, 140, 164, 188, 200] |
| chunks = [lines[splits[i - 1]:splits[i]] for i in range(1, len(splits))] |
| compressed_chunks = [] |
| for c in chunks: |
| out = io.BytesIO() |
| with gzip.GzipFile(fileobj=out, mode="wb") as f: |
| f.write(b'\n'.join(c)) |
| compressed_chunks.append(out.getvalue()) |
| file_pattern = write_prepared_pattern( |
| compressed_chunks, suffixes=['.gz'] * len(chunks)) |
| with TestPipeline() as pipeline: |
| pcoll = pipeline | 'Read' >> beam.io.Read( |
| LineSource(file_pattern, compression_type=CompressionTypes.AUTO)) |
| assert_that(pcoll, equal_to(lines)) |
| |
| def test_read_auto_pattern_compressed_and_uncompressed(self): |
| _, lines = write_data(200) |
| splits = [0, 34, 100, 140, 164, 188, 200] |
| chunks = [lines[splits[i - 1]:splits[i]] for i in range(1, len(splits))] |
| chunks_to_write = [] |
| for i, c in enumerate(chunks): |
| if i % 2 == 0: |
| out = io.BytesIO() |
| with gzip.GzipFile(fileobj=out, mode="wb") as f: |
| f.write(b'\n'.join(c)) |
| chunks_to_write.append(out.getvalue()) |
| else: |
| chunks_to_write.append(b'\n'.join(c)) |
| file_pattern = write_prepared_pattern( |
| chunks_to_write, suffixes=(['.gz', ''] * 3)) |
| with TestPipeline() as pipeline: |
| pcoll = pipeline | 'Read' >> beam.io.Read( |
| LineSource(file_pattern, compression_type=CompressionTypes.AUTO)) |
| assert_that(pcoll, equal_to(lines)) |
| |
| def test_splits_get_coder_from_fbs(self): |
| class DummyCoder(object): |
| val = 12345 |
| |
| class FileBasedSourceWithCoder(LineSource): |
| def default_output_coder(self): |
| return DummyCoder() |
| |
| pattern, expected_data = write_pattern([34, 66, 40, 24, 24, 12]) |
| self.assertEqual(200, len(expected_data)) |
| fbs = FileBasedSourceWithCoder(pattern) |
| splits = [split for split in fbs.split(desired_bundle_size=50)] |
| self.assertTrue(len(splits)) |
| for split in splits: |
| self.assertEqual(DummyCoder.val, split.source.default_output_coder().val) |
| |
| |
| class TestSingleFileSource(unittest.TestCase): |
| def setUp(self): |
| # Reducing the size of thread pools. Without this test execution may fail in |
| # environments with limited amount of resources. |
| filebasedsource.MAX_NUM_THREADS_FOR_SIZE_ESTIMATION = 2 |
| |
| def test_source_creation_fails_for_non_number_offsets(self): |
| start_not_a_number_error = 'start_offset must be a number*' |
| stop_not_a_number_error = 'stop_offset must be a number*' |
| file_name = 'dummy_pattern' |
| fbs = LineSource(file_name, validate=False) |
| |
| with self.assertRaisesRegex(TypeError, start_not_a_number_error): |
| SingleFileSource( |
| fbs, file_name='dummy_file', start_offset='aaa', stop_offset='bbb') |
| with self.assertRaisesRegex(TypeError, start_not_a_number_error): |
| SingleFileSource( |
| fbs, file_name='dummy_file', start_offset='aaa', stop_offset=100) |
| with self.assertRaisesRegex(TypeError, stop_not_a_number_error): |
| SingleFileSource( |
| fbs, file_name='dummy_file', start_offset=100, stop_offset='bbb') |
| with self.assertRaisesRegex(TypeError, stop_not_a_number_error): |
| SingleFileSource( |
| fbs, file_name='dummy_file', start_offset=100, stop_offset=None) |
| with self.assertRaisesRegex(TypeError, start_not_a_number_error): |
| SingleFileSource( |
| fbs, file_name='dummy_file', start_offset=None, stop_offset=100) |
| |
| def test_source_creation_display_data(self): |
| file_name = 'dummy_pattern' |
| fbs = LineSource(file_name, validate=False) |
| dd = DisplayData.create_from(fbs) |
| expected_items = [ |
| DisplayDataItemMatcher('compression', 'auto'), |
| DisplayDataItemMatcher('file_pattern', file_name) |
| ] |
| hc.assert_that(dd.items, hc.contains_inanyorder(*expected_items)) |
| |
| def test_source_creation_fails_if_start_lg_stop(self): |
| start_larger_than_stop_error = ( |
| 'start_offset must be smaller than stop_offset*') |
| fbs = LineSource('dummy_pattern', validate=False) |
| SingleFileSource( |
| fbs, file_name='dummy_file', start_offset=99, stop_offset=100) |
| with self.assertRaisesRegex(ValueError, start_larger_than_stop_error): |
| SingleFileSource( |
| fbs, file_name='dummy_file', start_offset=100, stop_offset=99) |
| with self.assertRaisesRegex(ValueError, start_larger_than_stop_error): |
| SingleFileSource( |
| fbs, file_name='dummy_file', start_offset=100, stop_offset=100) |
| |
| def test_estimates_size(self): |
| fbs = LineSource('dummy_pattern', validate=False) |
| |
| # Should simply return stop_offset - start_offset |
| source = SingleFileSource( |
| fbs, file_name='dummy_file', start_offset=0, stop_offset=100) |
| self.assertEqual(100, source.estimate_size()) |
| |
| source = SingleFileSource( |
| fbs, file_name='dummy_file', start_offset=10, stop_offset=100) |
| self.assertEqual(90, source.estimate_size()) |
| |
| def test_read_range_at_beginning(self): |
| fbs = LineSource('dummy_pattern', validate=False) |
| |
| file_name, expected_data = write_data(10) |
| assert len(expected_data) == 10 |
| |
| source = SingleFileSource(fbs, file_name, 0, 10 * 6) |
| range_tracker = source.get_range_tracker(0, 20) |
| read_data = [value for value in source.read(range_tracker)] |
| self.assertCountEqual(expected_data[:4], read_data) |
| |
| def test_read_range_at_end(self): |
| fbs = LineSource('dummy_pattern', validate=False) |
| |
| file_name, expected_data = write_data(10) |
| assert len(expected_data) == 10 |
| |
| source = SingleFileSource(fbs, file_name, 0, 10 * 6) |
| range_tracker = source.get_range_tracker(40, 60) |
| read_data = [value for value in source.read(range_tracker)] |
| self.assertCountEqual(expected_data[-3:], read_data) |
| |
| def test_read_range_at_middle(self): |
| fbs = LineSource('dummy_pattern', validate=False) |
| |
| file_name, expected_data = write_data(10) |
| assert len(expected_data) == 10 |
| |
| source = SingleFileSource(fbs, file_name, 0, 10 * 6) |
| range_tracker = source.get_range_tracker(20, 40) |
| read_data = [value for value in source.read(range_tracker)] |
| self.assertCountEqual(expected_data[4:7], read_data) |
| |
| def test_produces_splits_desiredsize_large_than_size(self): |
| fbs = LineSource('dummy_pattern', validate=False) |
| |
| file_name, expected_data = write_data(10) |
| assert len(expected_data) == 10 |
| source = SingleFileSource(fbs, file_name, 0, 10 * 6) |
| splits = [split for split in source.split(desired_bundle_size=100)] |
| self.assertEqual(1, len(splits)) |
| self.assertEqual(60, splits[0].weight) |
| self.assertEqual(0, splits[0].start_position) |
| self.assertEqual(60, splits[0].stop_position) |
| |
| range_tracker = splits[0].source.get_range_tracker(None, None) |
| read_data = [value for value in splits[0].source.read(range_tracker)] |
| self.assertCountEqual(expected_data, read_data) |
| |
| def test_produces_splits_desiredsize_smaller_than_size(self): |
| fbs = LineSource('dummy_pattern', validate=False) |
| |
| file_name, expected_data = write_data(10) |
| assert len(expected_data) == 10 |
| source = SingleFileSource(fbs, file_name, 0, 10 * 6) |
| splits = [split for split in source.split(desired_bundle_size=25)] |
| self.assertEqual(3, len(splits)) |
| |
| read_data = [] |
| for split in splits: |
| source = split.source |
| range_tracker = source.get_range_tracker( |
| split.start_position, split.stop_position) |
| data_from_split = [data for data in source.read(range_tracker)] |
| read_data.extend(data_from_split) |
| self.assertCountEqual(expected_data, read_data) |
| |
| def test_produce_split_with_start_and_end_positions(self): |
| fbs = LineSource('dummy_pattern', validate=False) |
| |
| file_name, expected_data = write_data(10) |
| assert len(expected_data) == 10 |
| source = SingleFileSource(fbs, file_name, 0, 10 * 6) |
| splits = [ |
| split for split in source.split( |
| desired_bundle_size=15, start_offset=10, stop_offset=50) |
| ] |
| self.assertEqual(3, len(splits)) |
| |
| read_data = [] |
| for split in splits: |
| source = split.source |
| range_tracker = source.get_range_tracker( |
| split.start_position, split.stop_position) |
| data_from_split = [data for data in source.read(range_tracker)] |
| read_data.extend(data_from_split) |
| self.assertCountEqual(expected_data[2:9], read_data) |
| |
| |
| if __name__ == '__main__': |
| logging.getLogger().setLevel(logging.INFO) |
| unittest.main() |
