sdks/python/apache_beam/io/unbounded_source_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.
#

"""Tests for apache_beam.io.unbounded_source.

Semantics are covered by deterministic unit tests; the end-to-end DirectRunner
tests assert ordering and termination only (no flaky defer-timing assertions).
"""

import logging
import unittest

from typing_extensions import override

import apache_beam as beam
from apache_beam import coders
from apache_beam.io import unbounded_source as _unbounded_source_module
from apache_beam.io.unbounded_source import _EVENT_TIME_BASE
from apache_beam.io.unbounded_source import _NO_DATA
from apache_beam.io.unbounded_source import CheckpointMark
from apache_beam.io.unbounded_source import ReadFromUnboundedSource
from apache_beam.io.unbounded_source import UnboundedCountingSource
from apache_beam.io.unbounded_source import UnboundedReader
from apache_beam.io.unbounded_source import UnboundedSource
from apache_beam.io.unbounded_source import _CountingCheckpointMark
from apache_beam.io.unbounded_source import _CountingReader
from apache_beam.io.unbounded_source import _FinalizeCheckpointOnce
from apache_beam.io.unbounded_source import _ReaderCache
from apache_beam.io.unbounded_source import _ReadFromUnboundedSourceDoFn
from apache_beam.io.unbounded_source import _set_watermark_if_greater
from apache_beam.io.unbounded_source import _UnboundedSourceRestriction
from apache_beam.io.unbounded_source import _UnboundedSourceRestrictionCoder
from apache_beam.io.unbounded_source import _UnboundedSourceRestrictionProvider
from apache_beam.io.unbounded_source import _UnboundedSourceRestrictionTracker
from apache_beam.io.watermark_estimators import ManualWatermarkEstimator
from apache_beam.options.pipeline_options import PipelineOptions
from apache_beam.options.pipeline_options_context import scoped_pipeline_options
from apache_beam.runners import sdf_utils
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 import core
from apache_beam.transforms.window import FixedWindows
from apache_beam.utils.timestamp import MAX_TIMESTAMP
from apache_beam.utils.timestamp import MIN_TIMESTAMP
from apache_beam.utils.timestamp import Timestamp

# pylint: disable=expression-not-assigned


class _StringCountingReader(_CountingReader):
  @override
  def get_current(self):
    return 'v%s' % self._current


class _StringCountingSource(UnboundedCountingSource):
  @override
  def create_reader(self, options, checkpoint_mark):
    start_index = (
        0 if checkpoint_mark is None else checkpoint_mark.last_index + 1)
    self.last_reader = _StringCountingReader(
        self._count, start_index, finalize_log=self._finalize_log)
    return self.last_reader

  @override
  def default_output_coder(self):
    return coders.StrUtf8Coder()


class _PrefixStrCoder(coders.Coder):
  def __init__(self, prefix):
    self._prefix = prefix

  @override
  def encode(self, value):
    if not value.startswith(self._prefix):
      raise ValueError('expected %r prefix' % self._prefix)
    return value[len(self._prefix):].encode('utf-8')

  @override
  def decode(self, value):
    return self._prefix + value.decode('utf-8')

  @override
  def is_deterministic(self):
    return True

  @override
  def to_type_hint(self):
    return str


class _PrefixStringReader(_StringCountingReader):
  @override
  def get_current(self):
    return 'prefix:%s' % super().get_current()


class _PrefixStringSource(_StringCountingSource):
  @override
  def create_reader(self, options, checkpoint_mark):
    start_index = (
        0 if checkpoint_mark is None else checkpoint_mark.last_index + 1)
    self.last_reader = _PrefixStringReader(
        self._count, start_index, finalize_log=self._finalize_log)
    return self.last_reader

  @override
  def default_output_coder(self):
    return _PrefixStrCoder('prefix:')


class _NoDataReader(UnboundedReader):
  """Always reports temporary absence of data with watermark below MAX."""
  @override
  def start(self):
    return False

  @override
  def advance(self):
    return False

  @override
  def get_current(self):
    raise AssertionError('no data available')

  @override
  def get_current_timestamp(self):
    raise AssertionError('no data available')

  @override
  def get_watermark(self):
    return _EVENT_TIME_BASE

  @override
  def get_checkpoint_mark(self):
    return _CountingCheckpointMark(-1)


class _NoDataSource(UnboundedSource):
  @override
  def split(self, desired_num_splits, options=None):
    return [self]

  @override
  def create_reader(self, options, checkpoint_mark):
    return _NoDataReader()

  @override
  def get_checkpoint_mark_coder(self):
    return coders.PickleCoder()


class _MutatingCheckpointMark(CheckpointMark):
  """A mark that mutates itself on finalize, to test primary/residual mark
  isolation across a checkpoint cut."""
  def __init__(self, last_index):
    self.last_index = last_index

  @override
  def finalize_checkpoint(self):
    self.last_index = -999


class _MutatingReader(UnboundedReader):
  def __init__(self):
    self._index = -1

  @override
  def start(self):
    self._index = 0
    return True

  @override
  def advance(self):
    self._index += 1
    return True

  @override
  def get_current(self):
    return self._index

  @override
  def get_current_timestamp(self):
    return _EVENT_TIME_BASE + self._index

  @override
  def get_watermark(self):
    return _EVENT_TIME_BASE + self._index

  @override
  def get_checkpoint_mark(self):
    return _MutatingCheckpointMark(self._index)


class _MutatingSource(UnboundedSource):
  @override
  def split(self, desired_num_splits, options=None):
    return [self]

  @override
  def create_reader(self, options, checkpoint_mark):
    return _MutatingReader()

  @override
  def get_checkpoint_mark_coder(self):
    return coders.PickleCoder()


