#

# 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 core python file."""

# pytype: skip-file

import logging

import os

import tempfile

import unittest

from typing import Iterable

from typing import Literal

from typing import TypeVar

import pytest

import apache_beam as beam

from apache_beam.coders import coders

from apache_beam.options.pipeline_options import PipelineOptions

from apache_beam.testing.util import assert_that

from apache_beam.testing.util import equal_to

from apache_beam.transforms.core import DeadLetter

from apache_beam.transforms.resources import ResourceHint

from apache_beam.transforms.userstate import BagStateSpec

from apache_beam.transforms.userstate import ReadModifyWriteStateSpec

from apache_beam.transforms.userstate import TimerSpec

from apache_beam.transforms.userstate import on_timer

from apache_beam.transforms.window import FixedWindows

from apache_beam.typehints import TypeCheckError

from apache_beam.typehints import row_type

from apache_beam.typehints import typehints

RETURN_NONE_PARTIAL_WARNING = "Process method returned None"

class TestDoFn0(beam.DoFn):

"""Returning without a value is allowed"""

def process(self, element):

if not element:

return

yield element

class TestDoFn1(beam.DoFn):

def process(self, element):

yield element

class TestDoFn2(beam.DoFn):

def process(self, element):

def inner_func(x):

yield x

return inner_func(element)

class TestDoFn3(beam.DoFn):

"""mixing return and yield is not allowed"""

def process(self, element):

if not element:

return -1

yield element

class TestDoFn4(beam.DoFn):

"""test the variable name containing return"""

def process(self, element):

my_return = element

yield my_return

class TestDoFn5(beam.DoFn):

"""test the variable name containing yield"""

def process(self, element):

my_yield = element

return my_yield

class TestDoFn6(beam.DoFn):

"""test the variable name containing return"""

def process(self, element):

return_test = element

yield return_test

class TestDoFn7(beam.DoFn):

"""test the variable name containing yield"""

def process(self, element):

yield_test = element

return yield_test

class TestDoFn8(beam.DoFn):

"""test the code containing yield and yield from"""

def process(self, element):

if not element:

yield from [1, 2, 3]

else:

yield element

class TestDoFn9(beam.DoFn):

def process(self, element):

if len(element) > 3:

raise ValueError('Not allowed to have long elements')

yield element

class TestDoFn10(beam.DoFn):

"""test process returning None explicitly"""

def process(self, element):

return None

class TestDoFn11(beam.DoFn):

"""test process returning None (no return and no yield)"""

def process(self, element):

pass

class TestDoFn12(beam.DoFn):

"""test process returning None in a filter pattern"""

def process(self, element):

if element == 0:

return

return element

class TestDoFnStateful(beam.DoFn):

STATE_SPEC = ReadModifyWriteStateSpec('num_elements', coders.VarIntCoder())

"""test process with a stateful dofn"""

def process(self, element, state=beam.DoFn.StateParam(STATE_SPEC)):

if len(element[1]) > 3:

raise ValueError('Not allowed to have long elements')

current_value = state.read() or 1

state.write(current_value + 1)

yield current_value

class TestDoFnWithTimer(beam.DoFn):

ALL_ELEMENTS = BagStateSpec('buffer', coders.VarIntCoder())

TIMER = TimerSpec('timer', beam.TimeDomain.WATERMARK)

"""test process with a stateful dofn"""

def process(

self,

element,

t=beam.DoFn.TimestampParam,

state=beam.DoFn.StateParam(ALL_ELEMENTS),

timer=beam.DoFn.TimerParam(TIMER)):

if element[1] > 3:

raise ValueError('Not allowed to have large numbers')

state.add(element[1])

timer.set(t)

return []

@on_timer(TIMER)

def expiry_callback(self, state=beam.DoFn.StateParam(ALL_ELEMENTS)):

unique_elements = list(state.read())

state.clear()

return unique_elements

class CreateTest(unittest.TestCase):

@pytest.fixture(autouse=True)

def inject_fixtures(self, caplog):

self._caplog = caplog

def test_dofn_with_yield_and_return(self):

