sdks/python/apache_beam/transforms/external_test.py - beam - Git at Google

 #
 # 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 transform.util classes."""

 from __future__ import absolute_import

 import argparse
 import subprocess
 import sys
 import unittest

 import grpc
 from mock import patch
 from past.builtins import unicode

 import apache_beam as beam
 from apache_beam import Pipeline
 from apache_beam.options.pipeline_options import PipelineOptions
 from apache_beam.options.pipeline_options import SetupOptions
 from apache_beam.portability import python_urns
 from apache_beam.runners.portability import expansion_service
 from apache_beam.runners.portability.expansion_service_test import FibTransform
 from apache_beam.testing.util import assert_that
 from apache_beam.testing.util import equal_to

 # Protect against environments where apitools library is not available.
 # pylint: disable=wrong-import-order, wrong-import-position
 try:
   from apache_beam.runners.dataflow.internal import apiclient
 except ImportError:
   apiclient = None
 # pylint: enable=wrong-import-order, wrong-import-position


 class ExternalTransformTest(unittest.TestCase):

   # This will be overwritten if set via a flag.
   expansion_service_jar = None
   expansion_service_port = None

   class _RunWithExpansion(object):

     def __init__(self, port, expansion_service_jar):
       self._port = port
       self._expansion_service_jar = expansion_service_jar

     def __enter__(self):
       if not ExternalTransformTest.expansion_service_jar:
         raise unittest.SkipTest('No expansion service jar provided.')

       # Start the java server and wait for it to be ready.
       self._server = subprocess.Popen(
           ['java', '-jar', self._expansion_service_jar, str(self._port)])

       port = ExternalTransformTest.expansion_service_port or 8091
       address = 'localhost:%s' % str(port)

       with grpc.insecure_channel(address) as channel:
         grpc.channel_ready_future(channel).result()

     def __exit__(self, type, value, traceback):
       self._server.kill()

   def test_pipeline_generation(self):
     pipeline = beam.Pipeline()
     res = (pipeline
            | beam.Create(['a', 'b'])
            | beam.ExternalTransform(
                'simple',
                None,
                expansion_service.ExpansionServiceServicer()))
     assert_that(res, equal_to(['Simple(a)', 'Simple(b)']))

     proto, _ = pipeline.to_runner_api(
         return_context=True)
     pipeline_from_proto = Pipeline.from_runner_api(
         proto, pipeline.runner, pipeline._options)

     # Original pipeline has the un-expanded external transform
     self.assertEqual([], pipeline.transforms_stack[0].parts[1].parts)

     # new pipeline has the expanded external transform
     self.assertNotEqual(
         [], pipeline_from_proto.transforms_stack[0].parts[1].parts)
     self.assertEqual(
         u'ExternalTransform(simple)/TestLabel',
         pipeline_from_proto.transforms_stack[0].parts[1].parts[0].full_label)

   def test_simple(self):
     with beam.Pipeline() as p:
       res = (
           p
           | beam.Create(['a', 'b'])
           | beam.ExternalTransform(
               'simple',
               None,
               expansion_service.ExpansionServiceServicer()))
       assert_that(res, equal_to(['Simple(a)', 'Simple(b)']))

   def test_multi(self):
     with beam.Pipeline() as p:
       main1 = p | 'Main1' >> beam.Create(['a', 'bb'], reshuffle=False)
       main2 = p | 'Main2' >> beam.Create(['x', 'yy', 'zzz'], reshuffle=False)
       side = p | 'Side' >> beam.Create(['s'])
       res = dict(main1=main1, main2=main2, side=side) | beam.ExternalTransform(
           'multi', None, expansion_service.ExpansionServiceServicer())
       assert_that(res['main'], equal_to(['as', 'bbs', 'xs', 'yys', 'zzzs']))
       assert_that(res['side'], equal_to(['ss']), label='CheckSide')

   def test_payload(self):
     with beam.Pipeline() as p:
       res = (
           p
           | beam.Create(['a', 'bb'], reshuffle=False)
           | beam.ExternalTransform(
               'payload', b's',
               expansion_service.ExpansionServiceServicer()))
       assert_that(res, equal_to(['as', 'bbs']))

   def test_nested(self):
     with beam.Pipeline() as p:
       assert_that(p | FibTransform(6), equal_to([8]))

   def test_java_expansion_portable_runner(self):
     pipeline_options = PipelineOptions(
         ['--runner=PortableRunner',
          '--experiments=beam_fn_api',
          '--environment_type=%s' % python_urns.EMBEDDED_PYTHON,
          '--job_endpoint=embed'])

     # We use the save_main_session option because one or more DoFn's in this
     # workflow rely on global context (e.g., a module imported at module level).
     pipeline_options.view_as(SetupOptions).save_main_session = True

     ExternalTransformTest.run_pipeline_with_portable_runner(pipeline_options)

   @unittest.skipIf(apiclient is None, 'GCP dependencies are not installed')
   def test_java_expansion_dataflow(self):
     # This test does not actually running the pipeline in Dataflow. It just
     # tests the translation to a Dataflow job request.

     with patch.object(
         apiclient.DataflowApplicationClient, 'create_job') as mock_create_job:
       port = ExternalTransformTest.expansion_service_port or 8091
       with self._RunWithExpansion(port, self.expansion_service_jar):
         pipeline_options = PipelineOptions(
             ['--runner=DataflowRunner',
              '--project=dummyproject',
              '--experiments=beam_fn_api',
              '--temp_location=gs://dummybucket/'])

         # We use the save_main_session option because one or more DoFn's in this
         # workflow rely on global context (e.g., a module imported at module
         # level).
         pipeline_options.view_as(SetupOptions).save_main_session = True

         # Run a simple count-filtered-letters pipeline.
         self.run_pipeline(pipeline_options, port, False)

         mock_args = mock_create_job.call_args_list
         assert mock_args
         args, kwargs = mock_args[0]
         job = args[0]
         job_str = '%s' % job
         self.assertIn('pytest:beam:transforms:filter_less_than', job_str)

   @staticmethod
   def run_pipeline_with_portable_runner(pipeline_options):

     port = ExternalTransformTest.expansion_service_port or 8091

     with ExternalTransformTest._RunWithExpansion(
         port, ExternalTransformTest.expansion_service_jar):
       # Run a simple count-filtered-letters pipeline.
       ExternalTransformTest.run_pipeline(pipeline_options, port, True)

   @staticmethod
   def run_pipeline(
       pipeline_options, expansion_service_port, wait_until_finish=True):
     # The actual definitions of these transforms is in
     # org.apache.beam.runners.core.construction.TestExpansionService.
     TEST_COUNT_URN = "pytest:beam:transforms:count"
     TEST_FILTER_URN = "pytest:beam:transforms:filter_less_than"

     # Run a simple count-filtered-letters pipeline.
     p = beam.Pipeline(options=pipeline_options)
     address = 'localhost:%s' % str(expansion_service_port)
     res = (
         p
         | beam.Create(list('aaabccxyyzzz'))
         | beam.Map(unicode)
         # TODO(BEAM-6587): Use strings directly rather than ints.
         | beam.Map(lambda x: int(ord(x)))
         | beam.ExternalTransform(TEST_FILTER_URN, b'middle', address)
         | beam.ExternalTransform(TEST_COUNT_URN, None, address)
         # # TODO(BEAM-6587): Remove when above is removed.
         | beam.Map(lambda kv: (chr(kv[0]), kv[1]))
         | beam.Map(lambda kv: '%s: %s' % kv))

     assert_that(res, equal_to(['a: 3', 'b: 1', 'c: 2']))

     result = p.run()
     if wait_until_finish:
       result.wait_until_finish()


 if __name__ == '__main__':
   parser = argparse.ArgumentParser()
   parser.add_argument('--expansion_service_jar')
   parser.add_argument('--expansion_service_port')
   known_args, pipeline_args = parser.parse_known_args(sys.argv)

   if known_args.expansion_service_jar:
     ExternalTransformTest.expansion_service_jar = (
         known_args.expansion_service_jar)
     ExternalTransformTest.expansion_service_port = int(
         known_args.expansion_service_port)
     pipeline_options = PipelineOptions(pipeline_args)
     ExternalTransformTest.run_pipeline_with_portable_runner(pipeline_options)
   else:
     sys.argv = pipeline_args
     unittest.main()
	#
	# 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 transform.util classes."""

	from __future__ import absolute_import

	import argparse
	import subprocess
	import sys
	import unittest

	import grpc
	from mock import patch
	from past.builtins import unicode

	import apache_beam as beam
	from apache_beam import Pipeline
	from apache_beam.options.pipeline_options import PipelineOptions
	from apache_beam.options.pipeline_options import SetupOptions
	from apache_beam.portability import python_urns
	from apache_beam.runners.portability import expansion_service
	from apache_beam.runners.portability.expansion_service_test import FibTransform
	from apache_beam.testing.util import assert_that
	from apache_beam.testing.util import equal_to

	# Protect against environments where apitools library is not available.
	# pylint: disable=wrong-import-order, wrong-import-position
	try:
	from apache_beam.runners.dataflow.internal import apiclient
	except ImportError:
	apiclient = None
	# pylint: enable=wrong-import-order, wrong-import-position


	class ExternalTransformTest(unittest.TestCase):

	# This will be overwritten if set via a flag.
	expansion_service_jar = None
	expansion_service_port = None

	class _RunWithExpansion(object):

	def __init__(self, port, expansion_service_jar):
	self._port = port
	self._expansion_service_jar = expansion_service_jar

	def __enter__(self):
	if not ExternalTransformTest.expansion_service_jar:
	raise unittest.SkipTest('No expansion service jar provided.')

	# Start the java server and wait for it to be ready.
	self._server = subprocess.Popen(
	['java', '-jar', self._expansion_service_jar, str(self._port)])

	port = ExternalTransformTest.expansion_service_port or 8091
	address = 'localhost:%s' % str(port)

	with grpc.insecure_channel(address) as channel:
	grpc.channel_ready_future(channel).result()

	def __exit__(self, type, value, traceback):
	self._server.kill()

	def test_pipeline_generation(self):
	pipeline = beam.Pipeline()
	res = (pipeline
	\| beam.Create(['a', 'b'])
	\| beam.ExternalTransform(
	'simple',
	None,
	expansion_service.ExpansionServiceServicer()))
	assert_that(res, equal_to(['Simple(a)', 'Simple(b)']))

	proto, _ = pipeline.to_runner_api(
	return_context=True)
	pipeline_from_proto = Pipeline.from_runner_api(
	proto, pipeline.runner, pipeline._options)

	# Original pipeline has the un-expanded external transform
	self.assertEqual([], pipeline.transforms_stack[0].parts[1].parts)

	# new pipeline has the expanded external transform
	self.assertNotEqual(
	[], pipeline_from_proto.transforms_stack[0].parts[1].parts)
	self.assertEqual(
	u'ExternalTransform(simple)/TestLabel',
	pipeline_from_proto.transforms_stack[0].parts[1].parts[0].full_label)

	def test_simple(self):
	with beam.Pipeline() as p:
	res = (
	p
	\| beam.Create(['a', 'b'])
	\| beam.ExternalTransform(
	'simple',
	None,
	expansion_service.ExpansionServiceServicer()))
	assert_that(res, equal_to(['Simple(a)', 'Simple(b)']))

	def test_multi(self):
	with beam.Pipeline() as p:
	main1 = p \| 'Main1' >> beam.Create(['a', 'bb'], reshuffle=False)
	main2 = p \| 'Main2' >> beam.Create(['x', 'yy', 'zzz'], reshuffle=False)
	side = p \| 'Side' >> beam.Create(['s'])
	res = dict(main1=main1, main2=main2, side=side) \| beam.ExternalTransform(
	'multi', None, expansion_service.ExpansionServiceServicer())
	assert_that(res['main'], equal_to(['as', 'bbs', 'xs', 'yys', 'zzzs']))
	assert_that(res['side'], equal_to(['ss']), label='CheckSide')

	def test_payload(self):
	with beam.Pipeline() as p:
	res = (
	p
	\| beam.Create(['a', 'bb'], reshuffle=False)
	\| beam.ExternalTransform(
	'payload', b's',
	expansion_service.ExpansionServiceServicer()))
	assert_that(res, equal_to(['as', 'bbs']))

	def test_nested(self):
	with beam.Pipeline() as p:
	assert_that(p \| FibTransform(6), equal_to([8]))

	def test_java_expansion_portable_runner(self):
	pipeline_options = PipelineOptions(
	['--runner=PortableRunner',
	'--experiments=beam_fn_api',
	'--environment_type=%s' % python_urns.EMBEDDED_PYTHON,
	'--job_endpoint=embed'])

	# We use the save_main_session option because one or more DoFn's in this
	# workflow rely on global context (e.g., a module imported at module level).
	pipeline_options.view_as(SetupOptions).save_main_session = True

	ExternalTransformTest.run_pipeline_with_portable_runner(pipeline_options)

	@unittest.skipIf(apiclient is None, 'GCP dependencies are not installed')
	def test_java_expansion_dataflow(self):
	# This test does not actually running the pipeline in Dataflow. It just
	# tests the translation to a Dataflow job request.

	with patch.object(
	apiclient.DataflowApplicationClient, 'create_job') as mock_create_job:
	port = ExternalTransformTest.expansion_service_port or 8091
	with self._RunWithExpansion(port, self.expansion_service_jar):
	pipeline_options = PipelineOptions(
	['--runner=DataflowRunner',
	'--project=dummyproject',
	'--experiments=beam_fn_api',
	'--temp_location=gs://dummybucket/'])

	# We use the save_main_session option because one or more DoFn's in this
	# workflow rely on global context (e.g., a module imported at module
	# level).
	pipeline_options.view_as(SetupOptions).save_main_session = True

	# Run a simple count-filtered-letters pipeline.
	self.run_pipeline(pipeline_options, port, False)

	mock_args = mock_create_job.call_args_list
	assert mock_args
	args, kwargs = mock_args[0]
	job = args[0]
	job_str = '%s' % job
	self.assertIn('pytest:beam:transforms:filter_less_than', job_str)

	@staticmethod
	def run_pipeline_with_portable_runner(pipeline_options):

	port = ExternalTransformTest.expansion_service_port or 8091

	with ExternalTransformTest._RunWithExpansion(
	port, ExternalTransformTest.expansion_service_jar):
	# Run a simple count-filtered-letters pipeline.
	ExternalTransformTest.run_pipeline(pipeline_options, port, True)

	@staticmethod
	def run_pipeline(
	pipeline_options, expansion_service_port, wait_until_finish=True):
	# The actual definitions of these transforms is in
	# org.apache.beam.runners.core.construction.TestExpansionService.
	TEST_COUNT_URN = "pytest:beam:transforms:count"
	TEST_FILTER_URN = "pytest:beam:transforms:filter_less_than"

	# Run a simple count-filtered-letters pipeline.
	p = beam.Pipeline(options=pipeline_options)
	address = 'localhost:%s' % str(expansion_service_port)
	res = (
	p
	\| beam.Create(list('aaabccxyyzzz'))
	\| beam.Map(unicode)
	# TODO(BEAM-6587): Use strings directly rather than ints.
	\| beam.Map(lambda x: int(ord(x)))
	\| beam.ExternalTransform(TEST_FILTER_URN, b'middle', address)
	\| beam.ExternalTransform(TEST_COUNT_URN, None, address)
	# # TODO(BEAM-6587): Remove when above is removed.
	\| beam.Map(lambda kv: (chr(kv[0]), kv[1]))
	\| beam.Map(lambda kv: '%s: %s' % kv))

	assert_that(res, equal_to(['a: 3', 'b: 1', 'c: 2']))

	result = p.run()
	if wait_until_finish:
	result.wait_until_finish()


	if __name__ == '__main__':
	parser = argparse.ArgumentParser()
	parser.add_argument('--expansion_service_jar')
	parser.add_argument('--expansion_service_port')
	known_args, pipeline_args = parser.parse_known_args(sys.argv)

	if known_args.expansion_service_jar:
	ExternalTransformTest.expansion_service_jar = (
	known_args.expansion_service_jar)
	ExternalTransformTest.expansion_service_port = int(
	known_args.expansion_service_port)
	pipeline_options = PipelineOptions(pipeline_args)
	ExternalTransformTest.run_pipeline_with_portable_runner(pipeline_options)
	else:
	sys.argv = pipeline_args
	unittest.main()