class _MaxOnLastReader(UnboundedReader):
  """Returns its only record with a MAX_TIMESTAMP watermark on the same claim,
  then EOF on the next claim."""
  @override
  def start(self):
    return True

  @override
  def advance(self):
    return False

  @override
  def get_current(self):
    return 7

  @override
  def get_current_timestamp(self):
    return _EVENT_TIME_BASE

  @override
  def get_watermark(self):
    return MAX_TIMESTAMP

  @override
  def get_checkpoint_mark(self):
    return _CountingCheckpointMark(0)


class _MaxOnLastSource(UnboundedSource):
  @override
  def split(self, desired_num_splits, options=None):
    return [self]

  @override
  def create_reader(self, options, checkpoint_mark):
    return _MaxOnLastReader()

  @override
  def get_checkpoint_mark_coder(self):
    return coders.PickleCoder()


def _new_tracker(source, checkpoint=None):
  restriction = _UnboundedSourceRestriction(
      source=source, checkpoint_mark=checkpoint)
  return _UnboundedSourceRestrictionTracker(restriction)


def _claim(tracker):
  """Claims once; returns (claimed_bool, holder_value)."""
  holder = [None]
  claimed = tracker.try_claim(holder)
  return claimed, holder[0]


# ------------------------------------------------------------------------------
# Tests
# ------------------------------------------------------------------------------


class AbcContractTest(unittest.TestCase):
  def test_checkpointmark_default_finalize_is_noop(self):
    self.assertIsNone(CheckpointMark().finalize_checkpoint())

  def test_unboundedsource_is_bounded_false(self):
    self.assertFalse(UnboundedCountingSource(3).is_bounded())

  def test_reader_lifecycle_start_advance_eof(self):
    reader = UnboundedCountingSource(3).create_reader(None, None)
    self.assertTrue(reader.start())
    self.assertEqual(reader.get_current(), 0)
    self.assertEqual(reader.get_current_timestamp(), _EVENT_TIME_BASE)
    self.assertTrue(reader.advance())
    self.assertEqual(reader.get_current(), 1)
    self.assertTrue(reader.advance())
    self.assertEqual(reader.get_current(), 2)
    self.assertFalse(reader.advance())
    self.assertEqual(reader.get_watermark(), MAX_TIMESTAMP)


class RestrictionCoderTest(unittest.TestCase):
  def test_roundtrip_no_checkpoint(self):
    source = UnboundedCountingSource(3)
    coder = _UnboundedSourceRestrictionCoder()
    decoded = coder.decode(
        coder.encode(_UnboundedSourceRestriction(source=source)))
    self.assertIsNone(decoded.checkpoint_mark)
    self.assertEqual(decoded.watermark, MIN_TIMESTAMP)
    self.assertFalse(decoded.is_done)
    reader = decoded.source.create_reader(None, None)
    self.assertTrue(reader.start())
    self.assertEqual(reader.get_current(), 0)

  def test_roundtrip_with_checkpoint_resumes(self):
    source = UnboundedCountingSource(5)
    coder = _UnboundedSourceRestrictionCoder()
    restriction = _UnboundedSourceRestriction(
        source=source,
        checkpoint_mark=_CountingCheckpointMark(1),
        watermark=Timestamp(1),
        is_done=False)
    decoded = coder.decode(coder.encode(restriction))
    self.assertEqual(decoded.checkpoint_mark.last_index, 1)
    self.assertEqual(decoded.watermark, Timestamp(1))
    self.assertFalse(decoded.is_done)
    # A reader built from the decoded checkpoint resumes at the next index.
    reader = decoded.source.create_reader(None, decoded.checkpoint_mark)
    self.assertTrue(reader.start())
    self.assertEqual(reader.get_current(), 2)