warning_text = 'Using yield and return'

with self._caplog.at_level(logging.WARNING):

assert beam.ParDo(sum)

assert beam.ParDo(TestDoFn0())

assert beam.ParDo(TestDoFn1())

assert beam.ParDo(TestDoFn2())

assert beam.ParDo(TestDoFn4())

assert beam.ParDo(TestDoFn5())

assert beam.ParDo(TestDoFn6())

assert beam.ParDo(TestDoFn7())

assert beam.ParDo(TestDoFn8())

assert warning_text not in self._caplog.text

with self._caplog.at_level(logging.WARNING):

beam.ParDo(TestDoFn3())

assert warning_text in self._caplog.text

def test_dofn_with_explicit_return_none(self):

with self._caplog.at_level(logging.WARNING):

beam.ParDo(TestDoFn10())

assert RETURN_NONE_PARTIAL_WARNING in self._caplog.text

assert str(TestDoFn10) in self._caplog.text

def test_dofn_with_implicit_return_none_missing_return_and_yield(self):

with self._caplog.at_level(logging.WARNING):

beam.ParDo(TestDoFn11())

assert RETURN_NONE_PARTIAL_WARNING not in self._caplog.text

def test_dofn_with_implicit_return_none_and_value(self):

with self._caplog.at_level(logging.WARNING):

beam.ParDo(TestDoFn12())

assert RETURN_NONE_PARTIAL_WARNING not in self._caplog.text

class PartitionTest(unittest.TestCase):

def test_partition_with_bools(self):

with pytest.raises(

(ValueError, RuntimeError),

match=

r"PartitionFn yielded a '([^']*)' when it should only yield integers"):

# Check for RuntimeError too since the portable runner casts

# all exceptions to RuntimeError

invalid_inputs = [True, 1.2, 'string', None]

for input_value in invalid_inputs:

with beam.testing.test_pipeline.TestPipeline() as p:

_ = (

p | beam.Create([input_value])

| beam.Partition(lambda x, _: x, 2))

def test_partition_with_numpy_integers(self):

# Test that numpy integer types are correctly accepted by the

# ApplyPartitionFnFn class

import numpy as np

# Create an instance of the ApplyPartitionFnFn class

apply_partition_fn = beam.Partition.ApplyPartitionFnFn()

# Define a simple partition function

class SimplePartitionFn(beam.PartitionFn):

def partition_for(self, element, num_partitions):

return element % num_partitions

partition_fn = SimplePartitionFn()

# Test with numpy.int32

# This should not raise an exception

outputs = list(apply_partition_fn.process(np.int32(1), partition_fn, 3))

self.assertEqual(len(outputs), 1)

self.assertEqual(outputs[0].tag, '1') # 1 % 3 = 1

# Test with numpy.int64

# This should not raise an exception

outputs = list(apply_partition_fn.process(np.int64(2), partition_fn, 3))

self.assertEqual(len(outputs), 1)

self.assertEqual(outputs[0].tag, '2') # 2 % 3 = 2

def test_partition_fn_returning_numpy_integers(self):

# Test that partition functions can return numpy integer types

import numpy as np

# Create an instance of the ApplyPartitionFnFn class

apply_partition_fn = beam.Partition.ApplyPartitionFnFn()

# Define partition functions that return numpy integer types

class Int32PartitionFn(beam.PartitionFn):

def partition_for(self, element, num_partitions):

return np.int32(element % num_partitions)

class Int64PartitionFn(beam.PartitionFn):

def partition_for(self, element, num_partitions):

return np.int64(element % num_partitions)

# Test with partition function returning numpy.int32

# This should not raise an exception

outputs = list(apply_partition_fn.process(1, Int32PartitionFn(), 3))

self.assertEqual(len(outputs), 1)

self.assertEqual(outputs[0].tag, '1') # 1 % 3 = 1

# Test with partition function returning numpy.int64

# This should not raise an exception

outputs = list(apply_partition_fn.process(2, Int64PartitionFn(), 3))

self.assertEqual(len(outputs), 1)

self.assertEqual(outputs[0].tag, '2') # 2 % 3 = 2

