def test_partition(self, test_pipeline=None): if not test_pipeline: test_pipeline = TestPipeline() test_pipeline.get_pipeline_options().view_as( DebugOptions).experiments.append( 'jar_packages=' + ValidateRunnerXlangTest.expansion_jar) test_pipeline.not_use_test_runner_api = True CrossLanguageTestPipelines().run_partition(test_pipeline)
def test_kafkaio_write(self): local_kafka_jar = os.environ.get('LOCAL_KAFKA_JAR') with self.local_kafka_service(local_kafka_jar) as kafka_port: options = self.get_options() p = TestPipeline(options=options) p.not_use_test_runner_api = True CrossLanguageKafkaIO('localhost:%s' % kafka_port, 'xlang_kafkaio_test').build_write_pipeline(p) job = p.run() job.wait_until_finish()
def test_kafkaio_write(self): local_kafka_jar = os.environ.get('LOCAL_KAFKA_JAR') with self.local_kafka_service(local_kafka_jar) as kafka_port: p = TestPipeline() p.not_use_test_runner_api = True xlang_kafkaio = CrossLanguageKafkaIO( '%s:%s' % (self.get_platform_localhost(), kafka_port), 'xlang_kafkaio_test') xlang_kafkaio.build_write_pipeline(p) job = p.run() job.wait_until_finish()
def test_combine_globally(self): test_pipeline = TestPipeline() test_pipeline.get_pipeline_options().view_as( DebugOptions).experiments.append( 'jar_packages=' + ValidateRunnerXlangTest.expansion_jar) test_pipeline.not_use_test_runner_api = True with test_pipeline as p: res = (p | beam.Create([1, 2, 3]).with_output_types(int) | beam.ExternalTransform( TEST_COMGL_URN, None, ValidateRunnerXlangTest.expansion_service)) assert_that(res, equal_to([6]))
def test_prefix(self): test_pipeline = TestPipeline() test_pipeline.get_pipeline_options().view_as( DebugOptions).experiments.append( 'jar_packages=' + ValidateRunnerXlangTest.expansion_jar) test_pipeline.not_use_test_runner_api = True with test_pipeline as p: res = (p | beam.Create(['a', 'b']).with_output_types(unicode) | beam.ExternalTransform( TEST_PREFIX_URN, ImplicitSchemaPayloadBuilder({'data': u'0'}), ValidateRunnerXlangTest.expansion_service)) assert_that(res, equal_to(['0a', '0b']))
def test_partition(self): test_pipeline = TestPipeline() test_pipeline.get_pipeline_options().view_as( DebugOptions).experiments.append( 'jar_packages=' + ValidateRunnerXlangTest.expansion_jar) test_pipeline.not_use_test_runner_api = True with test_pipeline as p: res = (p | beam.Create([1, 2, 3, 4, 5, 6]).with_output_types(int) | beam.ExternalTransform( TEST_PARTITION_URN, None, ValidateRunnerXlangTest.expansion_service)) assert_that(res['0'], equal_to([2, 4, 6]), label='check_even') assert_that(res['1'], equal_to([1, 3, 5]), label='check_odd')
def test_flatten(self): test_pipeline = TestPipeline() test_pipeline.get_pipeline_options().view_as( DebugOptions).experiments.append( 'jar_packages=' + ValidateRunnerXlangTest.expansion_jar) test_pipeline.not_use_test_runner_api = True with test_pipeline as p: col1 = p | 'col1' >> beam.Create([1, 2, 3]).with_output_types(int) col2 = p | 'col2' >> beam.Create([4, 5, 6]).with_output_types(int) res = ((col1, col2) | beam.ExternalTransform( TEST_FLATTEN_URN, None, ValidateRunnerXlangTest.expansion_service)) assert_that(res, equal_to([1, 2, 3, 4, 5, 6]))
def test_combine_per_key(self): test_pipeline = TestPipeline() test_pipeline.get_pipeline_options().view_as( DebugOptions).experiments.append( 'jar_packages=' + ValidateRunnerXlangTest.expansion_jar) test_pipeline.not_use_test_runner_api = True with test_pipeline as p: res = ( p | beam.Create([('a', 1), ('a', 2), ('b', 3)]).with_output_types( typing.Tuple[unicode, int]) | beam.ExternalTransform( TEST_COMPK_URN, None, ValidateRunnerXlangTest.expansion_service)) assert_that(res, equal_to([('a', 3), ('b', 3)]))
def test_group_by_key(self): test_pipeline = TestPipeline() test_pipeline.get_pipeline_options().view_as( DebugOptions).experiments.append( 'jar_packages=' + ValidateRunnerXlangTest.expansion_jar) test_pipeline.not_use_test_runner_api = True with test_pipeline as p: res = ( p | beam.Create([(0, "1"), (0, "2"), (1, "3")], reshuffle=False).with_output_types( typing.Tuple[int, unicode]) | beam.ExternalTransform( TEST_GBK_URN, None, ValidateRunnerXlangTest.expansion_service) | beam.Map(lambda x: "{}:{}".format(x[0], ','.join(sorted(x[1]))))) assert_that(res, equal_to(['0:1,2', '1:3']))
def test_multi_input_output_with_sideinput(self): test_pipeline = TestPipeline() test_pipeline.get_pipeline_options().view_as( DebugOptions).experiments.append( 'jar_packages=' + ValidateRunnerXlangTest.expansion_jar) test_pipeline.not_use_test_runner_api = True with test_pipeline as p: main1 = p | 'Main1' >> beam.Create( ['a', 'bb'], reshuffle=False).with_output_types(unicode) main2 = p | 'Main2' >> beam.Create( ['x', 'yy', 'zzz'], reshuffle=False).with_output_types(unicode) side = p | 'Side' >> beam.Create(['s']).with_output_types(unicode) res = dict( main1=main1, main2=main2, side=side) | beam.ExternalTransform( TEST_MULTI_URN, None, ValidateRunnerXlangTest.expansion_service) assert_that(res['main'], equal_to(['as', 'bbs', 'xs', 'yys', 'zzzs'])) assert_that(res['side'], equal_to(['ss']), label='CheckSide')
def test_write(self): expansion_jar = os.environ.get('EXPANSION_JAR') port = os.environ.get('EXPANSION_PORT') address = 'localhost:%s' % port try: test_pipeline = TestPipeline() test_pipeline.get_pipeline_options().view_as( DebugOptions).experiments.append('jar_packages='+expansion_jar) test_pipeline.not_use_test_runner_api = True with test_pipeline as p: _ = p \ | beam.Create([ AvroRecord({"name": "abc"}), AvroRecord({"name": "def"}), AvroRecord({"name": "ghi"})]) \ | beam.ExternalTransform( PARQUET_WRITE_URN, b'/tmp/test.parquet', address) except RuntimeError as e: if re.search(PARQUET_WRITE_URN, str(e)): print("looks like URN not implemented in expansion service, skipping.") else: raise e
def create_pipeline(self): test_pipeline = TestPipeline() test_pipeline.not_use_test_runner_api = True return test_pipeline