class RestrictionProviderTest(unittest.TestCase):
  def test_initial_split_calls_source_split(self):
    split_log = []

    class _NamedSource(UnboundedCountingSource):
      def __init__(self, name):
        super().__init__(0)
        self.name = name

      @override
      def split(self, desired_num_splits, options=None):
        split_log.append((desired_num_splits, options))
        return [_NamedSource('a'), _NamedSource('b')]

    source = _NamedSource('root')
    provider = _UnboundedSourceRestrictionProvider()
    restriction = _UnboundedSourceRestriction(
        source=source, watermark=Timestamp(7))

    splits = list(provider.split(source, restriction))

    # The provider is a stateless module-level singleton, so it always
    # passes ``None`` as the ``options`` argument to ``UnboundedSource.split``.
    self.assertEqual(split_log, [(20, None)])
    self.assertEqual([split.source.name for split in splits], ['a', 'b'])
    self.assertEqual([split.watermark for split in splits], [Timestamp(7)] * 2)
    self.assertTrue(all(split.checkpoint_mark is None for split in splits))
    self.assertTrue(
        all(split.finalization_checkpoint_mark is None for split in splits))

  def test_initial_split_uses_desired_num_unbounded_source_splits_option(self):
    split_log = []

    class _SplitSource(UnboundedCountingSource):
      @override
      def split(self, desired_num_splits, options=None):
        split_log.append((desired_num_splits, options))
        return [self]

    source = _SplitSource(5)
    provider = _UnboundedSourceRestrictionProvider()
    restriction = _UnboundedSourceRestriction(source=source)
    options = PipelineOptions(['--desired_num_unbounded_source_splits=5'])

    with scoped_pipeline_options(options):
      list(provider.split(source, restriction))

    self.assertEqual(split_log, [(5, None)])

  def test_initial_split_does_not_split_checkpointed_restriction(self):
    split_log = []

    class _SplitSource(UnboundedCountingSource):
      @override
      def split(self, desired_num_splits, options=None):
        split_log.append((desired_num_splits, options))
        return [self]

    source = _SplitSource(5)
    provider = _UnboundedSourceRestrictionProvider()
    restriction = _UnboundedSourceRestriction(
        source=source, checkpoint_mark=_CountingCheckpointMark(2))

    self.assertEqual(list(provider.split(source, restriction)), [restriction])
    self.assertEqual(split_log, [])

  def test_initial_split_falls_back_to_original_on_split_error(self):
    class _BoomSource(UnboundedCountingSource):
      @override
      def split(self, desired_num_splits, options=None):
        raise RuntimeError('split boom')

    source = _BoomSource(5)
    provider = _UnboundedSourceRestrictionProvider()
    restriction = _UnboundedSourceRestriction(source=source)

    self.assertEqual(list(provider.split(source, restriction)), [restriction])

  def test_truncate_returns_none_for_drain(self):
    # On drain the SDF stops emitting; truncate yields no residual.
    provider = _UnboundedSourceRestrictionProvider()
    source = UnboundedCountingSource(5)
    restriction = _UnboundedSourceRestriction(source=source)
    self.assertIsNone(provider.truncate(source, restriction))

  def test_splittable_source_partitions_into_independent_subsources(self):
    # A splittable source fans out into two sub-sources; reading each in
    # isolation yields the even and the odd integers, and their union is the
    # full sequence with no overlap.
    source = UnboundedCountingSource(6, is_splittable=True)
    provider = _UnboundedSourceRestrictionProvider()
    restriction = _UnboundedSourceRestriction(source=source)

    splits = list(provider.split(source, restriction))
    self.assertEqual(len(splits), 2)

    shards = []
    for split in splits:
      tracker = _UnboundedSourceRestrictionTracker(split)
      shard = []
      while True:
        claimed, record = _claim(tracker)
        if not claimed:
          break
        if record is not _NO_DATA:
          shard.append(record[0])
      shards.append(shard)
    self.assertEqual(sorted(shards), [[0, 2, 4], [1, 3, 5]])


class RestrictionTrackerTest(unittest.TestCase):
  def test_claim_emits_in_order(self):
    tracker = _new_tracker(UnboundedCountingSource(3))
    values = []
    while True:
      claimed, record = _claim(tracker)
      if not claimed:
        break
      self.assertIsNot(record, _NO_DATA)
      values.append(record[0])
    self.assertEqual(values, [0, 1, 2])
    self.assertTrue(tracker.check_done())

  def test_claim_emits_final_record_when_watermark_is_max(self):
    # A reader may return its last record with a MAX_TIMESTAMP watermark on the
    # same call; the record must still be emitted (EOF comes on the next claim).
    class _FinalRecordReader(UnboundedReader):
      @override
      def start(self):
        return True

      @override
      def advance(self):
        return False

      @override
      def get_current(self):
        return 'last'

      @override
      def get_current_timestamp(self):
        return _EVENT_TIME_BASE

      @override
      def get_watermark(self):
        return MAX_TIMESTAMP

      @override
      def get_checkpoint_mark(self):
        return _CountingCheckpointMark(0)

    class _FinalSource(UnboundedSource):
      @override
      def split(self, desired_num_splits, options=None):
        return [self]

      @override
      def create_reader(self, options, checkpoint_mark):
        return _FinalRecordReader()

      @override
      def get_checkpoint_mark_coder(self):
        return coders.PickleCoder()

    tracker = _new_tracker(_FinalSource())
    claimed, record = _claim(tracker)
    self.assertTrue(claimed)
    self.assertIsNot(record, _NO_DATA)
    self.assertEqual(record[0], 'last')
    # The next claim observes EOF and finishes (no second, phantom record).
    claimed_again, _ = _claim(tracker)
    self.assertFalse(claimed_again)
    self.assertTrue(tracker.check_done())

  def test_try_split_zero_produces_resumable_residual(self):
    source = UnboundedCountingSource(5)
    tracker = _new_tracker(source)
    # Claim 0 and 1.
    self.assertEqual(_claim(tracker)[1][0], 0)
    self.assertEqual(_claim(tracker)[1][0], 1)

    split = tracker.try_split(0)
    self.assertIsNotNone(split)
    primary, residual = split
    self.assertTrue(primary.is_done)
    self.assertFalse(residual.is_done)
    # Resume / finalize channel separation: primary carries only the
    # finalize hook, residual carries only the resume state.
    self.assertIsNone(primary.checkpoint_mark)
    self.assertIsNotNone(primary.finalization_checkpoint_mark)
    self.assertEqual(primary.finalization_checkpoint_mark.last_index, 1)
    self.assertEqual(residual.checkpoint_mark.last_index, 1)
    self.assertIsNone(residual.finalization_checkpoint_mark)
    # check_done passes on the (now done) primary.
    self.assertTrue(tracker.check_done())

    # Resuming from the residual continues at index 2.
    resumed = _new_tracker(source, checkpoint=residual.checkpoint_mark)
    self.assertEqual(_claim(resumed)[1][0], 2)

  def test_try_split_isolates_residual_from_finalize_mutation(self):
    # The primary's finalize hook and the residual's resume state must not
    # share one object, so a mutating finalize_checkpoint() cannot corrupt the
    # residual's resume position.
    tracker = _new_tracker(_MutatingSource())
    _claim(tracker)  # start -> index 0
    split = tracker.try_split(0)
    self.assertIsNotNone(split)
    primary, residual = split
    self.assertIsNot(
        primary.finalization_checkpoint_mark, residual.checkpoint_mark)
    self.assertEqual(residual.checkpoint_mark.last_index, 0)
    primary.finalization_checkpoint_mark.finalize_checkpoint()
    self.assertEqual(residual.checkpoint_mark.last_index, 0)

  def test_try_split_nonzero_declined(self):
    source = UnboundedCountingSource(5)
    tracker = _new_tracker(source)
    self.assertEqual(_claim(tracker)[1][0], 0)

    self.assertIsNone(tracker.try_split(0.5))
    self.assertFalse(tracker.current_restriction().is_done)
    self.assertIsNotNone(tracker._reader)
    self.assertEqual(_claim(tracker)[1][0], 1)

  def test_no_data_returns_sentinel_without_finishing(self):
    tracker = _new_tracker(_NoDataSource())
    claimed, record = _claim(tracker)
    self.assertTrue(claimed)
    self.assertIs(record, _NO_DATA)
    # A self-checkpoint is still possible (poll/resume path).
    self.assertIsNotNone(tracker.try_split(0))

  def test_check_done_raises_when_not_done(self):
    tracker = _new_tracker(UnboundedCountingSource(3))
    with self.assertRaises(ValueError):
      tracker.check_done()

  def test_is_bounded_false(self):
    self.assertFalse(_new_tracker(UnboundedCountingSource(3)).is_bounded())