def test_partition_boundedness(self):

def partition_fn(val, num_partitions):

return val % num_partitions

class UnboundedDoFn(beam.DoFn):

@beam.DoFn.unbounded_per_element()

def process(self, element):

yield element

with beam.testing.test_pipeline.TestPipeline() as p:

source = p | beam.Create([1, 2, 3, 4, 5])

p1, p2, p3 = source | "bounded" >> beam.Partition(partition_fn, 3)

self.assertEqual(source.is_bounded, True)

self.assertEqual(p1.is_bounded, True)

self.assertEqual(p2.is_bounded, True)

self.assertEqual(p3.is_bounded, True)

unbounded = source | beam.ParDo(UnboundedDoFn())

p4, p5, p6 = unbounded | "unbounded" >> beam.Partition(partition_fn, 3)

self.assertEqual(unbounded.is_bounded, False)

self.assertEqual(p4.is_bounded, False)

self.assertEqual(p5.is_bounded, False)

self.assertEqual(p6.is_bounded, False)

class FlattenTest(unittest.TestCase):

def test_flatten_identical_windows(self):

with beam.testing.test_pipeline.TestPipeline() as p:

source1 = p | "c1" >> beam.Create(

[1, 2, 3, 4, 5]) | "w1" >> beam.WindowInto(FixedWindows(100))

source2 = p | "c2" >> beam.Create([6, 7, 8]) | "w2" >> beam.WindowInto(

FixedWindows(100))

source3 = p | "c3" >> beam.Create([9, 10]) | "w3" >> beam.WindowInto(

FixedWindows(100))

out = (source1, source2, source3) | "flatten" >> beam.Flatten()

assert_that(out, equal_to([1, 2, 3, 4, 5, 6, 7, 8, 9, 10]))

def test_flatten_no_windows(self):

with beam.testing.test_pipeline.TestPipeline() as p:

source1 = p | "c1" >> beam.Create([1, 2, 3, 4, 5])

source2 = p | "c2" >> beam.Create([6, 7, 8])

source3 = p | "c3" >> beam.Create([9, 10])

out = (source1, source2, source3) | "flatten" >> beam.Flatten()

assert_that(out, equal_to([1, 2, 3, 4, 5, 6, 7, 8, 9, 10]))

def test_flatten_mismatched_windows(self):

with beam.testing.test_pipeline.TestPipeline() as p:

source1 = p | "c1" >> beam.Create(

[1, 2, 3, 4, 5]) | "w1" >> beam.WindowInto(FixedWindows(25))

source2 = p | "c2" >> beam.Create([6, 7, 8]) | "w2" >> beam.WindowInto(

FixedWindows(100))

source3 = p | "c3" >> beam.Create([9, 10]) | "w3" >> beam.WindowInto(

FixedWindows(100))

_ = (source1, source2, source3) | "flatten" >> beam.Flatten()

class ExceptionHandlingTest(unittest.TestCase):

def test_routes_failures(self):

with beam.Pipeline() as pipeline:

good, bad = (

pipeline | beam.Create(['abc', 'long_word', 'foo', 'bar', 'foobar'])

| beam.ParDo(TestDoFn9()).with_exception_handling()

)

bad_elements = bad | beam.Keys()

assert_that(good, equal_to(['abc', 'foo', 'bar']), 'good')

assert_that(bad_elements, equal_to(['long_word', 'foobar']), 'bad')

def test_handles_callbacks(self):

with tempfile.TemporaryDirectory() as tmp_dirname:

tmp_path = os.path.join(tmp_dirname, 'tmp_filename')

file_contents = 'random content'

def failure_callback(e, el):

if type(e) is not ValueError:

raise Exception(f'Failed to pass in correct exception, received {e}')

if el != 'foobar':

raise Exception(f'Failed to pass in correct element, received {el}')

f = open(tmp_path, "a")

logging.warning(tmp_path)

f.write(file_contents)

f.close()

with beam.Pipeline() as pipeline:

