github.com/apache/beam/sdks/v2@v2.48.2/python/apache_beam/examples/streaming_wordcount_debugging_test.py (about)

     1  #
     2  # Licensed to the Apache Software Foundation (ASF) under one or more
     3  # contributor license agreements.  See the NOTICE file distributed with
     4  # this work for additional information regarding copyright ownership.
     5  # The ASF licenses this file to You under the Apache License, Version 2.0
     6  # (the "License"); you may not use this file except in compliance with
     7  # the License.  You may obtain a copy of the License at
     8  #
     9  #    http://www.apache.org/licenses/LICENSE-2.0
    10  #
    11  # Unless required by applicable law or agreed to in writing, software
    12  # distributed under the License is distributed on an "AS IS" BASIS,
    13  # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    14  # See the License for the specific language governing permissions and
    15  # limitations under the License.
    16  #
    17  
    18  """Unit test for the streaming wordcount example with debug."""
    19  
    20  # pytype: skip-file
    21  
    22  import unittest
    23  
    24  import mock
    25  
    26  import apache_beam as beam
    27  from apache_beam.examples import streaming_wordcount_debugging
    28  from apache_beam.testing.test_stream import TestStream
    29  from apache_beam.testing.util import assert_that
    30  from apache_beam.testing.util import equal_to
    31  
    32  # Protect against environments where the PubSub library is not available.
    33  # pylint: disable=wrong-import-order, wrong-import-position
    34  try:
    35    from google.cloud import pubsub
    36  except ImportError:
    37    pubsub = None
    38  # pylint: enable=wrong-import-order, wrong-import-position
    39  
    40  
    41  class StreamingWordcountDebugging(unittest.TestCase):
    42    @unittest.skipIf(pubsub is None, 'GCP dependencies are not installed')
    43    @mock.patch('apache_beam.io.ReadFromPubSub')
    44    @mock.patch('apache_beam.io.WriteToPubSub')
    45    def test_streaming_wordcount_debugging(self, *unused_mocks):
    46      def FakeReadFromPubSub(topic=None, subscription=None, values=None):
    47        expected_topic = topic
    48        expected_subscription = subscription
    49  
    50        def _inner(topic=None, subscription=None):
    51          assert topic == expected_topic
    52          assert subscription == expected_subscription
    53          return TestStream().add_elements(values)
    54  
    55        return _inner
    56  
    57      class AssertTransform(beam.PTransform):
    58        def __init__(self, matcher):
    59          self.matcher = matcher
    60  
    61        def expand(self, pcoll):
    62          assert_that(pcoll, self.matcher)
    63  
    64      def FakeWriteToPubSub(topic=None, values=None):
    65        expected_topic = topic
    66  
    67        def _inner(topic=None, subscription=None):
    68          assert topic == expected_topic
    69          return AssertTransform(equal_to(values))
    70  
    71        return _inner
    72  
    73      input_topic = 'projects/fake-beam-test-project/topic/intopic'
    74      input_values = [
    75          '150', '151', '152', '153', '154', '210', '211', '212', '213', '214'
    76      ]
    77      output_topic = 'projects/fake-beam-test-project/topic/outtopic'
    78      output_values = [
    79          '150: 1',
    80          '151: 1',
    81          '152: 1',
    82          '153: 1',
    83          '154: 1',
    84          '210: 1',
    85          '211: 1',
    86          '212: 1',
    87          '213: 1',
    88          '214: 1'
    89      ]
    90      beam.io.ReadFromPubSub = (
    91          FakeReadFromPubSub(
    92              topic=input_topic,
    93              values=list(x.encode('utf-8') for x in input_values)))
    94      beam.io.WriteToPubSub = (
    95          FakeWriteToPubSub(
    96              topic=output_topic,
    97              values=list(x.encode('utf-8') for x in output_values)))
    98      streaming_wordcount_debugging.run([
    99          '--input_topic',
   100          'projects/fake-beam-test-project/topic/intopic',
   101          '--output_topic',
   102          'projects/fake-beam-test-project/topic/outtopic'
   103      ],
   104                                        save_main_session=False)
   105  
   106  
   107  if __name__ == '__main__':
   108    unittest.main()