class _RecordingBundleFinalizer:
  def __init__(self):
    self.registered = []

  def register(self, callback):
    self.registered.append(callback)


class _ManualClock:
  """A deterministic monotonic clock for the time-cap tests."""
  def __init__(self, now=0.0):
    self.now = now

  def __call__(self):
    return self.now


class BundleCapTest(unittest.TestCase):
  """A busy reader self-checkpoints once the per-bundle record or time cap is
  reached, so the runner can commit progress and run finalization."""
  def _bundle(self, dofn, source, checkpoint=None, estimator=None):
    """Builds the SDF tracker chain and returns the process() generator plus the
    tracker, threadsafe tracker, finalizer, and watermark estimator."""
    tracker = _UnboundedSourceRestrictionTracker(
        _UnboundedSourceRestriction(source=source, checkpoint_mark=checkpoint))
    threadsafe = sdf_utils.ThreadsafeRestrictionTracker(tracker)
    view = sdf_utils.RestrictionTrackerView(threadsafe)
    finalizer = _RecordingBundleFinalizer()
    estimator = estimator or ManualWatermarkEstimator(None)
    gen = dofn.process(
        None,
        bundle_finalizer=finalizer,
        tracker=view,
        watermark_estimator=estimator)
    return gen, tracker, threadsafe, finalizer, estimator

  def test_record_cap_checkpoints_busy_source(self):
    finalize_log = []
    dofn = _ReadFromUnboundedSourceDoFn(
        poll_interval=0, max_records_per_bundle=5, max_read_time_seconds=1e9)
    # 1000 records is effectively unbounded against a cap of 5.
    gen, tracker, threadsafe, finalizer, estimator = self._bundle(
        dofn, UnboundedCountingSource(1000, finalize_log=finalize_log))
    outputs = list(gen)

    self.assertEqual([tv.value for tv in outputs], [0, 1, 2, 3, 4])
    self.assertTrue(tracker.current_restriction().is_done)
    self.assertTrue(tracker.check_done())
    # The estimator holds the last emitted record's source watermark.
    self.assertEqual(estimator.current_watermark(), _EVENT_TIME_BASE + 4)
    # Residual resumes after the cut and carries no finalize hook.
    residual, _ = threadsafe.deferred_status()
    self.assertEqual(residual.checkpoint_mark.last_index, 4)
    self.assertIsNone(residual.finalization_checkpoint_mark)
    # Exactly one finalizer is registered; firing it commits the cut index once.
    self.assertEqual(len(finalizer.registered), 1)
    finalizer.registered[0]()
    finalizer.registered[0]()
    self.assertEqual(finalize_log, [4])

  def test_time_cap_checkpoints_busy_source(self):
    clock = _ManualClock(1000.0)
    dofn = _ReadFromUnboundedSourceDoFn(
        poll_interval=0,
        max_records_per_bundle=10**9,
        max_read_time_seconds=5.0,
        _now=clock)
    gen, tracker, threadsafe, _, _ = self._bundle(
        dofn, UnboundedCountingSource(1000))

    # The deadline arms at 1000 + 5 after the first record and is checked
    # between records, so records keep flowing until the clock passes it.
    self.assertEqual(next(gen).value, 0)
    self.assertEqual(next(gen).value, 1)
    self.assertEqual(next(gen).value, 2)
    clock.now = 1006.0
    with self.assertRaises(StopIteration):
      next(gen)

    self.assertTrue(tracker.current_restriction().is_done)
    residual, _ = threadsafe.deferred_status()
    self.assertEqual(residual.checkpoint_mark.last_index, 2)

  def test_cap_residual_resumes_in_next_bundle(self):
    dofn = _ReadFromUnboundedSourceDoFn(
        poll_interval=0, max_records_per_bundle=5, max_read_time_seconds=1e9)
    source = UnboundedCountingSource(1000)
    # Bundle 1 emits 0-4 and cuts a residual at index 4.
    gen1, _, threadsafe1, _, _ = self._bundle(dofn, source)
    self.assertEqual([tv.value for tv in gen1], [0, 1, 2, 3, 4])
    residual1, _ = threadsafe1.deferred_status()

    # Bundle 2 rebuilds the reader from the residual and emits 5-9.
    gen2, _, threadsafe2, _, _ = self._bundle(
        dofn, source, checkpoint=residual1.checkpoint_mark)
    self.assertEqual([tv.value for tv in gen2], [5, 6, 7, 8, 9])
    residual2, _ = threadsafe2.deferred_status()
    self.assertEqual(residual2.checkpoint_mark.last_index, 9)

  def test_busy_reader_is_reused_across_bundles(self):
    # The self-checkpoint parks the reader; the resuming bundle reclaims the
    # same started reader.
    dofn = _ReadFromUnboundedSourceDoFn(
        poll_interval=0, max_records_per_bundle=5, max_read_time_seconds=1e9)
    dofn.setup()  # creates the cross-bundle reader cache
    source = UnboundedCountingSource(1000)
    gen1, _, threadsafe1, _, _ = self._bundle(dofn, source)
    self.assertEqual([tv.value for tv in gen1], [0, 1, 2, 3, 4])
    reader1 = source.last_reader
    self.assertFalse(reader1.closed)  # parked, not closed
    residual1, _ = threadsafe1.deferred_status()

    gen2, _, _, _, _ = self._bundle(
        dofn, source, checkpoint=residual1.checkpoint_mark)
    self.assertEqual([tv.value for tv in gen2], [5, 6, 7, 8, 9])
    # No new reader was created; source.last_reader is unchanged.
    self.assertIs(source.last_reader, reader1)

  def test_teardown_closes_parked_readers(self):
    dofn = _ReadFromUnboundedSourceDoFn(
        poll_interval=0, max_records_per_bundle=5, max_read_time_seconds=1e9)
    dofn.setup()
    source = UnboundedCountingSource(1000)
    gen, _, _, _, _ = self._bundle(dofn, source)
    list(gen)  # the bundle parks its reader
    reader = source.last_reader
    self.assertFalse(reader.closed)
    dofn.teardown()
    self.assertTrue(reader.closed)

  def test_eof_exactly_at_cap_resumes_then_finishes(self):
    dofn = _ReadFromUnboundedSourceDoFn(
        poll_interval=0, max_records_per_bundle=5, max_read_time_seconds=1e9)
    source = UnboundedCountingSource(5)  # exactly cap records
    # Bundle 1 hits the cap on the last record before observing EOF.
    gen1, t1, threadsafe1, _, _ = self._bundle(dofn, source)
    self.assertEqual([tv.value for tv in gen1], [0, 1, 2, 3, 4])
    self.assertTrue(t1.current_restriction().is_done)
    residual1, _ = threadsafe1.deferred_status()
    self.assertEqual(residual1.checkpoint_mark.last_index, 4)

    # Bundle 2 resumes at index 5, finds EOF, and finishes with no output.
    gen2, t2, threadsafe2, _, _ = self._bundle(
        dofn, source, checkpoint=residual1.checkpoint_mark)
    self.assertEqual(list(gen2), [])
    self.assertTrue(t2.current_restriction().is_done)
    self.assertIsNone(threadsafe2.deferred_status())

  def test_eof_before_cap_finishes_without_residual(self):
    dofn = _ReadFromUnboundedSourceDoFn(
        poll_interval=0, max_records_per_bundle=100, max_read_time_seconds=1e9)
    gen, tracker, threadsafe, _, _ = self._bundle(
        dofn, UnboundedCountingSource(3))

    self.assertEqual([tv.value for tv in gen], [0, 1, 2])
    self.assertTrue(tracker.current_restriction().is_done)
    self.assertIsNone(threadsafe.deferred_status())

  def test_max_watermark_on_final_record_emitted_before_estimator_advances(
      self):
    # A reader that returns its final record with a MAX_TIMESTAMP watermark on
    # the same claim must have that record emitted before the estimator reaches
    # MAX, so the element is not stranded behind the output watermark.
    dofn = _ReadFromUnboundedSourceDoFn(poll_interval=0)
    gen, _, _, _, estimator = self._bundle(dofn, _MaxOnLastSource())

    first = next(gen)
    self.assertEqual(first.value, 7)
    # The estimator has not yet been advanced to MAX when the record is yielded.
    self.assertNotEqual(estimator.current_watermark(), MAX_TIMESTAMP)
    # Draining the bundle then advances the estimator to MAX on EOF.
    self.assertEqual(list(gen), [])
    self.assertEqual(estimator.current_watermark(), MAX_TIMESTAMP)