good, bad = (

pipeline | beam.Create(['abc', 'bcd', 'foo', 'bar', 'foobar'])

| beam.ParDo(TestDoFn9()).with_exception_handling(

on_failure_callback=failure_callback)

)

bad_elements = bad | beam.Keys()

assert_that(good, equal_to(['abc', 'bcd', 'foo', 'bar']), 'good')

assert_that(bad_elements, equal_to(['foobar']), 'bad')

with open(tmp_path) as f:

s = f.read()

self.assertEqual(s, file_contents)

def test_handles_no_callback_triggered(self):

with tempfile.TemporaryDirectory() as tmp_dirname:

tmp_path = os.path.join(tmp_dirname, 'tmp_filename')

file_contents = 'random content'

def failure_callback(e, el):

f = open(tmp_path, "a")

logging.warning(tmp_path)

f.write(file_contents)

f.close()

with beam.Pipeline() as pipeline:

good, bad = (

pipeline | beam.Create(['abc', 'bcd', 'foo', 'bar'])

| beam.ParDo(TestDoFn9()).with_exception_handling(

on_failure_callback=failure_callback)

)

bad_elements = bad | beam.Keys()

assert_that(good, equal_to(['abc', 'bcd', 'foo', 'bar']), 'good')

assert_that(bad_elements, equal_to([]), 'bad')

self.assertFalse(os.path.isfile(tmp_path))

def test_stateful_exception_handling(self):

with beam.Pipeline() as pipeline:

good, bad = (

pipeline | beam.Create([(1, 'abc'), (1, 'long_word'),

(1, 'foo'), (1, 'bar'), (1, 'foobar')])

| beam.ParDo(TestDoFnStateful()).with_exception_handling(

allow_unsafe_userstate_in_process=True)

)

bad_elements = bad | beam.Keys()

assert_that(good, equal_to([1, 2, 3]), 'good')

assert_that(

bad_elements, equal_to([(1, 'long_word'), (1, 'foobar')]), 'bad')

def test_timer_exception_handling(self):

with beam.Pipeline() as pipeline:

good, bad = (

pipeline | beam.Create([(1, 0), (1, 1), (1, 2), (1, 5), (1, 10)])

| beam.ParDo(TestDoFnWithTimer()).with_exception_handling(

allow_unsafe_userstate_in_process=True)

)

bad_elements = bad | beam.Keys()

assert_that(good, equal_to([0, 1, 2]), 'good')

assert_that(bad_elements, equal_to([(1, 5), (1, 10)]), 'bad')

def test_tags_with_exception_handling_then_resource_hint(self):

class TagHint(ResourceHint):

urn = 'beam:resources:tags:v1'

ResourceHint.register_resource_hint('tags', TagHint)

with beam.Pipeline() as pipeline:

ok, unused_errors = (

pipeline

| beam.Create([1])

| beam.Map(lambda x: x)

.with_exception_handling()

.with_resource_hints(tags='test_tag')

)

pd = ok.producer.transform

self.assertIsInstance(pd, beam.transforms.core.ParDo)

while hasattr(pd.fn, 'fn'):

pd = pd.fn

self.assertEqual(

pd.get_resource_hints(),

{'beam:resources:tags:v1': b'test_tag'},

)

def test_tags_with_exception_handling_timeout_then_resource_hint(self):

class TagHint(ResourceHint):

urn = 'beam:resources:tags:v1'

ResourceHint.register_resource_hint('tags', TagHint)

with beam.Pipeline() as pipeline:

ok, unused_errors = (

pipeline

| beam.Create([1])

| beam.Map(lambda x: x)

.with_exception_handling(timeout=1)

.with_resource_hints(tags='test_tag')

)

pd = ok.producer.transform

self.assertIsInstance(pd, beam.transforms.core.ParDo)

while hasattr(pd.fn, 'fn'):

pd = pd.fn

self.assertEqual(

pd.get_resource_hints(),

{'beam:resources:tags:v1': b'test_tag'},

)

def test_tags_with_resource_hint_then_exception_handling(self):

class TagHint(ResourceHint):

urn = 'beam:resources:tags:v1'

ResourceHint.register_resource_hint('tags', TagHint)

