iobase_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.
#

"""Unit tests for classes in iobase.py."""

# pytype: skip-file

import unittest

import mock

import apache_beam as beam
from apache_beam.io import iobase
from apache_beam.io import range_trackers
from apache_beam.io.concat_source import ConcatSource
from apache_beam.io.concat_source_test import RangeSource
from apache_beam.io.iobase import SourceBundle
from apache_beam.options.pipeline_options import DebugOptions
from apache_beam.testing.util import assert_that
from apache_beam.testing.util import equal_to
from apache_beam.utils import timestamp


class SDFBoundedSourceRestrictionProviderTest(unittest.TestCase):
  def setUp(self):
    self.initial_range_start = 0
    self.initial_range_stop = 4
    self.initial_range_source = RangeSource(
        self.initial_range_start, self.initial_range_stop)
    self.sdf_restriction_provider = (
        iobase._SDFBoundedSourceRestrictionProvider(desired_chunk_size=2))

  def test_initial_restriction(self):
    element = self.initial_range_source
    restriction = (self.sdf_restriction_provider.initial_restriction(element))
    self.assertTrue(
        isinstance(restriction, iobase._SDFBoundedSourceRestriction))
    self.assertTrue(isinstance(restriction._source_bundle, SourceBundle))
    self.assertEqual(
        self.initial_range_start, restriction._source_bundle.start_position)
    self.assertEqual(
        self.initial_range_stop, restriction._source_bundle.stop_position)
    self.assertTrue(isinstance(restriction._source_bundle.source, RangeSource))
    self.assertEqual(restriction._range_tracker, None)

  def test_create_tracker(self):
    expected_start = 1
    expected_stop = 3
    source_bundle = SourceBundle(
        expected_stop - expected_start,
        RangeSource(1, 3),
        expected_start,
        expected_stop)
    restriction_tracker = (
        self.sdf_restriction_provider.create_tracker(
            iobase._SDFBoundedSourceRestriction(source_bundle)))
    self.assertTrue(
        isinstance(
            restriction_tracker, iobase._SDFBoundedSourceRestrictionTracker))
    self.assertEqual(expected_start, restriction_tracker.start_pos())
    self.assertEqual(expected_stop, restriction_tracker.stop_pos())

  def test_simple_source_split(self):
    element = self.initial_range_source
    restriction = (self.sdf_restriction_provider.initial_restriction(element))
    expect_splits = [(0, 2), (2, 4)]
    split_bundles = list(
        self.sdf_restriction_provider.split(element, restriction))
    self.assertTrue(
        all(
            isinstance(bundle._source_bundle, SourceBundle)
            for bundle in split_bundles))

    splits = ([(
        bundle._source_bundle.start_position,
        bundle._source_bundle.stop_position) for bundle in split_bundles])
    self.assertEqual(expect_splits, list(splits))

  def test_concat_source_split(self):
    element = self.initial_range_source
    initial_concat_source = ConcatSource([self.initial_range_source])
    sdf_concat_restriction_provider = (
        iobase._SDFBoundedSourceRestrictionProvider(desired_chunk_size=2))
    restriction = (self.sdf_restriction_provider.initial_restriction(element))
    expect_splits = [(0, 2), (2, 4)]
    split_bundles = list(
        sdf_concat_restriction_provider.split(
            initial_concat_source, restriction))
    self.assertTrue(
        all(
            isinstance(bundle._source_bundle, SourceBundle)
            for bundle in split_bundles))
    splits = ([(
        bundle._source_bundle.start_position,
        bundle._source_bundle.stop_position) for bundle in split_bundles])
    self.assertEqual(expect_splits, list(splits))

  def test_restriction_size(self):
    element = self.initial_range_source
    restriction = (self.sdf_restriction_provider.initial_restriction(element))
    split_1, split_2 = self.sdf_restriction_provider.split(element, restriction)
    split_1_size = self.sdf_restriction_provider.restriction_size(
        element, split_1)
    split_2_size = self.sdf_restriction_provider.restriction_size(
        element, split_2)
    self.assertEqual(2, split_1_size)
    self.assertEqual(2, split_2_size)


