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