class WatermarkTest(unittest.TestCase):
  def test_set_watermark_is_monotonic(self):
    estimator = ManualWatermarkEstimator(None)
    _set_watermark_if_greater(estimator, Timestamp(5))
    self.assertEqual(estimator.current_watermark(), Timestamp(5))
    # A regression is ignored (would otherwise raise inside set_watermark).
    _set_watermark_if_greater(estimator, Timestamp(3))
    self.assertEqual(estimator.current_watermark(), Timestamp(5))
    _set_watermark_if_greater(estimator, Timestamp(7))
    self.assertEqual(estimator.current_watermark(), Timestamp(7))


class FinalizationTest(unittest.TestCase):
  def test_finalize_checkpoint_callback_is_at_most_once(self):
    finalize_log = []
    finalize_once = _FinalizeCheckpointOnce(
        _CountingCheckpointMark(1, finalize_log=finalize_log))

    finalize_once()
    finalize_once()

    self.assertEqual(finalize_log, [1])

  def test_finalize_checkpoint_invoked(self):
    # Unit-level finalize test (the e2e finalize may run in a worker process);
    # the hook lives on the primary, independent of the residual's resume state.
    finalize_log = []
    source = UnboundedCountingSource(5, finalize_log=finalize_log)
    tracker = _new_tracker(source)
    _claim(tracker)  # 0
    _claim(tracker)  # 1
    primary, _ = tracker.try_split(0)
    primary.finalization_checkpoint_mark.finalize_checkpoint()
    self.assertEqual(finalize_log, [1])


