#
# 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', typing.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(typing.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: typing.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: typing.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):
    SubprocessServer._cache._live_owners = set()

  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()