with beam.Pipeline() as pipeline:

ok, unused_errors = (

pipeline

| beam.Create([1])

| beam.Map(lambda x: x)

.with_resource_hints(tags='test_tag')

.with_exception_handling()

)

pd = ok.producer.transform

self.assertIsInstance(pd, beam.transforms.core.ParDo)

while hasattr(pd.fn, 'fn'):

pd = pd.fn

self.assertEqual(

pd.get_resource_hints(),

{'beam:resources:tags:v1': b'test_tag'},

)

def test_tags_with_resource_hint_then_exception_handling_timeout(self):

class TagHint(ResourceHint):

urn = 'beam:resources:tags:v1'

ResourceHint.register_resource_hint('tags', TagHint)

with beam.Pipeline() as pipeline:

ok, unused_errors = (

pipeline

| beam.Create([1])

| beam.Map(lambda x: x)

.with_resource_hints(tags='test_tag')

.with_exception_handling(timeout=1)

)

pd = ok.producer.transform

self.assertIsInstance(pd, beam.transforms.core.ParDo)

while hasattr(pd.fn, 'fn'):

pd = pd.fn

self.assertEqual(

pd.get_resource_hints(),

{'beam:resources:tags:v1': b'test_tag'},

)

class ExceptionHandlingWithOutputsTest(unittest.TestCase):

"""Tests for combining with_exception_handling() and with_outputs()."""

def _create_dofn_with_tagged_outputs(self):

"""A DoFn that yields tagged outputs and can raise on even numbers."""

class DoWithFailures(beam.DoFn):

def process(

self, element: int

) -> Iterable[int

| beam.pvalue.TaggedOutput[Literal['threes'], int]

| beam.pvalue.TaggedOutput[Literal['fives'], str]]:

if element % 2 == 0:

raise ValueError(f'Even numbers not allowed {element}')

if element % 3 == 0:

yield beam.pvalue.TaggedOutput('threes', element) # type: ignore[misc]

elif element % 5 == 0:

yield beam.pvalue.TaggedOutput('fives', str(element)) # type: ignore[misc]

else:

yield element

return DoWithFailures()

def test_with_exception_handling_then_with_outputs(self):

"""Direction 1: .with_exception_handling().with_outputs()"""

with beam.Pipeline() as p:

results = (

p

| beam.Create([1, 2, 3, 4, 5, 6, 7])

| beam.ParDo(self._create_dofn_with_tagged_outputs()).

with_exception_handling().with_outputs(

'threes', 'fives', main='main'))

assert_that(results.main, equal_to([1, 7]), 'main')

assert_that(results.threes, equal_to([3]), 'threes')

assert_that(results.fives, equal_to(['5']), 'fives')

bad_elements = results.bad | beam.Keys()

assert_that(bad_elements, equal_to([2, 4, 6]), 'bad')

# Verify type hints from annotations are propagated

self.assertEqual(results.main.element_type, int)

self.assertEqual(results.threes.element_type, int)

self.assertEqual(results.fives.element_type, str)

self.assertEqual(results.bad.element_type, DeadLetter[int])

def test_with_outputs_then_with_exception_handling(self):

"""Direction 2: .with_outputs().with_exception_handling()"""

with beam.Pipeline() as p:

results = (

p

| beam.Create([1, 2, 3, 4, 5, 6, 7])

| beam.ParDo(self._create_dofn_with_tagged_outputs()).with_outputs(

'threes', 'fives', main='main').with_exception_handling())

assert_that(results.main, equal_to([1, 7]), 'main')

assert_that(results.threes, equal_to([3]), 'threes')

assert_that(results.fives, equal_to(['5']), 'fives')

bad_elements = results.bad | beam.Keys()

assert_that(bad_elements, equal_to([2, 4, 6]), 'bad')

# Verify type hints from annotations are propagated

self.assertEqual(results.main.element_type, int)

self.assertEqual(results.threes.element_type, int)

self.assertEqual(results.fives.element_type, str)

self.assertEqual(results.bad.element_type, DeadLetter[int])