class SDFBoundedSourceRestrictionTrackerTest(unittest.TestCase):
  def setUp(self):
    self.initial_start_pos = 0
    self.initial_stop_pos = 4
    source_bundle = SourceBundle(
        self.initial_stop_pos - self.initial_start_pos,
        RangeSource(self.initial_start_pos, self.initial_stop_pos),
        self.initial_start_pos,
        self.initial_stop_pos)
    self.sdf_restriction_tracker = (
        iobase._SDFBoundedSourceRestrictionTracker(
            iobase._SDFBoundedSourceRestriction(source_bundle)))

  def test_current_restriction_before_split(self):
    current_restriction = (self.sdf_restriction_tracker.current_restriction())
    self.assertEqual(
        self.initial_start_pos,
        current_restriction._source_bundle.start_position)
    self.assertEqual(
        self.initial_stop_pos, current_restriction._source_bundle.stop_position)
    self.assertEqual(
        self.initial_start_pos,
        current_restriction._range_tracker.start_position())
    self.assertEqual(
        self.initial_stop_pos,
        current_restriction._range_tracker.stop_position())

  def test_current_restriction_after_split(self):
    fraction_of_remainder = 0.5
    self.sdf_restriction_tracker.try_claim(1)
    expected_restriction, _ = (
        self.sdf_restriction_tracker.try_split(fraction_of_remainder))
    current_restriction = self.sdf_restriction_tracker.current_restriction()
    self.assertEqual(
        expected_restriction._source_bundle, current_restriction._source_bundle)
    self.assertTrue(current_restriction._range_tracker)

  def test_try_split_at_remainder(self):
    fraction_of_remainder = 0.4
    expected_primary = (0, 2, 2.0)
    expected_residual = (2, 4, 2.0)
    self.sdf_restriction_tracker.try_claim(0)
    actual_primary, actual_residual = (
        self.sdf_restriction_tracker.try_split(fraction_of_remainder))
    self.assertEqual(
        expected_primary,
        (
            actual_primary._source_bundle.start_position,
            actual_primary._source_bundle.stop_position,
            actual_primary._source_bundle.weight))
    self.assertEqual(
        expected_residual,
        (
            actual_residual._source_bundle.start_position,
            actual_residual._source_bundle.stop_position,
            actual_residual._source_bundle.weight))
    self.assertEqual(
        actual_primary._source_bundle.weight,
        self.sdf_restriction_tracker.current_restriction().weight())

  def test_try_split_with_any_exception(self):
    source_bundle = SourceBundle(
        range_trackers.OffsetRangeTracker.OFFSET_INFINITY,
        RangeSource(0, range_trackers.OffsetRangeTracker.OFFSET_INFINITY),
        0,
        range_trackers.OffsetRangeTracker.OFFSET_INFINITY)
    self.sdf_restriction_tracker = (
        iobase._SDFBoundedSourceRestrictionTracker(
            iobase._SDFBoundedSourceRestriction(source_bundle)))
    self.sdf_restriction_tracker.try_claim(0)
    self.assertIsNone(self.sdf_restriction_tracker.try_split(0.5))


class UseSdfBoundedSourcesTests(unittest.TestCase):
  def _run_sdf_wrapper_pipeline(self, source, expected_values):
    with beam.Pipeline() as p:
      experiments = (p._options.view_as(DebugOptions).experiments or [])

      # Setup experiment option to enable using SDFBoundedSourceWrapper
      if 'beam_fn_api' not in experiments:
        # Required so mocking below doesn't mock Create used in assert_that.
        experiments.append('beam_fn_api')

      p._options.view_as(DebugOptions).experiments = experiments

      actual = p | beam.io.Read(source)
      assert_that(actual, equal_to(expected_values))

  @mock.patch('apache_beam.io.iobase.SDFBoundedSourceReader.expand')
  def test_sdf_wrapper_overrides_read(self, sdf_wrapper_mock_expand):
    def _fake_wrapper_expand(pbegin):
      return pbegin | beam.Map(lambda x: 'fake')

    sdf_wrapper_mock_expand.side_effect = _fake_wrapper_expand
    self._run_sdf_wrapper_pipeline(RangeSource(0, 4), ['fake'])

  def test_sdf_wrap_range_source(self):
    self._run_sdf_wrapper_pipeline(RangeSource(0, 4), [0, 1, 2, 3])


class UnboundedSourceTest(unittest.TestCase):
  """Basic tests for UnboundedSource, UnboundedReader, and CheckpointMark.

  These tests verify the foundational API structure. Full integration with
  Read transform and SDF wrappers will be tested in future implementation.
  """

  def test_checkpoint_mark_finalize(self):
    """Test that CheckpointMark can be subclassed and finalized."""
    class TestCheckpointMark(iobase.CheckpointMark):
      def __init__(self):
        self.finalized = False

      def finalize(self):
        self.finalized = True

    checkpoint = TestCheckpointMark()
    self.assertFalse(checkpoint.finalized)
    checkpoint.finalize()
    self.assertTrue(checkpoint.finalized)

  def test_unbounded_source_basic_interface(self):
    """Test that UnboundedSource can be subclassed with basic methods."""
    class TestUnboundedSource(iobase.UnboundedSource):
      def reader(self, checkpoint=None):
        return TestUnboundedReader()

      def default_output_coder(self):
        return beam.coders.VarIntCoder()

    class TestUnboundedReader(iobase.UnboundedReader):
      def __init__(self):
        self.index = -1
        self.data = [1, 2, 3]

      def start(self):
        self.index = 0
        return True

      def advance(self):
        self.index += 1
        return self.index < len(self.data)

      def get_current(self):
        return self.data[self.index]

      def get_current_timestamp(self):
        return timestamp.Timestamp.of(self.index)

      def get_watermark(self):
        return timestamp.Timestamp.of(self.index)

      def get_checkpoint_mark(self):
        return iobase.CheckpointMark()

    source = TestUnboundedSource()
    self.assertFalse(source.is_bounded())

    # Test reader basic operations
    reader = source.reader()
    self.assertTrue(reader.start())
    self.assertEqual(1, reader.get_current())
    self.assertEqual(timestamp.Timestamp.of(0), reader.get_current_timestamp())

    self.assertTrue(reader.advance())
    self.assertEqual(2, reader.get_current())

    self.assertTrue(reader.advance())
    self.assertEqual(3, reader.get_current())

    self.assertFalse(reader.advance())

  def test_unbounded_source_split_default(self):
    """Test that UnboundedSource.split() returns [self] by default."""
    class SimpleUnboundedSource(iobase.UnboundedSource):
      def reader(self, checkpoint=None):
        pass

      def default_output_coder(self):
        return beam.coders.VarIntCoder()

    source = SimpleUnboundedSource()
    splits = source.split(desired_num_splits=10)
    self.assertEqual([source], splits)


if __name__ == '__main__':
  unittest.main()