| # | |
| # 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.external classes.""" | |
| # pytype: skip-file | |
| import dataclasses | |
| import logging | |
| import os | |
| import tempfile | |
| import typing | |
| import unittest | |
| import mock | |
| import apache_beam as beam | |
| from apache_beam import ManagedReplacement | |
| from apache_beam import Pipeline | |
| from apache_beam.coders import RowCoder | |
| from apache_beam.options.pipeline_options import PipelineOptions | |
| from apache_beam.portability.api import beam_expansion_api_pb2 | |
| from apache_beam.portability.api import external_transforms_pb2 | |
| from apache_beam.portability.api import schema_pb2 | |
| from apache_beam.portability.common_urns import ManagedTransforms | |
| from apache_beam.runners import pipeline_context | |
| 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 | |
| from apache_beam.transforms import external | |
| from apache_beam.transforms.external import MANAGED_SCHEMA_TRANSFORM_IDENTIFIER | |
| from apache_beam.transforms.external import AnnotationBasedPayloadBuilder | |
| from apache_beam.transforms.external import ImplicitSchemaPayloadBuilder | |
| from apache_beam.transforms.external import JavaClassLookupPayloadBuilder | |
| from apache_beam.transforms.external import JavaExternalTransform | |
| from apache_beam.transforms.external import JavaJarExpansionService | |
| from apache_beam.transforms.external import NamedTupleBasedPayloadBuilder | |
| from apache_beam.transforms.external import SchemaTransformPayloadBuilder | |
| from apache_beam.typehints import typehints | |
| from apache_beam.typehints.native_type_compatibility import convert_to_beam_type | |
| from apache_beam.utils import proto_utils | |
| from apache_beam.utils.subprocess_server import JavaJarServer | |
| from apache_beam.utils.subprocess_server import SubprocessServer | |
| # 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 # type: ignore | |
| # pylint: enable=wrong-import-order, wrong-import-position | |
| def get_payload(cls): | |
| payload = external_transforms_pb2.ExternalConfigurationPayload() | |
| payload.ParseFromString(cls._payload) | |
| return payload | |
| class PayloadBase(object): | |
| values = { | |
| 'integer_example': 1, | |
| 'boolean': True, | |
| 'string_example': 'thing', | |
| 'list_of_strings': ['foo', 'bar'], | |
| 'mapping': { | |
| 'key': 1.1 | |
| }, | |
| 'optional_integer': None, | |
| } | |
| bytes_values = { | |
| 'integer_example': 1, | |
| 'boolean': True, | |
| 'string_example': 'thing', | |
| 'list_of_strings': ['foo', 'bar'], | |
| 'mapping': { | |
| 'key': 1.1 | |
| }, | |
| 'optional_integer': None, | |
| } | |
| def get_payload_from_typing_hints(self, values): | |
| """Return ExternalConfigurationPayload based on python typing hints""" | |
| raise NotImplementedError | |
| def get_payload_from_beam_typehints(self, values): | |
| """Return ExternalConfigurationPayload based on beam typehints""" | |
| raise NotImplementedError | |
| def test_typing_payload_builder(self): | |
| result = self.get_payload_from_typing_hints(self.values) | |
| decoded = RowCoder(result.schema).decode(result.payload) | |
| for key, value in self.values.items(): | |
| self.assertEqual(getattr(decoded, key), value) | |
| def test_typehints_payload_builder(self): | |
| result = self.get_payload_from_typing_hints(self.values) | |
| decoded = RowCoder(result.schema).decode(result.payload) | |
| for key, value in self.values.items(): | |
| self.assertEqual(getattr(decoded, key), value) | |
| def test_optional_error(self): | |
| """ | |
| value can only be None if typehint is Optional | |
| """ | |
| with self.assertRaises(ValueError): | |
| self.get_payload_from_typing_hints({k: None for k in self.values}) | |
| class ExternalTuplePayloadTest(PayloadBase, unittest.TestCase): | |
| def get_payload_from_typing_hints(self, values): | |
| TestSchema = typing.NamedTuple( | |
| 'TestSchema', | |
| [ | |
| ('integer_example', int), | |
| ('boolean', bool), | |
| ('string_example', str), | |
| ('list_of_strings', list[str]), | |
| ('mapping', typing.Mapping[str, float]), | |
| ('optional_integer', typing.Optional[int]), | |
| ]) | |
| builder = NamedTupleBasedPayloadBuilder(TestSchema(**values)) | |
| return builder.build() | |
| def get_payload_from_beam_typehints(self, values): | |
| raise unittest.SkipTest( | |
| "Beam typehints cannot be used with " | |
| "typing.NamedTuple") | |
| class ExternalImplicitPayloadTest(unittest.TestCase): | |
| """ | |
| ImplicitSchemaPayloadBuilder works very differently than the other payload | |
| builders | |
| """ | |
| def test_implicit_payload_builder(self): | |
| builder = ImplicitSchemaPayloadBuilder(PayloadBase.values) | |
| result = builder.build() | |
| decoded = RowCoder(result.schema).decode(result.payload) | |
| for key, value in PayloadBase.values.items(): | |
| # Note the default value in the getattr call. | |
| # ImplicitSchemaPayloadBuilder omits fields with valu=None since their | |
| # type cannot be inferred. | |
| self.assertEqual(getattr(decoded, key, None), value) | |
| def test_implicit_payload_builder_with_bytes(self): | |
| values = PayloadBase.bytes_values | |
| builder = ImplicitSchemaPayloadBuilder(values) | |
| result = builder.build() | |
| decoded = RowCoder(result.schema).decode(result.payload) | |
| for key, value in PayloadBase.values.items(): | |
| # Note the default value in the getattr call. | |
| # ImplicitSchemaPayloadBuilder omits fields with valu=None since their | |
| # type cannot be inferred. | |
| self.assertEqual(getattr(decoded, key, None), value) | |
| # Verify we have not modified a cached type (BEAM-10766) | |
| # TODO(BEAM-7372): Remove when bytes coercion code is removed. | |
| self.assertEqual(typehints.List[bytes], convert_to_beam_type(list[bytes])) | |
| class ExternalTransformTest(unittest.TestCase): | |
| def test_pipeline_generation(self): | |
| pipeline = beam.Pipeline() | |
| _ = ( | |
| pipeline | |
| | beam.Create(['a', 'b']) | |
| | beam.ExternalTransform( | |
| 'beam:transforms:xlang:test:prefix', | |
| ImplicitSchemaPayloadBuilder({'data': '0'}), | |
| expansion_service.ExpansionServiceServicer())) | |
| 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( | |
| 'ExternalTransform(beam:transforms:xlang:test:prefix)/TestLabel', | |
| pipeline_from_proto.transforms_stack[0].parts[1].parts[0].full_label) | |
| @unittest.skipIf(apiclient is None, 'GCP dependencies are not installed') | |
| def test_pipeline_generation_with_runner_overrides(self): | |
| pipeline_properties = [ | |
| '--job_name=test-job', | |
| '--project=test-project', | |
| '--temp_location=gs://beam/tmp', | |
| '--no_auth', | |
| '--dry_run=True', | |
| '--sdk_location=container', | |
| '--runner=DataflowRunner', | |
| '--streaming', | |
| '--region=us-central1' | |
| ] | |
| with beam.Pipeline(options=PipelineOptions(pipeline_properties)) as p: | |
| _ = ( | |
| p | |
| | beam.io.ReadFromPubSub( | |
| subscription= | |
| 'projects/dummy-project/subscriptions/dummy-subscription') | |
| | beam.ExternalTransform( | |
| 'beam:transforms:xlang:test:prefix', | |
| ImplicitSchemaPayloadBuilder({'data': '0'}), | |
| expansion_service.ExpansionServiceServicer())) | |
| pipeline_proto, _ = p.to_runner_api(return_context=True) | |
| pubsub_read_transform = None | |
| external_transform = None | |
| proto_transforms = pipeline_proto.components.transforms | |
| for id in proto_transforms: | |
| if 'beam:transforms:xlang:test:prefix' in proto_transforms[ | |
| id].unique_name: | |
| external_transform = proto_transforms[id] | |
| if 'ReadFromPubSub' in proto_transforms[id].unique_name: | |
| pubsub_read_transform = proto_transforms[id] | |
| if not (pubsub_read_transform and external_transform): | |
| raise ValueError( | |
| 'Could not find an external transform and the PubSub read transform ' | |
| 'in the pipeline') | |
| self.assertEqual(1, len(list(pubsub_read_transform.outputs.values()))) | |
| self.assertEqual(1, len(list(external_transform.inputs.values()))) | |
| # Verify that the PubSub read transform output is connected to the | |
| # external transform input. Instead of comparing exact PCollection | |
| # reference IDs (which can be non-deterministic), we verify that both | |
| # transforms reference valid PCollections in the pipeline components | |
| pubsub_output_id = list(pubsub_read_transform.outputs.values())[0] | |
| external_input_id = list(external_transform.inputs.values())[0] | |
| # Both should reference valid PCollections in the pipeline components | |
| self.assertIn(pubsub_output_id, pipeline_proto.components.pcollections) | |
| self.assertIn(external_input_id, pipeline_proto.components.pcollections) | |
| # Verify that the pipeline structure is correct by checking that | |
| # we have exactly 2 PCollections total (the intermediate one between | |
| # the transforms, and the final output from external transform) | |
| total_pcollections = len(pipeline_proto.components.pcollections) | |
| self.assertGreaterEqual( | |
| total_pcollections, 1, "Pipeline should have at least 1 PCollection") | |
| 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_output_coder(self): | |
| external_transform = beam.ExternalTransform( | |
| 'map_to_union_types', | |
| None, | |
| expansion_service.ExpansionServiceServicer()).with_output_types(int) | |
| with beam.Pipeline() as p: | |
| res = (p | beam.Create([2, 2], reshuffle=False) | external_transform) | |
| assert_that(res, equal_to([2, 2])) | |
| context = pipeline_context.PipelineContext( | |
| external_transform._expanded_components) | |
| self.assertEqual(len(external_transform._expanded_transform.outputs), 1) | |
| for _, pcol_id in external_transform._expanded_transform.outputs.items(): | |
| pcol = context.pcollections.get_by_id(pcol_id) | |
| self.assertEqual(pcol.element_type, int) | |
| def test_no_output_coder(self): | |
| external_transform = beam.ExternalTransform( | |
| 'map_to_union_types', | |
| None, | |
| expansion_service.ExpansionServiceServicer()) | |
| with beam.Pipeline() as p: | |
| res = (p | beam.Create([2, 2], reshuffle=False) | external_transform) | |
| assert_that(res, equal_to([2, 2])) | |
| context = pipeline_context.PipelineContext( | |
| external_transform._expanded_components) | |
| self.assertEqual(len(external_transform._expanded_transform.outputs), 1) | |
| for _, pcol_id in external_transform._expanded_transform.outputs.items(): | |
| pcol = context.pcollections.get_by_id(pcol_id) | |
| self.assertEqual(pcol.element_type, typehints.Any) | |
| def test_nested(self): | |
| with beam.Pipeline() as p: | |
| assert_that(p | FibTransform(6), equal_to([8])) | |
| def test_external_empty_spec_translation(self): | |
| pipeline = beam.Pipeline() | |
| external_transform = beam.ExternalTransform( | |
| 'beam:transforms:xlang:test:prefix', | |
| ImplicitSchemaPayloadBuilder({'data': '0'}), | |
| expansion_service.ExpansionServiceServicer()) | |
| _ = (pipeline | beam.Create(['a', 'b']) | external_transform) | |
| pipeline.run().wait_until_finish() | |
| external_transform_label = ( | |
| 'ExternalTransform(beam:transforms:xlang:test:prefix)/TestLabel') | |
| for transform in external_transform._expanded_components.transforms.values( | |
| ): | |
| # We clear the spec of one of the external transforms. | |
| if transform.unique_name == external_transform_label: | |
| transform.spec.Clear() | |
| context = pipeline_context.PipelineContext() | |
| proto_pipeline = pipeline.to_runner_api(context=context) | |
| proto_transform = None | |
| for transform in proto_pipeline.components.transforms.values(): | |
| if (transform.unique_name == | |
| 'ExternalTransform(beam:transforms:xlang:test:prefix)/TestLabel'): | |
| proto_transform = transform | |
| self.assertIsNotNone(proto_transform) | |
| self.assertTrue(str(proto_transform).strip().find('spec {') == -1) | |
| def test_unique_name(self): | |
| p = beam.Pipeline() | |
| _ = p | FibTransform(6) | |
| proto = p.to_runner_api() | |
| xforms = [x.unique_name for x in proto.components.transforms.values()] | |
| self.assertEqual( | |
| len(set(xforms)), len(xforms), msg='Transform names are not unique.') | |
| pcolls = [x.unique_name for x in proto.components.pcollections.values()] | |
| self.assertEqual( | |
| len(set(pcolls)), len(pcolls), msg='PCollection names are not unique.') | |
| def test_external_transform_finder_non_leaf(self): | |
| pipeline = beam.Pipeline() | |
| _ = ( | |
| pipeline | |
| | beam.Create(['a', 'b']) | |
| | beam.ExternalTransform( | |
| 'beam:transforms:xlang:test:prefix', | |
| ImplicitSchemaPayloadBuilder({'data': '0'}), | |
| expansion_service.ExpansionServiceServicer()) | |
| | beam.Map(lambda x: x)) | |
| pipeline.run().wait_until_finish() | |
| self.assertTrue(pipeline.contains_external_transforms) | |
| def test_external_transform_finder_leaf(self): | |
| pipeline = beam.Pipeline() | |
| _ = ( | |
| pipeline | |
| | beam.Create(['a', 'b']) | |
| | beam.ExternalTransform( | |
| 'beam:transforms:xlang:test:nooutput', | |
| ImplicitSchemaPayloadBuilder({'data': '0'}), | |
| expansion_service.ExpansionServiceServicer())) | |
| pipeline.run().wait_until_finish() | |
| self.assertTrue(pipeline.contains_external_transforms) | |
| def test_sanitize_java_traceback(self): | |
| error_string = ''' | |
| java.lang.RuntimeException: ACTUAL \n MULTILINE \n ERROR | |
| \tat org.apache.beam.sdk.expansion.service.ExpansionService$TransformProviderForBuilder.getTransform(ExpansionService.java:308) | |
| \tat org.apache.beam.sdk.expansion.service.TransformProvider.apply(TransformProvider.java:121) | |
| \tat org.apache.beam.sdk.expansion.service.ExpansionService.expand(ExpansionService.java:627) | |
| \tat org.apache.beam.sdk.expansion.service.ExpansionService.expand(ExpansionService.java:729) | |
| \tat org.apache.beam.model.expansion.v1.ExpansionServiceGrpc$MethodHandlers.invoke(ExpansionServiceGrpc.java:306) | |
| \tat org.apache.beam.vendor.grpc.v1p69p0.io.grpc.stub.ServerCalls$UnaryServerCallHandler$UnaryServerCallListener.onHalfClose(ServerCalls.java:182) | |
| \tat org.apache.beam.vendor.grpc.v1p69p0.io.grpc.internal.ServerCallImpl$ServerStreamListenerImpl.halfClosed(ServerCallImpl.java:351) | |
| \tat org.apache.beam.vendor.grpc.v1p69p0.io.grpc.internal.ServerImpl$JumpToApplicationThreadServerStreamListener$1HalfClosed.runInContext(ServerImpl.java:861) | |
| \tat org.apache.beam.vendor.grpc.v1p69p0.io.grpc.internal.ContextRunnable.run(ContextRunnable.java:37) | |
| \tat org.apache.beam.vendor.grpc.v1p69p0.io.grpc.internal.SerializingExecutor.run(SerializingExecutor.java:133) | |
| \tat java.util.concurrent.ThreadPoolExecutor.runWorker(ThreadPoolExecutor.java:1149) | |
| \tat java.util.concurrent.ThreadPoolExecutor$Worker.run(ThreadPoolExecutor.java:624) | |
| \tat java.lang.Thread.run(Thread.java:748) | |
| Caused by: java.lang.IllegalArgumentException: Received unknown SQL Dialect 'X'. Known dialects: [calcite] | |
| \tat org.apache.beam.sdk.extensions.sql.expansion.ExternalSqlTransformRegistrar$Builder.buildExternal(ExternalSqlTransformRegistrar.java:73) | |
| \tat org.apache.beam.sdk.extensions.sql.expansion.ExternalSqlTransformRegistrar$Builder.buildExternal(ExternalSqlTransformRegistrar.java:63) | |
| \tat org.apache.beam.sdk.expansion.service.ExpansionService$TransformProviderForBuilder.getTransform(ExpansionService.java:303) | |
| \t... 12 more | |
| '''.strip() | |
| core_msg = 'java.lang.RuntimeException: ACTUAL \n MULTILINE \n ERROR' | |
| self.assertEqual( | |
| f"{error_string}\n\n{core_msg}", | |
| external._sanitize_java_traceback(error_string)) | |
| class ExternalAnnotationPayloadTest(PayloadBase, unittest.TestCase): | |
| def get_payload_from_typing_hints(self, values): | |
| class AnnotatedTransform(beam.ExternalTransform): | |
| URN = 'beam:external:fakeurn:v1' | |
| def __init__( | |
| self, | |
| integer_example: int, | |
| boolean: bool, | |
| string_example: str, | |
| list_of_strings: list[str], | |
| mapping: typing.Mapping[str, float], | |
| optional_integer: typing.Optional[int] = None, | |
| expansion_service=None): | |
| super().__init__( | |
| self.URN, | |
| AnnotationBasedPayloadBuilder( | |
| self, | |
| integer_example=integer_example, | |
| boolean=boolean, | |
| string_example=string_example, | |
| list_of_strings=list_of_strings, | |
| mapping=mapping, | |
| optional_integer=optional_integer, | |
| ), | |
| expansion_service) | |
| return get_payload(AnnotatedTransform(**values)) | |
| def get_payload_from_beam_typehints(self, values): | |
| class AnnotatedTransform(beam.ExternalTransform): | |
| URN = 'beam:external:fakeurn:v1' | |
| def __init__( | |
| self, | |
| integer_example: int, | |
| boolean: bool, | |
| string_example: str, | |
| list_of_strings: typehints.List[str], | |
| mapping: typehints.Dict[str, float], | |
| optional_integer: typehints.Optional[int] = None, | |
| expansion_service=None): | |
| super().__init__( | |
| self.URN, | |
| AnnotationBasedPayloadBuilder( | |
| self, | |
| integer_example=integer_example, | |
| boolean=boolean, | |
| string_example=string_example, | |
| list_of_strings=list_of_strings, | |
| mapping=mapping, | |
| optional_integer=optional_integer, | |
| ), | |
| expansion_service) | |
| return get_payload(AnnotatedTransform(**values)) | |
| class ExternalDataclassesPayloadTest(PayloadBase, unittest.TestCase): | |
| def get_payload_from_typing_hints(self, values): | |
| @dataclasses.dataclass | |
| class DataclassTransform(beam.ExternalTransform): | |
| URN = 'beam:external:fakeurn:v1' | |
| integer_example: int | |
| boolean: bool | |
| string_example: str | |
| list_of_strings: list[str] | |
| mapping: typing.Mapping[str, float] = dataclasses.field(default=dict) | |
| optional_integer: typing.Optional[int] = None | |
| expansion_service: dataclasses.InitVar[typing.Optional[str]] = None | |
| return get_payload(DataclassTransform(**values)) | |
| def get_payload_from_beam_typehints(self, values): | |
| @dataclasses.dataclass | |
| class DataclassTransform(beam.ExternalTransform): | |
| URN = 'beam:external:fakeurn:v1' | |
| integer_example: int | |
| boolean: bool | |
| string_example: str | |
| list_of_strings: typehints.List[str] | |
| mapping: typehints.Dict[str, float] = {} | |
| optional_integer: typehints.Optional[int] = None | |
| expansion_service: dataclasses.InitVar[typehints.Optional[str]] = None | |
| return get_payload(DataclassTransform(**values)) | |
| class SchemaTransformPayloadBuilderTest(unittest.TestCase): | |
| def test_build_payload(self): | |
| ComplexType = typing.NamedTuple( | |
| "ComplexType", [ | |
| ("str_sub_field", str), | |
| ("int_sub_field", int), | |
| ]) | |
| payload_builder = SchemaTransformPayloadBuilder( | |
| identifier='dummy_id', | |
| str_field='aaa', | |
| int_field=123, | |
| object_field=ComplexType(str_sub_field="bbb", int_sub_field=456)) | |
| payload_bytes = payload_builder.payload() | |
| payload_from_bytes = proto_utils.parse_Bytes( | |
| payload_bytes, external_transforms_pb2.SchemaTransformPayload) | |
| self.assertEqual('dummy_id', payload_from_bytes.identifier) | |
| expected_coder = RowCoder(payload_from_bytes.configuration_schema) | |
| schema_transform_config = expected_coder.decode( | |
| payload_from_bytes.configuration_row) | |
| self.assertEqual('aaa', schema_transform_config.str_field) | |
| self.assertEqual(123, schema_transform_config.int_field) | |
| self.assertEqual('bbb', schema_transform_config.object_field.str_sub_field) | |
| self.assertEqual(456, schema_transform_config.object_field.int_sub_field) | |
| class SchemaAwareExternalTransformTest(unittest.TestCase): | |
| class MockDiscoveryService: | |
| # define context manager enter and exit functions | |
| def __enter__(self): | |
| return self | |
| def __exit__(self, unusued1, unused2, unused3): | |
| pass | |
| def DiscoverSchemaTransform(self, unused_request=None): | |
| test_config = beam_expansion_api_pb2.SchemaTransformConfig( | |
| config_schema=schema_pb2.Schema( | |
| fields=[ | |
| schema_pb2.Field( | |
| name="str_field", | |
| type=schema_pb2.FieldType(atomic_type="STRING")), | |
| schema_pb2.Field( | |
| name="int_field", | |
| type=schema_pb2.FieldType(atomic_type="INT64")) | |
| ], | |
| id="test-id"), | |
| input_pcollection_names=["input"], | |
| output_pcollection_names=["output"]) | |
| test_managed_config = beam_expansion_api_pb2.SchemaTransformConfig( | |
| config_schema=schema_pb2.Schema( | |
| fields=[ | |
| schema_pb2.Field( | |
| name="transform_identifier", | |
| type=schema_pb2.FieldType(atomic_type="STRING")), | |
| schema_pb2.Field( | |
| name="config_url", | |
| type=schema_pb2.FieldType(atomic_type="STRING")), | |
| schema_pb2.Field( | |
| name="config", | |
| type=schema_pb2.FieldType(atomic_type="STRING")) | |
| ], | |
| id="test-id1"), | |
| input_pcollection_names=["input"], | |
| output_pcollection_names=["output"]) | |
| return beam_expansion_api_pb2.DiscoverSchemaTransformResponse( | |
| schema_transform_configs={ | |
| "test_schematransform": test_config, | |
| MANAGED_SCHEMA_TRANSFORM_IDENTIFIER: test_managed_config | |
| }) | |
| @mock.patch("apache_beam.transforms.external.ExternalTransform.service") | |
| def test_discover_one_config(self, mock_service): | |
| _mock = self.MockDiscoveryService() | |
| mock_service.return_value = _mock | |
| config = beam.SchemaAwareExternalTransform.discover_config( | |
| "test_service", name="test_schematransform") | |
| self.assertEqual(config.outputs[0], "output") | |
| self.assertEqual(config.inputs[0], "input") | |
| self.assertEqual(config.identifier, "test_schematransform") | |
| @mock.patch("apache_beam.transforms.external.ExternalTransform.service") | |
| def test_discover_one_config_fails_with_no_configs_found(self, mock_service): | |
| mock_service.return_value = self.MockDiscoveryService() | |
| with self.assertRaises(ValueError): | |
| beam.SchemaAwareExternalTransform.discover_config( | |
| "test_service", name="non_existent") | |
| @mock.patch("apache_beam.transforms.external.ExternalTransform.service") | |
| def test_rearrange_kwargs_based_on_discovery(self, mock_service): | |
| mock_service.return_value = self.MockDiscoveryService() | |
| identifier = "test_schematransform" | |
| expansion_service = "test_service" | |
| kwargs = {"int_field": 0, "str_field": "str"} | |
| transform = beam.SchemaAwareExternalTransform( | |
| identifier=identifier, | |
| expansion_service=expansion_service, | |
| rearrange_based_on_discovery=True, | |
| **kwargs) | |
| payload = transform._payload_builder.build() | |
| ordered_fields = [f.name for f in payload.configuration_schema.fields] | |
| schematransform_config = beam.SchemaAwareExternalTransform.discover_config( | |
| expansion_service, identifier) | |
| external_config_fields = schematransform_config.configuration_schema._fields | |
| self.assertNotEqual(tuple(kwargs.keys()), external_config_fields) | |
| self.assertEqual(tuple(ordered_fields), external_config_fields) | |
| @mock.patch("apache_beam.transforms.external.ExternalTransform.service") | |
| def test_managed_replacement_unknown_id(self, mock_service): | |
| mock_service.return_value = self.MockDiscoveryService() | |
| identifier = "test_schematransform" | |
| kwargs = {"int_field": 0, "str_field": "str"} | |
| managed_replacement = ManagedReplacement( | |
| underlying_transform_identifier="unknown_id", | |
| update_compatibility_version="2.50.0") | |
| with self.assertRaises(ValueError): | |
| beam.SchemaAwareExternalTransform( | |
| identifier=identifier, | |
| expansion_service=expansion_service, | |
| rearrange_based_on_discovery=True, | |
| managed_replacement=managed_replacement, | |
| **kwargs) | |
| @mock.patch("apache_beam.transforms.external.ExternalTransform.service") | |
| @mock.patch("apache_beam.transforms.external.BeamJarExpansionService") | |
| def test_managed_replacement_known_id( | |
| self, mock_service, mock_beam_jar_service): | |
| mock_service.return_value = self.MockDiscoveryService() | |
| mock_beam_jar_service.return_value = self.MockDiscoveryService() | |
| identifier = "test_schematransform" | |
| kwargs = {"int_field": 0, "str_field": "str"} | |
| managed_replacement = ManagedReplacement( | |
| underlying_transform_identifier=ManagedTransforms.Urns.ICEBERG_READ.urn, | |
| update_compatibility_version="2.50.0") | |
| external_transform = beam.SchemaAwareExternalTransform( | |
| identifier=identifier, | |
| expansion_service=expansion_service, | |
| rearrange_based_on_discovery=True, | |
| managed_replacement=managed_replacement, | |
| **kwargs) | |
| self.assertIsNotNone(external_transform._managed_payload_builder) | |
| class JavaClassLookupPayloadBuilderTest(unittest.TestCase): | |
| def _verify_row(self, schema, row_payload, expected_values): | |
| row = RowCoder(schema).decode(row_payload) | |
| for attr_name, expected_value in expected_values.items(): | |
| self.assertTrue(hasattr(row, attr_name)) | |
| value = getattr(row, attr_name) | |
| self.assertEqual(expected_value, value) | |
| def test_build_payload_with_constructor(self): | |
| payload_builder = JavaClassLookupPayloadBuilder('dummy_class_name') | |
| payload_builder.with_constructor('abc', 123, str_field='def', int_field=456) | |
| payload_bytes = payload_builder.payload() | |
| payload_from_bytes = proto_utils.parse_Bytes( | |
| payload_bytes, external_transforms_pb2.JavaClassLookupPayload) | |
| self.assertTrue( | |
| isinstance( | |
| payload_from_bytes, external_transforms_pb2.JavaClassLookupPayload)) | |
| self.assertFalse(payload_from_bytes.constructor_method) | |
| self._verify_row( | |
| payload_from_bytes.constructor_schema, | |
| payload_from_bytes.constructor_payload, { | |
| 'ignore0': 'abc', | |
| 'ignore1': 123, | |
| 'str_field': 'def', | |
| 'int_field': 456 | |
| }) | |
| def test_build_payload_with_constructor_method(self): | |
| payload_builder = JavaClassLookupPayloadBuilder('dummy_class_name') | |
| payload_builder.with_constructor_method( | |
| 'dummy_constructor_method', 'abc', 123, str_field='def', int_field=456) | |
| payload_bytes = payload_builder.payload() | |
| payload_from_bytes = proto_utils.parse_Bytes( | |
| payload_bytes, external_transforms_pb2.JavaClassLookupPayload) | |
| self.assertTrue( | |
| isinstance( | |
| payload_from_bytes, external_transforms_pb2.JavaClassLookupPayload)) | |
| self.assertEqual( | |
| 'dummy_constructor_method', payload_from_bytes.constructor_method) | |
| self._verify_row( | |
| payload_from_bytes.constructor_schema, | |
| payload_from_bytes.constructor_payload, { | |
| 'ignore0': 'abc', | |
| 'ignore1': 123, | |
| 'str_field': 'def', | |
| 'int_field': 456 | |
| }) | |
| def test_build_payload_with_builder_methods(self): | |
| payload_builder = JavaClassLookupPayloadBuilder('dummy_class_name') | |
| payload_builder.with_constructor('abc', 123, str_field='def', int_field=456) | |
| payload_builder.add_builder_method( | |
| 'builder_method1', 'abc1', 1234, str_field1='abc2', int_field1=2345) | |
| payload_builder.add_builder_method( | |
| 'builder_method2', 'abc3', 3456, str_field2='abc4', int_field2=4567) | |
| payload_bytes = payload_builder.payload() | |
| payload_from_bytes = proto_utils.parse_Bytes( | |
| payload_bytes, external_transforms_pb2.JavaClassLookupPayload) | |
| self.assertTrue( | |
| isinstance( | |
| payload_from_bytes, external_transforms_pb2.JavaClassLookupPayload)) | |
| self._verify_row( | |
| payload_from_bytes.constructor_schema, | |
| payload_from_bytes.constructor_payload, { | |
| 'ignore0': 'abc', | |
| 'ignore1': 123, | |
| 'str_field': 'def', | |
| 'int_field': 456 | |
| }) | |
| self.assertEqual(2, len(payload_from_bytes.builder_methods)) | |
| builder_method = payload_from_bytes.builder_methods[0] | |
| self.assertTrue( | |
| isinstance(builder_method, external_transforms_pb2.BuilderMethod)) | |
| self.assertEqual('builder_method1', builder_method.name) | |
| self._verify_row( | |
| builder_method.schema, | |
| builder_method.payload, | |
| { | |
| 'ignore0': 'abc1', | |
| 'ignore1': 1234, | |
| 'str_field1': 'abc2', | |
| 'int_field1': 2345 | |
| }) | |
| builder_method = payload_from_bytes.builder_methods[1] | |
| self.assertTrue( | |
| isinstance(builder_method, external_transforms_pb2.BuilderMethod)) | |
| self.assertEqual('builder_method2', builder_method.name) | |
| self._verify_row( | |
| builder_method.schema, | |
| builder_method.payload, | |
| { | |
| 'ignore0': 'abc3', | |
| 'ignore1': 3456, | |
| 'str_field2': 'abc4', | |
| 'int_field2': 4567 | |
| }) | |
| def test_build_payload_with_constructor_twice_fails(self): | |
| payload_builder = JavaClassLookupPayloadBuilder('dummy_class_name') | |
| payload_builder.with_constructor('abc') | |
| with self.assertRaises(ValueError): | |
| payload_builder.with_constructor('def') | |
| def test_implicit_builder_with_constructor(self): | |
| constructor_transform = ( | |
| JavaExternalTransform('org.pkg.MyTransform')('abc').withIntProperty(5)) | |
| payload_bytes = constructor_transform._payload_builder.payload() | |
| payload_from_bytes = proto_utils.parse_Bytes( | |
| payload_bytes, external_transforms_pb2.JavaClassLookupPayload) | |
| self.assertEqual('org.pkg.MyTransform', payload_from_bytes.class_name) | |
| self._verify_row( | |
| payload_from_bytes.constructor_schema, | |
| payload_from_bytes.constructor_payload, {'ignore0': 'abc'}) | |
| builder_method = payload_from_bytes.builder_methods[0] | |
| self.assertEqual('withIntProperty', builder_method.name) | |
| self._verify_row( | |
| builder_method.schema, builder_method.payload, {'ignore0': 5}) | |
| def test_implicit_builder_with_constructor_method(self): | |
| constructor_transform = JavaExternalTransform('org.pkg.MyTransform').of( | |
| str_field='abc').withProperty(int_field=1234).build() | |
| payload_bytes = constructor_transform._payload_builder.payload() | |
| payload_from_bytes = proto_utils.parse_Bytes( | |
| payload_bytes, external_transforms_pb2.JavaClassLookupPayload) | |
| self.assertEqual('of', payload_from_bytes.constructor_method) | |
| self._verify_row( | |
| payload_from_bytes.constructor_schema, | |
| payload_from_bytes.constructor_payload, {'str_field': 'abc'}) | |
| with_property_method = payload_from_bytes.builder_methods[0] | |
| self.assertEqual('withProperty', with_property_method.name) | |
| self._verify_row( | |
| with_property_method.schema, | |
| with_property_method.payload, {'int_field': 1234}) | |
| build_method = payload_from_bytes.builder_methods[1] | |
| self.assertEqual('build', build_method.name) | |
| self._verify_row(build_method.schema, build_method.payload, {}) | |
| class JavaJarExpansionServiceTest(unittest.TestCase): | |
| def setUp(self): | |
| # Temporarily override _live_owners with an empty set for this test, | |
| # preventing contamination of the process-wide global cache and avoiding | |
| # side effects on other tests. | |
| patcher = mock.patch.object( | |
| SubprocessServer._cache, '_live_owners', new=set()) | |
| patcher.start() | |
| self.addCleanup(patcher.stop) | |
| def test_classpath(self): | |
| with tempfile.TemporaryDirectory() as temp_dir: | |
| try: | |
| # Avoid having to prefix everything in our test strings. | |
| oldwd = os.getcwd() | |
| os.chdir(temp_dir) | |
| # Touch some files for globing. | |
| with open('a1.jar', 'w') as _: | |
| pass | |
| service = JavaJarExpansionService( | |
| 'main.jar', classpath=['a*.jar', 'b.jar']) | |
| self.assertEqual( | |
| service._default_args(), | |
| ['{{PORT}}', '--filesToStage=main.jar,a1.jar,b.jar']) | |
| finally: | |
| os.chdir(oldwd) | |
| @mock.patch.object(JavaJarServer, 'local_jar') | |
| def test_classpath_with_url(self, local_jar): | |
| def _side_effect_fn(path): | |
| return path[path.rindex('/') + 1:] | |
| local_jar.side_effect = _side_effect_fn | |
| with tempfile.TemporaryDirectory() as temp_dir: | |
| try: | |
| # Avoid having to prefix everything in our test strings. | |
| oldwd = os.getcwd() | |
| os.chdir(temp_dir) | |
| service = JavaJarExpansionService( | |
| 'main.jar', classpath=['https://dummy_path/dummyjar.jar']) | |
| self.assertEqual( | |
| service._default_args(), | |
| ['{{PORT}}', '--filesToStage=main.jar,dummyjar.jar']) | |
| finally: | |
| os.chdir(oldwd) | |
| @mock.patch.object(JavaJarServer, 'local_jar') | |
| def test_classpath_with_gradle_artifact(self, local_jar): | |
| def _side_effect_fn(path, user_agent=None): | |
| return path[path.rindex('/') + 1:] | |
| local_jar.side_effect = _side_effect_fn | |
| with tempfile.TemporaryDirectory() as temp_dir: | |
| try: | |
| # Avoid having to prefix everything in our test strings. | |
| oldwd = os.getcwd() | |
| os.chdir(temp_dir) | |
| service = JavaJarExpansionService( | |
| 'main.jar', classpath=['dummy_group:dummy_artifact:dummy_version']) | |
| self.assertEqual( | |
| service._default_args(), | |
| [ | |
| '{{PORT}}', | |
| '--filesToStage=main.jar,dummy_artifact-dummy_version.jar' | |
| ]) | |
| finally: | |
| os.chdir(oldwd) | |
| def test_classpath_with_glob(self): | |
| with tempfile.TemporaryDirectory() as temp_dir: | |
| try: | |
| # Avoid having to prefix everything in our test strings. | |
| oldwd = os.getcwd() | |
| os.chdir(temp_dir) | |
| # Touch some files for globing. | |
| with open('a1.jar', 'w') as _: | |
| pass | |
| service = JavaJarExpansionService( | |
| 'main.jar', classpath=['a*.jar', 'b.jar']) | |
| self.assertEqual( | |
| service._default_args(), | |
| ['{{PORT}}', '--filesToStage=main.jar,a1.jar,b.jar']) | |
| finally: | |
| os.chdir(oldwd) | |
| if __name__ == '__main__': | |
| logging.getLogger().setLevel(logging.INFO) | |
| unittest.main() |