def test_with_outputs_then_with_exception_handling_custom_dead_letter_tag(

self):

"""Direction 2 with custom dead_letter_tag."""

with beam.Pipeline() as p:

results = (

p

| beam.Create([1, 2, 3])

| beam.ParDo(self._create_dofn_with_tagged_outputs()).with_outputs(

'threes',

main='main').with_exception_handling(dead_letter_tag='errors'))

assert_that(results.main, equal_to([1]), 'main')

assert_that(results.threes, equal_to([3]), 'threes')

bad_elements = results.errors | beam.Keys()

assert_that(bad_elements, equal_to([2]), 'errors')

self.assertEqual(results.threes.element_type, int)

self.assertEqual(results.errors.element_type, DeadLetter[int])

def test_with_exception_handling_then_with_outputs_custom_dead_letter_tag(

self):

"""Direction 1 with custom dead_letter_tag."""

with beam.Pipeline() as p:

results = (

p

| beam.Create([1, 2, 3])

| beam.ParDo(

self._create_dofn_with_tagged_outputs()).with_exception_handling(

dead_letter_tag='errors').with_outputs('threes', main='main'))

assert_that(results.main, equal_to([1]), 'main')

assert_that(results.threes, equal_to([3]), 'threes')

bad_elements = results.errors | beam.Keys()

assert_that(bad_elements, equal_to([2]), 'errors')

self.assertEqual(results.threes.element_type, int)

self.assertEqual(results.errors.element_type, DeadLetter[int])

def test_exception_handling_no_with_outputs_backward_compat(self):

"""Without with_outputs(), behavior is unchanged."""

with beam.Pipeline() as p:

good, bad = (

p

| beam.Create([1, 2, 7])

| beam.ParDo(self._create_dofn_with_tagged_outputs())

.with_exception_handling())

assert_that(good, equal_to([1, 7]), 'good')

bad_elements = bad | beam.Keys()

assert_that(bad_elements, equal_to([2]), 'bad')

def test_exception_handling_compat_version_uses_old_behavior(self):

"""With compat version < 2.73.0, old expand path is used."""

options = PipelineOptions(update_compatibility_version="2.72.0")

with beam.Pipeline(options=options) as p:

good, bad = (

p

| beam.Create([1, 2, 7])

| beam.ParDo(self._create_dofn_with_tagged_outputs())

.with_exception_handling())

assert_that(good, equal_to([1, 7]), 'good')

bad_elements = bad | beam.Keys()

assert_that(bad_elements, equal_to([2]), 'bad')

def test_exception_handling_compat_version_element_type_set_manually(self):

"""With compat version < 2.73.0, element_type is set via manual override

(the old behavior) rather than via with_output_types."""

options = PipelineOptions(update_compatibility_version="2.72.0")

with beam.Pipeline(options=options) as p:

results = (

p

| beam.Create([1, 2, 3])

| beam.ParDo(self._create_dofn_with_tagged_outputs()).

with_exception_handling().with_outputs('threes', main='main'))

# In old path, dead letter type is Any (no with_output_types call)

self.assertEqual(results.bad.element_type, typehints.Any)

# Tagged outputs still get types from DoFn Literal annotations

# (via DoOutputsTuple.__getitem__ reading tagged_output_types)

self.assertEqual(results.threes.element_type, int)

# Main output type should still be inferred via manual override

assert_that(results.main, equal_to([1]), 'main')

def test_with_outputs_then_exception_handling_with_map(self):

"""with_outputs().with_exception_handling() also works on Map."""

with beam.Pipeline() as p:

results = (

p

| beam.Create([1, 2, 3, 4, 5])

| beam.Map(lambda x: x if x % 2 != 0 else 1 / 0).with_outputs(

main='main').with_exception_handling())

assert_that(results.main, equal_to([1, 3, 5]), 'main')

bad_elements = results.bad | beam.Keys()

assert_that(bad_elements, equal_to([2, 4]), 'bad')

def test_with_output_types_chained_on_pardo(self):

"""When type hints are chained on the ParDo (not annotations on the DoFn),

tagged output types should still be propagated through

with_exception_handling().with_outputs()."""

