[Python] Add UnboundedSource ValidatesRunner test (#38892)

Author: Eliaaazzz

.github/trigger_files/beam_PostCommit_Python_ValidatesRunner_Dataflow.json

@@ -1,4 +1,5 @@
 {
   "comment": "Modify this file in a trivial way to cause this test suite to run",
-  "modification": 5
+  "modification": 5,
+  "https://github.com/apache/beam/pull/38892": "UnboundedSource portable VR test"
 }

.github/trigger_files/beam_PostCommit_Python_ValidatesRunner_Flink.json

@@ -1,5 +1,6 @@
 {
   "https://github.com/apache/beam/pull/32648": "testing addition of Flink 1.19 support",
   "https://github.com/apache/beam/pull/34830": "testing",
-  "trigger-2026-04-04": "portable_runner expand_sdf opt-in"
+  "trigger-2026-04-04": "portable_runner expand_sdf opt-in",
+  "https://github.com/apache/beam/pull/38892": "UnboundedSource portable VR test"
 }

.github/trigger_files/beam_PostCommit_Python_ValidatesRunner_Spark.json

@@ -1,5 +1,6 @@
 {
   "https://github.com/apache/beam/pull/34830": "testing",
   "https://github.com/apache/beam/issues/35429": "testing",
-  "trigger-2026-04-04": "portable_runner expand_sdf opt-in"
+  "trigger-2026-04-04": "portable_runner expand_sdf opt-in",
+  "https://github.com/apache/beam/pull/38892": "UnboundedSource portable VR test"
 }

sdks/python/apache_beam/io/iobase_test.py

@@ -224,22 +224,22 @@ def test_sdf_wrap_range_source(self):
 class UseSdfUnboundedSourcesTests(unittest.TestCase):
   """Covers the UnboundedSource branch in
   ``iobase.Read.expand()``. Uses ``UnboundedCountingSource`` from
-  ``unbounded_source_test`` as a finite fake source (no network).
+  ``unbounded_source`` as a finite fake source (no network).
   """
   def test_read_end_to_end_unbounded(self):
-    from apache_beam.io.unbounded_source_test import UnboundedCountingSource
+    from apache_beam.io.unbounded_source import UnboundedCountingSource
     with beam.Pipeline() as p:
       out = p | beam.io.Read(UnboundedCountingSource(5))
       assert_that(out, equal_to([0, 1, 2, 3, 4]))
 
   def test_read_unbounded_pcollection_is_unbounded(self):
-    from apache_beam.io.unbounded_source_test import UnboundedCountingSource
+    from apache_beam.io.unbounded_source import UnboundedCountingSource
     p = beam.Pipeline()
     out = p | beam.io.Read(UnboundedCountingSource(3))
     self.assertFalse(out.is_bounded)
 
   def test_read_unbounded_serializes_as_expanded_composite(self):
-    from apache_beam.io.unbounded_source_test import UnboundedCountingSource
+    from apache_beam.io.unbounded_source import UnboundedCountingSource
     p = beam.Pipeline()
     p | 'ReadIt' >> beam.io.Read(UnboundedCountingSource(3))
 

sdks/python/apache_beam/io/unbounded_source.py

@@ -93,6 +93,8 @@ def get_checkpoint_mark_coder(self):
 from typing import Iterable
 from typing import Optional
 
+from typing_extensions import override
+
 from apache_beam import coders
 from apache_beam.coders.coders import BooleanCoder
 from apache_beam.coders.coders import Coder
@@ -994,3 +996,141 @@ def expand(self, pbegin):
 
   def _infer_output_coder(self, input_type=None, input_coder=None):
     return self._source.default_output_coder()
+
+
+# ------------------------------------------------------------------------------
+# Internal in-memory counting source, shared by Beam's own tests and examples.
+# ------------------------------------------------------------------------------
+
+# Realistic event-time base away from the Unix epoch.
+_EVENT_TIME_BASE = Timestamp(1729987200)  # 2024-10-27T00:00:00Z
+
+
+class _CountingCheckpointMark(CheckpointMark):
+  def __init__(self, last_index, finalize_log=None):
+    self.last_index = last_index
+    self._finalize_log = finalize_log
+
+  @override
+  def finalize_checkpoint(self):
+    if self._finalize_log is not None:
+      self._finalize_log.append(self.last_index)
+
+  def __eq__(self, other):
+    return (
+        isinstance(other, _CountingCheckpointMark) and
+        other.last_index == self.last_index)
+
+  def __hash__(self):
+    return hash(self.last_index)
+
+  def __repr__(self):
+    return '_CountingCheckpointMark(last_index=%r)' % (self.last_index, )
+
+
+class _CountingReader(UnboundedReader):
+  def __init__(
+      self, count, start_index, finalize_log=None, modulus=1, residue=0):
+    self._count = count
+    self._next = start_index
+    self._modulus = modulus
+    self._residue = residue
+    self._current = None
+    self._exhausted = False
+    self._finalize_log = finalize_log
+    self.closed = False
+
+  def _read_next(self):
+    while self._next < self._count:
+      index = self._next
+      self._next += 1
+      if index % self._modulus == self._residue:
+        self._current = index
+        return True
+    self._exhausted = True
+    return False
+
+  @override
+  def start(self):
+    return self._read_next()
+
+  @override
+  def advance(self):
+    return self._read_next()
+
+  @override
+  def get_current(self):
+    return self._current
+
+  @override
+  def get_current_timestamp(self):
+    return _EVENT_TIME_BASE + self._current
+
+  @override
+  def get_watermark(self):
+    if self._exhausted:
+      return MAX_TIMESTAMP
+    if self._current is None:
+      return MIN_TIMESTAMP
+    return _EVENT_TIME_BASE + self._current
+
+  @override
+  def get_checkpoint_mark(self):
+    last = self._current if self._current is not None else self._next - 1
+    return _CountingCheckpointMark(last, finalize_log=self._finalize_log)
+
+  @override
+  def close(self):
+    self.closed = True
+
+
+class UnboundedCountingSource(UnboundedSource):
+  """In-memory counting source for Beam's own tests and examples.
+
+  Internal: not part of the public ``UnboundedSource`` API. Emits integers
+  ``0..count-1`` with event time ``_EVENT_TIME_BASE + index``, self-terminates
+  at EOF, resumes from ``last_index + 1``, and optionally splits into
+  independent even/odd sub-sources.
+  """
+  def __init__(
+      self,
+      count,
+      finalize_log=None,
+      is_splittable=False,
+      modulus=1,
+      residue=0):
+    self._count = count
+    self._finalize_log = finalize_log
+    self._is_splittable = is_splittable
+    self._modulus = modulus
+    self._residue = residue
+    self.last_reader = None
+
+  @override
+  def split(self, desired_num_splits, options=None):
+    if not self._is_splittable or desired_num_splits < 2:
+      return [self]
+    # Split into independent even/odd sub-sources (each non-splittable).
+    return [
+        UnboundedCountingSource(
+            self._count,
+            finalize_log=self._finalize_log,
+            modulus=2,
+            residue=residue) for residue in (0, 1)
+    ]
+
+  @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 = _CountingReader(
+        self._count,
+        start_index,
+        finalize_log=self._finalize_log,
+        modulus=self._modulus,
+        residue=self._residue)
+    return self.last_reader
+
+  @override
+  def get_checkpoint_mark_coder(self):
+    return coders.PickleCoder()

sdks/python/apache_beam/io/unbounded_source_test.py

@@ -29,11 +29,15 @@
 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
@@ -55,138 +59,6 @@
 
 # pylint: disable=expression-not-assigned
 
-# Realistic event-time base away from the Unix epoch.
-_EVENT_TIME_BASE = Timestamp(1729987200)  # 2024-10-27T00:00:00Z
-
-# ------------------------------------------------------------------------------
-# In-memory demo source emitting integers 0..count-1 with event time
-# ``_EVENT_TIME_BASE + index``. It self-terminates at EOF, resumes from
-# ``last_index + 1``, and splits into even/odd sub-sources when configured.
-# ------------------------------------------------------------------------------
-
-
-class _CountingCheckpointMark(CheckpointMark):
-  def __init__(self, last_index, finalize_log=None):
-    self.last_index = last_index
-    self._finalize_log = finalize_log
-
-  @override
-  def finalize_checkpoint(self):
-    if self._finalize_log is not None:
-      self._finalize_log.append(self.last_index)
-
-  def __eq__(self, other):
-    return (
-        isinstance(other, _CountingCheckpointMark) and
-        other.last_index == self.last_index)
-
-  def __hash__(self):
-    return hash(self.last_index)
-
-  def __repr__(self):
-    return '_CountingCheckpointMark(last_index=%r)' % (self.last_index, )
-
-
-class _CountingReader(UnboundedReader):
-  def __init__(
-      self, count, start_index, finalize_log=None, modulus=1, residue=0):
-    self._count = count
-    self._next = start_index
-    self._modulus = modulus
-    self._residue = residue
-    self._current = None
-    self._exhausted = False
-    self._finalize_log = finalize_log
-    self.closed = False
-
-  def _read_next(self):
-    while self._next < self._count:
-      index = self._next
-      self._next += 1
-      if index % self._modulus == self._residue:
-        self._current = index
-        return True
-    self._exhausted = True
-    return False
-
-  @override
-  def start(self):
-    return self._read_next()
-
-  @override
-  def advance(self):
-    return self._read_next()
-
-  @override
-  def get_current(self):
-    return self._current
-
-  @override
-  def get_current_timestamp(self):
-    return _EVENT_TIME_BASE + self._current
-
-  @override
-  def get_watermark(self):
-    if self._exhausted:
-      return MAX_TIMESTAMP
-    if self._current is None:
-      return MIN_TIMESTAMP
-    return _EVENT_TIME_BASE + self._current
-
-  @override
-  def get_checkpoint_mark(self):
-    last = self._current if self._current is not None else self._next - 1
-    return _CountingCheckpointMark(last, finalize_log=self._finalize_log)
-
-  @override
-  def close(self):
-    self.closed = True
-
-
-class UnboundedCountingSource(UnboundedSource):
-  def __init__(
-      self,
-      count,
-      finalize_log=None,
-      is_splittable=False,
-      modulus=1,
-      residue=0):
-    self._count = count
-    self._finalize_log = finalize_log
-    self._is_splittable = is_splittable
-    self._modulus = modulus
-    self._residue = residue
-    self.last_reader = None
-
-  @override
-  def split(self, desired_num_splits, options=None):
-    if not self._is_splittable or desired_num_splits < 2:
-      return [self]
-    # Split into independent even/odd sub-sources (each non-splittable).
-    return [
-        UnboundedCountingSource(
-            self._count,
-            finalize_log=self._finalize_log,
-            modulus=2,
-            residue=residue) for residue in (0, 1)
-    ]
-
-  @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 = _CountingReader(
-        self._count,
-        start_index,
-        finalize_log=self._finalize_log,
-        modulus=self._modulus,
-        residue=self._residue)
-    return self.last_reader
-
-  @override
-  def get_checkpoint_mark_coder(self):
-    return coders.PickleCoder()
-
 
 class _StringCountingReader(_CountingReader):
   @override