#
# 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 side inputs."""
# pytype: skip-file
from __future__ import absolute_import
import logging
import unittest
import apache_beam as beam
from apache_beam.testing.test_pipeline import TestPipeline
from apache_beam.testing.util import assert_that
from apache_beam.testing.util import equal_to
[docs]class KeywordOnlyArgsTests(unittest.TestCase):
# Enable nose tests running in parallel
_multiprocess_can_split_ = True
[docs] def test_combine_keyword_only_args(self):
with TestPipeline() as pipeline:
def bounded_sum(values, *s, bound=500):
return min(sum(values) + sum(s), bound)
def bounded_sum_without_default_values(values, *s, bound):
return min(sum(values) + sum(s), bound)
pcoll = pipeline | 'start' >> beam.Create([6, 3, 1])
result1 = pcoll | 'sum1' >> beam.CombineGlobally(
bounded_sum, 5, 8, bound=20)
result2 = pcoll | 'sum2' >> beam.CombineGlobally(bounded_sum, 0, 0)
result3 = pcoll | 'sum3' >> beam.CombineGlobally(bounded_sum)
result4 = pcoll | 'sum4' >> beam.CombineGlobally(bounded_sum, bound=5)
result5 = pcoll | 'sum5' >> beam.CombineGlobally(
bounded_sum_without_default_values, 5, 8, bound=20)
result6 = pcoll | 'sum6' >> beam.CombineGlobally(
bounded_sum_without_default_values, 0, 0, bound=500)
result7 = pcoll | 'sum7' >> beam.CombineGlobally(
bounded_sum_without_default_values, bound=500)
result8 = pcoll | 'sum8' >> beam.CombineGlobally(
bounded_sum_without_default_values, bound=5)
assert_that(result1, equal_to([20]), label='assert1')
assert_that(result2, equal_to([10]), label='assert2')
assert_that(result3, equal_to([10]), label='assert3')
assert_that(result4, equal_to([5]), label='assert4')
assert_that(result5, equal_to([20]), label='assert5')
assert_that(result6, equal_to([10]), label='assert6')
assert_that(result7, equal_to([10]), label='assert7')
assert_that(result8, equal_to([5]), label='assert8')
[docs] def test_do_fn_keyword_only_args(self):
with TestPipeline() as pipeline:
class MyDoFn(beam.DoFn):
def process(self, element, *s, bound=500):
return [min(sum(s) + element, bound)]
pcoll = pipeline | 'start' >> beam.Create([6, 3, 1])
result1 = pcoll | 'sum1' >> beam.ParDo(MyDoFn(), 5, 8, bound=15)
result2 = pcoll | 'sum2' >> beam.ParDo(MyDoFn(), 5, 8)
result3 = pcoll | 'sum3' >> beam.ParDo(MyDoFn())
result4 = pcoll | 'sum4' >> beam.ParDo(MyDoFn(), bound=5)
assert_that(result1, equal_to([15, 15, 14]), label='assert1')
assert_that(result2, equal_to([19, 16, 14]), label='assert2')
assert_that(result3, equal_to([6, 3, 1]), label='assert3')
assert_that(result4, equal_to([5, 3, 1]), label='assert4')
if __name__ == '__main__':
logging.getLogger().setLevel(logging.DEBUG)
unittest.main()