class EndToEndTest(unittest.TestCase):
  def test_direct_runner_emits_all_in_order(self):
    with TestPipeline() as p:
      out = p | ReadFromUnboundedSource(UnboundedCountingSource(5))
      self.assertFalse(out.is_bounded)
      assert_that(out, equal_to([0, 1, 2, 3, 4]))

  def test_eof_lets_event_time_window_fire(self):
    # On EOF the DoFn advances the watermark estimator to MAX_TIMESTAMP so the
    # downstream FixedWindow closes and the GroupByKey fires; otherwise the
    # output would be empty.
    with TestPipeline() as p:
      out = (
          p
          | ReadFromUnboundedSource(UnboundedCountingSource(5))
          | beam.WindowInto(FixedWindows(100))
          | beam.Map(lambda v: ('all', v))
          | beam.GroupByKey()
          | beam.MapTuple(lambda _key, values: sorted(values)))
      assert_that(out, equal_to([[0, 1, 2, 3, 4]]))

  def test_read_dispatches_through_iobase_read(self):
    # ``beam.io.Read(source)`` must produce the same records as
    # ``ReadFromUnboundedSource(source)``.
    with TestPipeline() as p:
      out = p | beam.io.Read(UnboundedCountingSource(5))
      self.assertFalse(out.is_bounded)
      assert_that(out, equal_to([0, 1, 2, 3, 4]))

  def test_splittable_source_reads_all_records_across_splits(self):
    # A splittable source fans out into even/odd sub-sources during initial
    # SDF splitting; the union of all sub-source reads is the full sequence.
    with TestPipeline() as p:
      out = p | beam.io.Read(UnboundedCountingSource(6, is_splittable=True))
      assert_that(out, equal_to([0, 1, 2, 3, 4, 5]))

  def test_source_default_output_coder_sets_output_type(self):
    with TestPipeline() as p:
      out = p | ReadFromUnboundedSource(_StringCountingSource(2))
      self.assertEqual(out.element_type, str)
      assert_that(out, equal_to(['v0', 'v1']))

  def test_small_cap_self_checkpoints_and_resumes_to_eof(self):
    # A small per-bundle cap forces several self-checkpoint/resume cycles
    # through the runner before EOF, exercising the cross-bundle reader cache
    # over real residual encode/decode. All records still arrive in order.
    with TestPipeline() as p:
      out = p | ReadFromUnboundedSource(
          UnboundedCountingSource(20), max_records_per_bundle=5)
      assert_that(out, equal_to(list(range(20))))


class ReadFromUnboundedSourceCoderTest(unittest.TestCase):
  def test_parameterized_output_coder_does_not_mutate_global_registry(self):
    try:
      p = beam.Pipeline()
      out = p | ReadFromUnboundedSource(_PrefixStringSource(1))

      self.assertNotEqual(out.element_type, str)
      self.assertEqual(coders.registry.get_coder(str), coders.StrUtf8Coder())
      self.assertEqual(
          ReadFromUnboundedSource(_PrefixStringSource(1))._infer_output_coder(),
          _PrefixStrCoder('prefix:'))
    finally:
      coders.registry.register_coder(str, coders.StrUtf8Coder)


# ------------------------------------------------------------------------------
# Reader lifecycle, watermark, and contract regression tests (reader close on
# every exit path, the NotImplementedError message, finalize idempotency).
# ------------------------------------------------------------------------------


class ReaderCloseTest(unittest.TestCase):
  """Reader lifecycle: close() must run on every tracker-driven exit path."""
  def test_tracker_closes_reader_on_eof(self):
    source = UnboundedCountingSource(0)  # immediately exhausted
    tracker = _new_tracker(source)
    holder = [None]
    self.assertFalse(tracker.try_claim(holder))
    self.assertIsNone(tracker._reader)
    self.assertTrue(source.last_reader.closed)

  def test_tracker_closes_reader_on_split_without_cache(self):
    # With no cache injected, a self-checkpoint closes the reader.
    source = UnboundedCountingSource(5)
    tracker = _new_tracker(source)
    _claim(tracker)  # creates reader, claims 0
    reader = source.last_reader
    self.assertFalse(reader.closed)
    split = tracker.try_split(0)
    self.assertIsNotNone(split)
    self.assertIsNone(tracker._reader)
    self.assertTrue(reader.closed)

  def test_tracker_parks_reader_on_split_with_cache(self):
    # With a cache, a self-checkpoint parks the reader for the residual to
    # reclaim.
    source = UnboundedCountingSource(5)
    cache = _ReaderCache()
    tracker = _new_tracker(source)
    tracker._reader_cache = cache
    _claim(tracker)
    reader = source.last_reader
    _, residual = tracker.try_split(0)
    self.assertIsNone(tracker._reader)
    self.assertFalse(reader.closed)
    # The residual's key reclaims the same started reader.
    self.assertEqual(
        cache.acquire(tracker._cache_key(residual)), (reader, True))

  def test_close_helper_is_idempotent_and_safe_on_empty_tracker(self):
    tracker = _new_tracker(UnboundedCountingSource(3))
    # No reader yet -- helper must be a no-op.
    tracker._close_reader_if_open()
    _claim(tracker)
    reader = tracker._reader
    tracker._close_reader_if_open()
    self.assertTrue(reader.closed)
    self.assertIsNone(tracker._reader)
    # Second call is a no-op (no reader to close).
    tracker._close_reader_if_open()

  def test_close_helper_swallows_reader_close_errors(self):
    class _BoomReader(UnboundedReader):
      @override
      def start(self):
        return True

      @override
      def advance(self):
        return False

      @override
      def get_current(self):
        return 'x'

      @override
      def get_current_timestamp(self):
        return _EVENT_TIME_BASE

      @override
      def get_watermark(self):
        return _EVENT_TIME_BASE

      @override
      def get_checkpoint_mark(self):
        return CheckpointMark()

      @override
      def close(self):
        raise RuntimeError('close blew up')

    class _BoomSource(UnboundedSource):
      @override
      def split(self, desired_num_splits, options=None):
        return [self]

      @override
      def create_reader(self, options, checkpoint_mark):
        return _BoomReader()

      @override
      def get_checkpoint_mark_coder(self):
        return coders.PickleCoder()

    tracker = _new_tracker(_BoomSource())
    _claim(tracker)
    with self.assertLogs(_unbounded_source_module._LOGGER, 'WARNING') as logs:
      tracker._close_reader_if_open()
    self.assertTrue(
        any('Error closing UnboundedReader' in line for line in logs.output))
    self.assertIsNone(tracker._reader)


