github.com/apache/beam/sdks/v2@v2.48.2/python/apache_beam/transforms/external_it_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  """Integration tests for cross-language transform expansion."""
    19  
    20  # pytype: skip-file
    21  
    22  import unittest
    23  
    24  import pytest
    25  
    26  import apache_beam as beam
    27  from apache_beam import Pipeline
    28  from apache_beam.runners.portability import expansion_service
    29  from apache_beam.testing.test_pipeline import TestPipeline
    30  from apache_beam.testing.util import assert_that
    31  from apache_beam.testing.util import equal_to
    32  from apache_beam.transforms import ptransform
    33  
    34  
    35  class ExternalTransformIT(unittest.TestCase):
    36    @pytest.mark.it_postcommit
    37    def test_job_python_from_python_it(self):
    38      @ptransform.PTransform.register_urn('simple', None)
    39      class SimpleTransform(ptransform.PTransform):
    40        def expand(self, pcoll):
    41          return pcoll | beam.Map(lambda x: 'Simple(%s)' % x)
    42  
    43        def to_runner_api_parameter(self, unused_context):
    44          return 'simple', None
    45  
    46        @staticmethod
    47        def from_runner_api_parameter(_0, _1, _2):
    48          return SimpleTransform()
    49  
    50      pipeline = TestPipeline(is_integration_test=True)
    51  
    52      res = (
    53          pipeline
    54          | beam.Create(['a', 'b'])
    55          | beam.ExternalTransform(
    56              'simple', None, expansion_service.ExpansionServiceServicer()))
    57      assert_that(res, equal_to(['Simple(a)', 'Simple(b)']))
    58  
    59      proto_pipeline, _ = pipeline.to_runner_api(return_context=True)
    60      pipeline_from_proto = Pipeline.from_runner_api(
    61          proto_pipeline, pipeline.runner, pipeline._options)
    62      pipeline_from_proto.run().wait_until_finish()
    63  
    64  
    65  if __name__ == '__main__':
    66    unittest.main()