github.com/apache/beam/sdks/v2@v2.48.2/python/apache_beam/transforms/create_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 the Create and _CreateSource classes."""
# pytype: skip-file

import logging
import unittest

from apache_beam import Create
from apache_beam import coders
from apache_beam.coders import FastPrimitivesCoder
from apache_beam.internal import pickler
from apache_beam.io import source_test_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


class CreateTest(unittest.TestCase):
  def setUp(self):
    self.coder = FastPrimitivesCoder()

  def test_create_transform(self):
    with TestPipeline() as p:
      assert_that(p | 'Empty' >> Create([]), equal_to([]), label='empty')
      assert_that(p | 'One' >> Create([None]), equal_to([None]), label='one')
      assert_that(p | Create(list(range(10))), equal_to(list(range(10))))

  def test_create_source_read(self):
    self.check_read([], self.coder)
    self.check_read([1], self.coder)
    # multiple values.
    self.check_read(list(range(10)), self.coder)

  def check_read(self, values, coder):
    source = Create._create_source_from_iterable(values, coder)
    read_values = source_test_utils.read_from_source(source)
    self.assertEqual(sorted(values), sorted(read_values))

  def test_create_source_read_with_initial_splits(self):
    self.check_read_with_initial_splits([], self.coder, num_splits=2)
    self.check_read_with_initial_splits([1], self.coder, num_splits=2)
    values = list(range(8))
    # multiple values with a single split.
    self.check_read_with_initial_splits(values, self.coder, num_splits=1)
    # multiple values with a single split with a large desired bundle size
    self.check_read_with_initial_splits(values, self.coder, num_splits=0.5)
    # multiple values with many splits.
    self.check_read_with_initial_splits(values, self.coder, num_splits=3)
    # multiple values with uneven sized splits.
    self.check_read_with_initial_splits(values, self.coder, num_splits=4)
    # multiple values with num splits equal to num values.
    self.check_read_with_initial_splits(
        values, self.coder, num_splits=len(values))
    # multiple values with num splits greater than to num values.
    self.check_read_with_initial_splits(values, self.coder, num_splits=30)

  def check_read_with_initial_splits(self, values, coder, num_splits):
    """A test that splits the given source into `num_splits` and verifies that
    the data read from original source is equal to the union of the data read
    from the split sources.
    """
    source = Create._create_source_from_iterable(values, coder)
    desired_bundle_size = source._total_size // num_splits
    splits = source.split(desired_bundle_size)
    splits_info = [(split.source, split.start_position, split.stop_position)
                   for split in splits]
    source_test_utils.assert_sources_equal_reference_source(
        (source, None, None), splits_info)

  def test_create_source_read_reentrant(self):
    source = Create._create_source_from_iterable(range(9), self.coder)
    source_test_utils.assert_reentrant_reads_succeed((source, None, None))

  def test_create_source_read_reentrant_with_initial_splits(self):
    source = Create._create_source_from_iterable(range(24), self.coder)
    for split in source.split(desired_bundle_size=5):
      source_test_utils.assert_reentrant_reads_succeed(
          (split.source, split.start_position, split.stop_position))

  def test_create_source_dynamic_splitting(self):
    # 2 values
    source = Create._create_source_from_iterable(range(2), self.coder)
    source_test_utils.assert_split_at_fraction_exhaustive(source)
    # Multiple values.
    source = Create._create_source_from_iterable(range(11), self.coder)
    source_test_utils.assert_split_at_fraction_exhaustive(
        source, perform_multi_threaded_test=True)

  def test_create_source_progress(self):
    num_values = 10
    source = Create._create_source_from_iterable(range(num_values), self.coder)
    splits = [split for split in source.split(desired_bundle_size=100)]
    assert len(splits) == 1
    fraction_consumed_report = []
    split_points_report = []
    range_tracker = splits[0].source.get_range_tracker(
        splits[0].start_position, splits[0].stop_position)
    for _ in splits[0].source.read(range_tracker):
      fraction_consumed_report.append(range_tracker.fraction_consumed())
      split_points_report.append(range_tracker.split_points())

    self.assertEqual([float(i) / num_values for i in range(num_values)],
                     fraction_consumed_report)

    expected_split_points_report = [((i - 1), num_values - (i - 1))
                                    for i in range(1, num_values + 1)]

    self.assertEqual(expected_split_points_report, split_points_report)

  def test_create_uses_coder_for_pickling(self):
    coders.registry.register_coder(_Unpicklable, _UnpicklableCoder)
    create = Create([_Unpicklable(1), _Unpicklable(2), _Unpicklable(3)])
    unpickled_create = pickler.loads(pickler.dumps(create))
    self.assertEqual(
        sorted(create.values, key=lambda v: v.value),
        sorted(unpickled_create.values, key=lambda v: v.value))

    with self.assertRaises(NotImplementedError):
      # As there is no special coder for Union types, this will fall back to
      # FastPrimitivesCoder, which in turn falls back to pickling.
      create_mixed_types = Create([_Unpicklable(1), 2])
      pickler.dumps(create_mixed_types)


class _Unpicklable(object):
  def __init__(self, value):
    self.value = value

  def __eq__(self, other):
    return self.value == other.value

  def __getstate__(self):
    raise NotImplementedError()

  def __setstate__(self, state):
    raise NotImplementedError()


class _UnpicklableCoder(coders.Coder):
  def encode(self, value):
    return str(value.value).encode()

  def decode(self, encoded):
    return _Unpicklable(int(encoded.decode()))

  def to_type_hint(self):
    return _Unpicklable


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