class ReaderCacheTest(unittest.TestCase):
  """The cache parks a reader under one key, hands it to the next acquirer,
  bounds itself by idle time and entry count, and closes on teardown."""
  def _reader(self):
    class _FakeReader(UnboundedReader):
      def __init__(self):
        self.closed = False

      @override
      def close(self):
        self.closed = True

    return _FakeReader()

  def test_park_then_acquire_returns_same_reader_and_started_flag(self):
    cache = _ReaderCache()
    reader = self._reader()
    cache.park('k', reader, True)
    self.assertEqual(cache.acquire('k'), (reader, True))
    # Acquire removes the entry so two trackers cannot share one reader.
    self.assertIsNone(cache.acquire('k'))
    self.assertFalse(reader.closed)

  def test_acquire_miss_returns_none(self):
    self.assertIsNone(_ReaderCache().acquire('absent'))

  def test_park_replacing_same_key_closes_displaced_reader(self):
    # Two parks under one key without an intervening acquire (e.g. a bundle
    # retry) must not leak the first reader.
    cache = _ReaderCache()
    old, new = self._reader(), self._reader()
    cache.park('k', old, True)
    cache.park('k', new, True)
    self.assertTrue(old.closed)
    self.assertFalse(new.closed)
    self.assertEqual(cache.acquire('k'), (new, True))

  def test_idle_reader_is_closed_on_next_touch(self):
    clock = _ManualClock(1000.0)
    cache = _ReaderCache(idle_seconds=30.0, now=clock)
    reader = self._reader()
    cache.park('k', reader, True)
    clock.now = 1031.0  # past the idle window
    cache.park('other', self._reader(), False)  # triggers idle eviction
    self.assertTrue(reader.closed)
    self.assertIsNone(cache.acquire('k'))

  def test_max_size_evicts_and_closes_oldest(self):
    cache = _ReaderCache(max_size=2)
    readers = [self._reader() for _ in range(3)]
    cache.park('a', readers[0], True)
    cache.park('b', readers[1], True)
    cache.park('c', readers[2], True)  # exceeds the cap, evicts 'a'
    self.assertTrue(readers[0].closed)
    self.assertIsNone(cache.acquire('a'))
    self.assertIsNotNone(cache.acquire('b'))
    self.assertIsNotNone(cache.acquire('c'))

  def test_close_all_closes_every_parked_reader(self):
    cache = _ReaderCache()
    readers = [self._reader() for _ in range(3)]
    for i, reader in enumerate(readers):
      cache.park(str(i), reader, True)
    cache.close_all()
    self.assertTrue(all(reader.closed for reader in readers))
    self.assertIsNone(cache.acquire('0'))

  def test_close_all_swallows_reader_close_errors(self):
    class _BoomReader(UnboundedReader):
      @override
      def close(self):
        raise RuntimeError('close blew up')

    cache = _ReaderCache()
    cache.park('k', _BoomReader(), True)
    with self.assertLogs(_unbounded_source_module._LOGGER, 'WARNING') as logs:
      cache.close_all()
    self.assertTrue(
        any('Error closing UnboundedReader' in line for line in logs.output))


