github.com/apache/beam/sdks/v2@v2.48.2/python/apache_beam/io/filebasedsource_test.py

# 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()