class DoWithFailuresNoAnnotations(beam.DoFn):

def process(self, element):

if element % 2 == 0:

raise ValueError(f'Even numbers not allowed {element}')

if element % 3 == 0:

yield beam.pvalue.TaggedOutput('threes', element)

else:

yield element

with beam.Pipeline() as p:

results = (

p

| beam.Create([1, 2, 3, 7])

| beam.ParDo(DoWithFailuresNoAnnotations()).with_output_types(

int, threes=int).with_exception_handling().with_outputs(

'threes', main='main'))

assert_that(results.main, equal_to([1, 7]), 'main')

assert_that(results.threes, equal_to([3]), 'threes')

bad_elements = results.bad | beam.Keys()

assert_that(bad_elements, equal_to([2]), 'bad')

self.assertEqual(results.main.element_type, int)

self.assertEqual(results.threes.element_type, int)

def test_with_outputs_and_error_handler(self):

"""with_outputs() + error_handler should return DoOutputsTuple, not a

bare PCollection."""

from apache_beam.transforms.error_handling import ErrorHandler

with beam.Pipeline() as p:

with ErrorHandler(beam.Map(lambda x: x)) as handler:

results = (

p

| beam.Create([1, 2, 3, 4, 5, 6, 7])

| beam.ParDo(self._create_dofn_with_tagged_outputs()).with_outputs(

'threes', 'fives',

main='main').with_exception_handling(error_handler=handler))

assert_that(results.main, equal_to([1, 7]), 'main')

assert_that(results.threes, equal_to([3]), 'threes')

assert_that(results.fives, equal_to(['5']), 'fives')

def test_callablewrapper_typehint():

T = TypeVar("T")

def identity(x: T) -> T:

return x

dofn = beam.core.CallableWrapperDoFn(identity)

assert dofn.get_type_hints().strip_iterable()[1][0][0] == typehints.Any

class FlatMapTest(unittest.TestCase):

def test_default(self):

with beam.Pipeline() as pipeline:

letters = (

pipeline

| beam.Create(['abc', 'def'], reshuffle=False)

| beam.FlatMap())

assert_that(letters, equal_to(['a', 'b', 'c', 'd', 'e', 'f']))

def test_default_identity_function_with_typehint(self):

with beam.Pipeline() as pipeline:

letters = (

pipeline

| beam.Create([["abc"]], reshuffle=False)

| beam.FlatMap()

| beam.Map(lambda s: s.upper()).with_input_types(str))

assert_that(letters, equal_to(["ABC"]))

def test_typecheck_with_default(self):

with pytest.raises(TypeCheckError):

with beam.Pipeline() as pipeline:

_ = (

pipeline

| beam.Create([[1, 2, 3]], reshuffle=False)

| beam.FlatMap()

| beam.Map(lambda s: s.upper()).with_input_types(str))

class CreateInferOutputSchemaTest(unittest.TestCase):

def test_multiple_types_for_field(self):

output_type = beam.Create([beam.Row(a=1),

beam.Row(a='foo')]).infer_output_type(None)

self.assertEqual(

output_type,

row_type.RowTypeConstraint.from_fields([

('a', typehints.Union[int, str])

]))

def test_single_type_for_field(self):

output_type = beam.Create([beam.Row(a=1),

beam.Row(a=2)]).infer_output_type(None)

self.assertEqual(

output_type, row_type.RowTypeConstraint.from_fields([('a', int)]))

def test_optional_type_for_field(self):

output_type = beam.Create([beam.Row(a=1),

beam.Row(a=None)]).infer_output_type(None)

self.assertEqual(

output_type,

row_type.RowTypeConstraint.from_fields([('a', typehints.Optional[int])

]))

def test_none_type_for_field_raises_error(self):

with self.assertRaisesRegex(TypeError,

"('No types found for field %s', 'a')"):

beam.Create([beam.Row(a=None), beam.Row(a=None)]).infer_output_type(None)

if __name__ == '__main__':

logging.getLogger().setLevel(logging.INFO)

unittest.main()