class TrackerContractRegressionTest(unittest.TestCase):
  """Tracker contract: source-watermark on the data path, finalize/resume
  channel separation, and reader close on a reader-method failure."""
  def test_data_path_holder_carries_source_watermark(self):
    class _LaggingReader(UnboundedReader):
      @override
      def start(self):
        return True

      @override
      def advance(self):
        return False

      @override
      def get_current(self):
        return 'rec'

      @override
      def get_current_timestamp(self):
        return Timestamp(1000)  # record event time

      @override
      def get_watermark(self):
        return Timestamp(990)  # source watermark lags 10us behind

      @override
      def get_checkpoint_mark(self):
        return _CountingCheckpointMark(0)

    class _LaggingSource(UnboundedSource):
      @override
      def split(self, desired_num_splits, options=None):
        return [self]

      @override
      def create_reader(self, options, checkpoint_mark):
        return _LaggingReader()

      @override
      def get_checkpoint_mark_coder(self):
        return coders.PickleCoder()

    tracker = _new_tracker(_LaggingSource())
    claimed, record = _claim(tracker)
    self.assertTrue(claimed)
    self.assertIsNot(record, _NO_DATA)
    value, record_timestamp, source_watermark = record
    self.assertEqual(value, 'rec')
    self.assertEqual(record_timestamp, Timestamp(1000))
    # Critical: watermark slot is the SOURCE watermark, NOT record timestamp.
    self.assertEqual(source_watermark, Timestamp(990))
    self.assertNotEqual(source_watermark, record_timestamp)

  def test_split_separates_finalize_and_resume_channels(self):
    source = UnboundedCountingSource(5)
    tracker = _new_tracker(source)
    _claim(tracker)  # claim 0 so reader has progress
    primary, residual = tracker.try_split(0)
    # Primary carries ONLY the finalize hook -- no resume state.
    self.assertIsNone(primary.checkpoint_mark)
    self.assertIsNotNone(primary.finalization_checkpoint_mark)
    self.assertTrue(primary.is_done)
    # Residual carries ONLY the resume state -- no finalize hook (a future
    # bundle that splits THIS residual will produce ITS own finalize mark).
    self.assertIsNotNone(residual.checkpoint_mark)
    self.assertIsNone(residual.finalization_checkpoint_mark)
    self.assertFalse(residual.is_done)
    # The two marks reference the same underlying checkpoint object.
    self.assertEqual(
        primary.finalization_checkpoint_mark.last_index,
        residual.checkpoint_mark.last_index)

  def test_eof_populates_finalize_and_clears_resume(self):
    # EOF transition: restriction.checkpoint_mark goes to None (no more
    # records to resume from), finalization_checkpoint_mark carries the
    # final commit hook.
    source = UnboundedCountingSource(0)  # immediately exhausted
    tracker = _new_tracker(source)
    holder = [None]
    self.assertFalse(tracker.try_claim(holder))
    r = tracker.current_restriction()
    self.assertTrue(r.is_done)
    self.assertEqual(r.watermark, MAX_TIMESTAMP)
    self.assertIsNone(r.checkpoint_mark)
    self.assertIsNotNone(r.finalization_checkpoint_mark)

  def test_tracker_closes_reader_when_advance_raises(self):
    # try_claim closes the reader before re-raising a reader-method failure, so
    # the DoFn's finally need not traverse the SDF chain for these.
    class _BoomReader(UnboundedReader):
      def __init__(self):
        self.closed = False

      @override
      def start(self):
        return True

      @override
      def advance(self):
        raise RuntimeError('advance boom')

      @override
      def get_current(self):
        return 'first'

      @override
      def get_current_timestamp(self):
        return _EVENT_TIME_BASE

      @override
      def get_watermark(self):
        return _EVENT_TIME_BASE

      @override
      def get_checkpoint_mark(self):
        return _CountingCheckpointMark(0)

      @override
      def close(self):
        self.closed = True

    class _BoomSource(UnboundedSource):
      @override
      def split(self, desired_num_splits, options=None):
        return [self]

      @override
      def create_reader(self, options, checkpoint_mark):
        return _BoomReader()

      @override
      def get_checkpoint_mark_coder(self):
        return coders.PickleCoder()

    src = _BoomSource()
    tracker = _new_tracker(src)
    # First claim succeeds (start returns True).
    self.assertTrue(tracker.try_claim([None]))
    reader_after_first = tracker._reader
    self.assertIsNotNone(reader_after_first)
    # The second claim's advance() raises; the tracker must close the reader
    # before propagating.
    with self.assertRaises(RuntimeError):
      tracker.try_claim([None])
    self.assertTrue(reader_after_first.closed)
    self.assertIsNone(tracker._reader)

  def test_tracker_closes_reader_when_get_watermark_raises(self):
    # Reader method failures other than advance() also trigger close.
    class _WatermarkBoomReader(UnboundedReader):
      def __init__(self):
        self.closed = False

      @override
      def start(self):
        return False  # no data -> drops into get_watermark path

      @override
      def advance(self):
        return False

      @override
      def get_current(self):
        raise AssertionError

      @override
      def get_current_timestamp(self):
        raise AssertionError

      @override
      def get_watermark(self):
        raise RuntimeError('watermark boom')

      @override
      def get_checkpoint_mark(self):
        return _CountingCheckpointMark(0)

      @override
      def close(self):
        self.closed = True

    class _WatermarkBoomSource(UnboundedSource):
      @override
      def split(self, desired_num_splits, options=None):
        return [self]

      @override
      def create_reader(self, options, checkpoint_mark):
        return _WatermarkBoomReader()

      @override
      def get_checkpoint_mark_coder(self):
        return coders.PickleCoder()

    src = _WatermarkBoomSource()
    tracker = _new_tracker(src)
    with self.assertRaises(RuntimeError):
      tracker.try_claim([None])
    self.assertIsNone(tracker._reader)


class UnboundedSourceContractTest(unittest.TestCase):
  def test_get_checkpoint_mark_coder_default_names_subclass(self):
    class MySource(UnboundedSource):
      pass

    with self.assertRaises(NotImplementedError) as cm:
      MySource().get_checkpoint_mark_coder()
    self.assertIn('MySource', str(cm.exception))


class ReadFromUnboundedSourceValidationTest(unittest.TestCase):
  def test_non_source_argument_raises(self):
    with self.assertRaises(TypeError):
      ReadFromUnboundedSource('not-a-source')  # type: ignore[arg-type]

  def test_invalid_caps_raise(self):
    source = UnboundedCountingSource(1)
    with self.assertRaises(ValueError):
      ReadFromUnboundedSource(source, max_records_per_bundle=0)
    with self.assertRaises(ValueError):
      ReadFromUnboundedSource(source, max_read_time_seconds=0)
    with self.assertRaises(ValueError):
      ReadFromUnboundedSource(source, poll_interval=-1)


if __name__ == '__main__':
  logging.getLogger().setLevel(logging.INFO)
  unittest.main()