################################################################################
#  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.
################################################################################
import json
import unittest
from typing import Dict

import pyflink.datastream.data_stream as data_stream
from pyflink.common import typeinfo
from pyflink.common.configuration import Configuration
from pyflink.common.serialization import SimpleStringSchema, ByteArraySchema, DeserializationSchema
from pyflink.common.typeinfo import Types
from pyflink.common.types import Row
from pyflink.common.watermark_strategy import WatermarkStrategy
from pyflink.datastream.connectors.base import DeliveryGuarantee
from pyflink.datastream.connectors.kafka import KafkaSource, KafkaTopicPartition, \
    KafkaOffsetsInitializer, KafkaOffsetResetStrategy, KafkaRecordSerializationSchema, KafkaSink, \
    FlinkKafkaProducer, FlinkKafkaConsumer
from pyflink.datastream.formats.avro import AvroRowDeserializationSchema, AvroRowSerializationSchema
from pyflink.datastream.formats.csv import CsvRowDeserializationSchema, CsvRowSerializationSchema
from pyflink.datastream.formats.json import JsonRowDeserializationSchema, JsonRowSerializationSchema
from pyflink.java_gateway import get_gateway
from pyflink.testing.test_case_utils import (
    PyFlinkStreamingTestCase,
    PyFlinkTestCase,
    invoke_java_object_method,
    to_java_data_structure,
)
from pyflink.util.java_utils import to_jarray, is_instance_of, get_field_value


@unittest.skip("Disable due to Kafka connector need to release a new version 2.0")
class KafkaSourceTests(PyFlinkStreamingTestCase):

    @unittest.skip("Disabled due to cyclic dependencies: Should be removed with Flink 2.0")
    def test_legacy_kafka_connector(self):
        source_topic = 'test_source_topic'
        sink_topic = 'test_sink_topic'
        props = {'bootstrap.servers': 'localhost:9092', 'group.id': 'test_group'}
        type_info = Types.ROW([Types.INT(), Types.STRING()])

        # Test for kafka consumer
        deserialization_schema = JsonRowDeserializationSchema.builder() \
            .type_info(type_info=type_info).build()

        flink_kafka_consumer = FlinkKafkaConsumer(source_topic, deserialization_schema, props)
        flink_kafka_consumer.set_start_from_earliest()
        flink_kafka_consumer.set_commit_offsets_on_checkpoints(True)

        j_properties = get_field_value(flink_kafka_consumer.get_java_function(), 'properties')
        self.assertEqual('localhost:9092', j_properties.getProperty('bootstrap.servers'))
        self.assertEqual('test_group', j_properties.getProperty('group.id'))
        self.assertTrue(get_field_value(flink_kafka_consumer.get_java_function(),
                                        'enableCommitOnCheckpoints'))
        j_start_up_mode = get_field_value(flink_kafka_consumer.get_java_function(), 'startupMode')

        j_deserializer = get_field_value(flink_kafka_consumer.get_java_function(), 'deserializer')
        j_deserialize_type_info = invoke_java_object_method(j_deserializer, "getProducedType")
        deserialize_type_info = typeinfo._from_java_type(j_deserialize_type_info)
        self.assertTrue(deserialize_type_info == type_info)
        self.assertTrue(j_start_up_mode.equals(get_gateway().jvm
                                               .org.apache.flink.streaming.connectors
                                               .kafka.config.StartupMode.EARLIEST))
        j_topic_desc = get_field_value(flink_kafka_consumer.get_java_function(),
                                       'topicsDescriptor')
        j_topics = invoke_java_object_method(j_topic_desc, 'getFixedTopics')
        self.assertEqual(['test_source_topic'], list(j_topics))

        # Test for kafka producer
        serialization_schema = JsonRowSerializationSchema.builder().with_type_info(type_info) \
            .build()
        flink_kafka_producer = FlinkKafkaProducer(sink_topic, serialization_schema, props)
        flink_kafka_producer.set_write_timestamp_to_kafka(False)

        j_producer_config = get_field_value(flink_kafka_producer.get_java_function(),
                                            'producerConfig')
        self.assertEqual('localhost:9092', j_producer_config.getProperty('bootstrap.servers'))
        self.assertEqual('test_group', j_producer_config.getProperty('group.id'))
        self.assertFalse(get_field_value(flink_kafka_producer.get_java_function(),
                                         'writeTimestampToKafka'))

    def test_compiling(self):
        source = KafkaSource.builder() \
            .set_bootstrap_servers('localhost:9092') \
            .set_topics('test_topic') \
            .set_value_only_deserializer(SimpleStringSchema()) \
            .build()

        ds = self.env.from_source(source=source,
                                  watermark_strategy=WatermarkStrategy.for_monotonous_timestamps(),
                                  source_name='kafka source')
        ds.print()
        plan = json.loads(self.env.get_execution_plan())
        self.assertEqual('Source: kafka source', plan['nodes'][0]['type'])

    def test_set_properties(self):
        source = KafkaSource.builder() \
            .set_bootstrap_servers('localhost:9092') \
            .set_group_id('test_group_id') \
            .set_client_id_prefix('test_client_id_prefix') \
            .set_property('test_property', 'test_value') \
            .set_topics('test_topic') \
            .set_value_only_deserializer(SimpleStringSchema()) \
            .build()
        conf = self._get_kafka_source_configuration(source)
        self.assertEqual(conf.get_string('bootstrap.servers', ''), 'localhost:9092')
        self.assertEqual(conf.get_string('group.id', ''), 'test_group_id')
        self.assertEqual(conf.get_string('client.id.prefix', ''), 'test_client_id_prefix')
        self.assertEqual(conf.get_string('test_property', ''), 'test_value')

    def test_set_topics(self):
        source = KafkaSource.builder() \
            .set_bootstrap_servers('localhost:9092') \
            .set_topics('test_topic1', 'test_topic2') \
            .set_value_only_deserializer(SimpleStringSchema()) \
            .build()
        kafka_subscriber = get_field_value(source.get_java_function(), 'subscriber')
        self.assertEqual(
            kafka_subscriber.getClass().getCanonicalName(),
            'org.apache.flink.connector.kafka.source.enumerator.subscriber.TopicListSubscriber'
        )
        topics = get_field_value(kafka_subscriber, 'topics')
        self.assertTrue(is_instance_of(topics, get_gateway().jvm.java.util.List))
        self.assertEqual(topics.size(), 2)
        self.assertEqual(topics[0], 'test_topic1')
        self.assertEqual(topics[1], 'test_topic2')

    def test_set_topic_pattern(self):
        source = KafkaSource.builder() \
            .set_bootstrap_servers('localhost:9092') \
            .set_topic_pattern('test_topic*') \
            .set_value_only_deserializer(SimpleStringSchema()) \
            .build()
        kafka_subscriber = get_field_value(source.get_java_function(), 'subscriber')
        self.assertEqual(
            kafka_subscriber.getClass().getCanonicalName(),
            'org.apache.flink.connector.kafka.source.enumerator.subscriber.TopicPatternSubscriber'
        )
        topic_pattern = get_field_value(kafka_subscriber, 'topicPattern')
        self.assertTrue(is_instance_of(topic_pattern, get_gateway().jvm.java.util.regex.Pattern))
        self.assertEqual(topic_pattern.toString(), 'test_topic*')

    def test_set_partitions(self):
        topic_partition_1 = KafkaTopicPartition('test_topic', 1)
        topic_partition_2 = KafkaTopicPartition('test_topic', 2)
        source = KafkaSource.builder() \
            .set_bootstrap_servers('localhost:9092') \
            .set_partitions({topic_partition_1, topic_partition_2}) \
            .set_value_only_deserializer(SimpleStringSchema()) \
            .build()
        kafka_subscriber = get_field_value(source.get_java_function(), 'subscriber')
        self.assertEqual(
            kafka_subscriber.getClass().getCanonicalName(),
            'org.apache.flink.connector.kafka.source.enumerator.subscriber.PartitionSetSubscriber'
        )
        partitions = get_field_value(kafka_subscriber, 'subscribedPartitions')
        self.assertTrue(is_instance_of(partitions, get_gateway().jvm.java.util.Set))
        self.assertTrue(topic_partition_1._to_j_topic_partition() in partitions)
        self.assertTrue(topic_partition_2._to_j_topic_partition() in partitions)

    def test_set_starting_offsets(self):
        def _build_source(initializer: KafkaOffsetsInitializer):
            return KafkaSource.builder() \
                .set_bootstrap_servers('localhost:9092') \
                .set_topics('test_topic') \
                .set_value_only_deserializer(SimpleStringSchema()) \
                .set_group_id('test_group') \
                .set_starting_offsets(initializer) \
                .build()

        self._check_reader_handled_offsets_initializer(
            _build_source(KafkaOffsetsInitializer.latest()), -1, KafkaOffsetResetStrategy.LATEST
        )
        self._check_reader_handled_offsets_initializer(
            _build_source(KafkaOffsetsInitializer.earliest()), -2,
            KafkaOffsetResetStrategy.EARLIEST
        )
        self._check_reader_handled_offsets_initializer(
            _build_source(KafkaOffsetsInitializer.committed_offsets()), -3,
            KafkaOffsetResetStrategy.NONE
        )
        self._check_reader_handled_offsets_initializer(
            _build_source(KafkaOffsetsInitializer.committed_offsets(
                KafkaOffsetResetStrategy.LATEST
            )), -3, KafkaOffsetResetStrategy.LATEST
        )
        self._check_timestamp_offsets_initializer(
            _build_source(KafkaOffsetsInitializer.timestamp(100)), 100
        )
        specified_offsets = {
            KafkaTopicPartition('test_topic1', 1): 1000,
            KafkaTopicPartition('test_topic2', 2): 2000
        }
        self._check_specified_offsets_initializer(
            _build_source(KafkaOffsetsInitializer.offsets(specified_offsets)), specified_offsets,
            KafkaOffsetResetStrategy.EARLIEST
        )
        self._check_specified_offsets_initializer(
            _build_source(KafkaOffsetsInitializer.offsets(
                specified_offsets,
                KafkaOffsetResetStrategy.LATEST
            )),
            specified_offsets,
            KafkaOffsetResetStrategy.LATEST
        )

    def test_bounded(self):
        def _build_source(initializer: KafkaOffsetsInitializer):
            return KafkaSource.builder() \
                .set_bootstrap_servers('localhost:9092') \
                .set_topics('test_topic') \
                .set_value_only_deserializer(SimpleStringSchema()) \
                .set_group_id('test_group') \
                .set_bounded(initializer) \
                .build()

        def _check_bounded(source: KafkaSource):
            self.assertEqual(
                get_field_value(source.get_java_function(), 'boundedness').toString(), 'BOUNDED'
            )

        self._test_set_bounded_or_unbounded(_build_source, _check_bounded)

    def test_unbounded(self):
        def _build_source(initializer: KafkaOffsetsInitializer):
            return KafkaSource.builder() \
                .set_bootstrap_servers('localhost:9092') \
                .set_topics('test_topic') \
                .set_value_only_deserializer(SimpleStringSchema()) \
                .set_group_id('test_group') \
                .set_unbounded(initializer) \
                .build()

        def _check_bounded(source: KafkaSource):
            self.assertEqual(
                get_field_value(source.get_java_function(), 'boundedness').toString(),
                'CONTINUOUS_UNBOUNDED'
            )

        self._test_set_bounded_or_unbounded(_build_source, _check_bounded)

    def _test_set_bounded_or_unbounded(self, _build_source, _check_boundedness):
        source = _build_source(KafkaOffsetsInitializer.latest())
        _check_boundedness(source)
        self._check_reader_handled_offsets_initializer(
            source, -1, KafkaOffsetResetStrategy.LATEST, False
        )
        source = _build_source(KafkaOffsetsInitializer.earliest())
        _check_boundedness(source)
        self._check_reader_handled_offsets_initializer(
            source, -2, KafkaOffsetResetStrategy.EARLIEST, False
        )
        source = _build_source(KafkaOffsetsInitializer.committed_offsets())
        _check_boundedness(source)
        self._check_reader_handled_offsets_initializer(
            source, -3, KafkaOffsetResetStrategy.NONE, False
        )
        source = _build_source(KafkaOffsetsInitializer.committed_offsets(
            KafkaOffsetResetStrategy.LATEST
        ))
        _check_boundedness(source)
        self._check_reader_handled_offsets_initializer(
            source, -3, KafkaOffsetResetStrategy.LATEST, False
        )
        source = _build_source(KafkaOffsetsInitializer.timestamp(100))
        _check_boundedness(source)
        self._check_timestamp_offsets_initializer(source, 100, False)
        specified_offsets = {
            KafkaTopicPartition('test_topic1', 1): 1000,
            KafkaTopicPartition('test_topic2', 2): 2000
        }
        source = _build_source(KafkaOffsetsInitializer.offsets(specified_offsets))
        _check_boundedness(source)
        self._check_specified_offsets_initializer(
            source, specified_offsets, KafkaOffsetResetStrategy.EARLIEST, False
        )
        source = _build_source(KafkaOffsetsInitializer.offsets(
            specified_offsets,
            KafkaOffsetResetStrategy.LATEST)
        )
        _check_boundedness(source)
        self._check_specified_offsets_initializer(
            source, specified_offsets, KafkaOffsetResetStrategy.LATEST, False
        )

    def test_set_value_only_deserializer(self):
        def _check(schema: DeserializationSchema, class_name: str):
            source = KafkaSource.builder() \
                .set_bootstrap_servers('localhost:9092') \
                .set_topics('test_topic') \
                .set_value_only_deserializer(schema) \
                .build()
            deserialization_schema_wrapper = get_field_value(source.get_java_function(),
                                                             'deserializationSchema')
            self.assertEqual(
                deserialization_schema_wrapper.getClass().getCanonicalName(),
                'org.apache.flink.connector.kafka.source.reader.deserializer'
                '.KafkaValueOnlyDeserializationSchemaWrapper'
            )
            deserialization_schema = get_field_value(deserialization_schema_wrapper,
                                                     'deserializationSchema')
            self.assertEqual(deserialization_schema.getClass().getCanonicalName(),
                             class_name)

        _check(SimpleStringSchema(), 'org.apache.flink.api.common.serialization.SimpleStringSchema')
        _check(ByteArraySchema(), 'org.apache.flink.api.common.serialization.ByteArraySchema')
        _check(
            JsonRowDeserializationSchema.builder().type_info(Types.ROW([Types.STRING()])).build(),
            'org.apache.flink.formats.json.JsonRowDeserializationSchema'
        )
        _check(
            CsvRowDeserializationSchema.Builder(Types.ROW([Types.STRING()])).build(),
            'org.apache.flink.formats.csv.CsvRowDeserializationSchema'
        )
        avro_schema_string = """
        {
            "type": "record",
            "name": "test_record",
            "fields": []
        }
        """
        _check(
            AvroRowDeserializationSchema(avro_schema_string=avro_schema_string),
            'org.apache.flink.formats.avro.AvroRowDeserializationSchema'
        )

    def _check_reader_handled_offsets_initializer(self,
                                                  source: KafkaSource,
                                                  offset: int,
                                                  reset_strategy: KafkaOffsetResetStrategy,
                                                  is_start: bool = True):
        if is_start:
            field_name = 'startingOffsetsInitializer'
        else:
            field_name = 'stoppingOffsetsInitializer'
        offsets_initializer = get_field_value(source.get_java_function(), field_name)
        self.assertEqual(
            offsets_initializer.getClass().getCanonicalName(),
            'org.apache.flink.connector.kafka.source.enumerator.initializer'
            '.ReaderHandledOffsetsInitializer'
        )

        starting_offset = get_field_value(offsets_initializer, 'startingOffset')
        self.assertEqual(starting_offset, offset)

        offset_reset_strategy = get_field_value(offsets_initializer, 'offsetResetStrategy')
        self.assertTrue(
            offset_reset_strategy.equals(reset_strategy._to_j_offset_reset_strategy())
        )

    def _check_timestamp_offsets_initializer(self,
                                             source: KafkaSource,
                                             timestamp: int,
                                             is_start: bool = True):
        if is_start:
            field_name = 'startingOffsetsInitializer'
        else:
            field_name = 'stoppingOffsetsInitializer'
        offsets_initializer = get_field_value(source.get_java_function(), field_name)
        self.assertEqual(
            offsets_initializer.getClass().getCanonicalName(),
            'org.apache.flink.connector.kafka.source.enumerator.initializer'
            '.TimestampOffsetsInitializer'
        )

        starting_timestamp = get_field_value(offsets_initializer, 'startingTimestamp')
        self.assertEqual(starting_timestamp, timestamp)

    def _check_specified_offsets_initializer(self,
                                             source: KafkaSource,
                                             offsets: Dict[KafkaTopicPartition, int],
                                             reset_strategy: KafkaOffsetResetStrategy,
                                             is_start: bool = True):
        if is_start:
            field_name = 'startingOffsetsInitializer'
        else:
            field_name = 'stoppingOffsetsInitializer'
        offsets_initializer = get_field_value(source.get_java_function(), field_name)
        self.assertEqual(
            offsets_initializer.getClass().getCanonicalName(),
            'org.apache.flink.connector.kafka.source.enumerator.initializer'
            '.SpecifiedOffsetsInitializer'
        )

        initial_offsets = get_field_value(offsets_initializer, 'initialOffsets')
        self.assertTrue(is_instance_of(initial_offsets, get_gateway().jvm.java.util.Map))
        self.assertEqual(initial_offsets.size(), len(offsets))
        for j_topic_partition in initial_offsets:
            topic_partition = KafkaTopicPartition(j_topic_partition.topic(),
                                                  j_topic_partition.partition())
            self.assertIsNotNone(offsets.get(topic_partition))
            self.assertEqual(initial_offsets[j_topic_partition], offsets[topic_partition])

        offset_reset_strategy = get_field_value(offsets_initializer, 'offsetResetStrategy')
        self.assertTrue(
            offset_reset_strategy.equals(reset_strategy._to_j_offset_reset_strategy())
        )

    @staticmethod
    def _get_kafka_source_configuration(source: KafkaSource):
        jvm = get_gateway().jvm
        j_source = source.get_java_function()
        j_to_configuration = j_source.getClass().getDeclaredMethod(
            'getConfiguration', to_jarray(jvm.java.lang.Class, [])
        )
        j_to_configuration.setAccessible(True)
        j_configuration = j_to_configuration.invoke(j_source, to_jarray(jvm.java.lang.Object, []))
        return Configuration(j_configuration=j_configuration)


@unittest.skip("Disable due to Kafka connector need to release a new version 2.0")
class KafkaSinkTests(PyFlinkStreamingTestCase):

    def test_compile(self):
        sink = KafkaSink.builder() \
            .set_bootstrap_servers('localhost:9092') \
            .set_record_serializer(self._build_serialization_schema()) \
            .build()

        ds = self.env.from_collection([], type_info=Types.STRING())
        ds.sink_to(sink)

        plan = json.loads(self.env.get_execution_plan())
        self.assertEqual(plan['nodes'][1]['type'], 'Sink: Writer')
        self.assertEqual(plan['nodes'][2]['type'], 'Sink: Committer')

    def test_set_bootstrap_severs(self):
        sink = KafkaSink.builder() \
            .set_bootstrap_servers('localhost:9092,localhost:9093') \
            .set_record_serializer(self._build_serialization_schema()) \
            .build()
        config = get_field_value(sink.get_java_function(), 'kafkaProducerConfig')
        self.assertEqual(config.get('bootstrap.servers'), 'localhost:9092,localhost:9093')

    def test_set_delivery_guarantee(self):
        sink = KafkaSink.builder() \
            .set_bootstrap_servers('localhost:9092') \
            .set_record_serializer(self._build_serialization_schema()) \
            .build()
        guarantee = get_field_value(sink.get_java_function(), 'deliveryGuarantee')
        self.assertEqual(guarantee.toString(), 'none')

        sink = KafkaSink.builder() \
            .set_bootstrap_servers('localhost:9092') \
            .set_delivery_guarantee(DeliveryGuarantee.AT_LEAST_ONCE) \
            .set_record_serializer(self._build_serialization_schema()) \
            .build()
        guarantee = get_field_value(sink.get_java_function(), 'deliveryGuarantee')
        self.assertEqual(guarantee.toString(), 'at-least-once')

        sink = KafkaSink.builder() \
            .set_bootstrap_servers('localhost:9092') \
            .set_delivery_guarantee(DeliveryGuarantee.EXACTLY_ONCE) \
            .set_record_serializer(self._build_serialization_schema()) \
            .build()
        guarantee = get_field_value(sink.get_java_function(), 'deliveryGuarantee')
        self.assertEqual(guarantee.toString(), 'exactly-once')

    def test_set_transactional_id_prefix(self):
        sink = KafkaSink.builder() \
            .set_bootstrap_servers('localhost:9092') \
            .set_transactional_id_prefix('test-prefix') \
            .set_record_serializer(self._build_serialization_schema()) \
            .build()
        prefix = get_field_value(sink.get_java_function(), 'transactionalIdPrefix')
        self.assertEqual(prefix, 'test-prefix')

    def test_set_property(self):
        sink = KafkaSink.builder() \
            .set_bootstrap_servers('localhost:9092') \
            .set_record_serializer(self._build_serialization_schema()) \
            .set_property('test-key', 'test-value') \
            .build()
        config = get_field_value(sink.get_java_function(), 'kafkaProducerConfig')
        self.assertEqual(config.get('test-key'), 'test-value')

    def test_set_record_serializer(self):
        sink = KafkaSink.builder() \
            .set_bootstrap_servers('localhost:9092') \
            .set_record_serializer(self._build_serialization_schema()) \
            .build()
        serializer = get_field_value(sink.get_java_function(), 'recordSerializer')
        self.assertEqual(serializer.getClass().getCanonicalName(),
                         'org.apache.flink.connector.kafka.sink.'
                         'KafkaRecordSerializationSchemaBuilder.'
                         'KafkaRecordSerializationSchemaWrapper')
        topic_selector = get_field_value(serializer, 'topicSelector')
        self.assertEqual(topic_selector.apply(None), 'test-topic')
        value_serializer = get_field_value(serializer, 'valueSerializationSchema')
        self.assertEqual(value_serializer.getClass().getCanonicalName(),
                         'org.apache.flink.api.common.serialization.SimpleStringSchema')

    @staticmethod
    def _build_serialization_schema() -> KafkaRecordSerializationSchema:
        return KafkaRecordSerializationSchema.builder() \
            .set_topic('test-topic') \
            .set_value_serialization_schema(SimpleStringSchema()) \
            .build()


@unittest.skip("Disable due to Kafka connector need to release a new version 2.0")
class KafkaRecordSerializationSchemaTests(PyFlinkTestCase):

    def test_set_topic(self):
        input_type = Types.ROW([Types.STRING()])

        serialization_schema = KafkaRecordSerializationSchema.builder() \
            .set_topic('test-topic') \
            .set_value_serialization_schema(
                JsonRowSerializationSchema.builder().with_type_info(input_type).build()) \
            .build()
        jvm = get_gateway().jvm
        serialization_schema._j_serialization_schema.open(
            jvm.org.apache.flink.connector.testutils.formats.DummyInitializationContext(),
            jvm.org.apache.flink.connector.kafka.sink.DefaultKafkaSinkContext(
                0, 1, jvm.java.util.Properties()))

        j_record = serialization_schema._j_serialization_schema.serialize(
            to_java_data_structure(Row('test')), None, None
        )
        self.assertEqual(j_record.topic(), 'test-topic')
        self.assertIsNone(j_record.key())
        self.assertEqual(j_record.value(), b'{"f0":"test"}')

    @unittest.skip("Disable due to Kafka connector need to release a new version 2.0")
    def test_set_topic_selector(self):
        def _select(data):
            data = data[0]
            if data == 'a':
                return 'topic-a'
            elif data == 'b':
                return 'topic-b'
            else:
                return 'topic-dead-letter'

        def _check_record(data, topic, serialized_data):
            input_type = Types.ROW([Types.STRING()])

            serialization_schema = KafkaRecordSerializationSchema.builder() \
                .set_topic_selector(_select) \
                .set_value_serialization_schema(
                    JsonRowSerializationSchema.builder().with_type_info(input_type).build()) \
                .build()
            jvm = get_gateway().jvm
            serialization_schema._j_serialization_schema.open(
                jvm.org.apache.flink.connector.testutils.formats.DummyInitializationContext(),
                jvm.org.apache.flink.connector.kafka.sink.DefaultKafkaSinkContext(
                    0, 1, jvm.java.util.Properties()))
            sink = KafkaSink.builder() \
                .set_bootstrap_servers('localhost:9092') \
                .set_record_serializer(serialization_schema) \
                .build()

            ds = MockDataStream(Types.ROW([Types.STRING()]))
            ds.sink_to(sink)
            row = Row(data)
            topic_row: Row = ds.feed(row)
            j_record = serialization_schema._j_serialization_schema.serialize(
                to_java_data_structure(topic_row), None, None
            )
            self.assertEqual(j_record.topic(), topic)
            self.assertIsNone(j_record.key())
            self.assertEqual(j_record.value(), serialized_data)

        _check_record('a', 'topic-a', b'{"f0":"a"}')
        _check_record('b', 'topic-b', b'{"f0":"b"}')
        _check_record('c', 'topic-dead-letter', b'{"f0":"c"}')
        _check_record('d', 'topic-dead-letter', b'{"f0":"d"}')

    def test_set_key_serialization_schema(self):
        def _check_key_serialization_schema(key_serialization_schema, expected_class):
            serialization_schema = KafkaRecordSerializationSchema.builder() \
                .set_topic('test-topic') \
                .set_key_serialization_schema(key_serialization_schema) \
                .set_value_serialization_schema(SimpleStringSchema()) \
                .build()
            schema_field = get_field_value(serialization_schema._j_serialization_schema,
                                           'keySerializationSchema')
            self.assertIsNotNone(schema_field)
            self.assertEqual(schema_field.getClass().getCanonicalName(), expected_class)

        self._check_serialization_schema_implementations(_check_key_serialization_schema)

    def test_set_value_serialization_schema(self):
        def _check_value_serialization_schema(value_serialization_schema, expected_class):
            serialization_schema = KafkaRecordSerializationSchema.builder() \
                .set_topic('test-topic') \
                .set_value_serialization_schema(value_serialization_schema) \
                .build()
            schema_field = get_field_value(serialization_schema._j_serialization_schema,
                                           'valueSerializationSchema')
            self.assertIsNotNone(schema_field)
            self.assertEqual(schema_field.getClass().getCanonicalName(), expected_class)

        self._check_serialization_schema_implementations(_check_value_serialization_schema)

    @staticmethod
    def _check_serialization_schema_implementations(check_function):
        input_type = Types.ROW([Types.STRING()])

        check_function(
            JsonRowSerializationSchema.builder().with_type_info(input_type).build(),
            'org.apache.flink.formats.json.JsonRowSerializationSchema'
        )
        check_function(
            CsvRowSerializationSchema.Builder(input_type).build(),
            'org.apache.flink.formats.csv.CsvRowSerializationSchema'
        )
        avro_schema_string = """
        {
            "type": "record",
            "name": "test_record",
            "fields": []
        }
        """
        check_function(
            AvroRowSerializationSchema(avro_schema_string=avro_schema_string),
            'org.apache.flink.formats.avro.AvroRowSerializationSchema'
        )
        check_function(
            SimpleStringSchema(),
            'org.apache.flink.api.common.serialization.SimpleStringSchema'
        )
        check_function(
            ByteArraySchema(),
            'org.apache.flink.api.common.serialization.ByteArraySchema'
        )


class MockDataStream(data_stream.DataStream):

    def __init__(self, original_type=None):
        super().__init__(None)
        self._operators = []
        self._type = original_type

    def feed(self, data):
        for op in self._operators:
            data = op(data)
        return data

    def get_type(self):
        return self._type

    def map(self, f, output_type=None):
        self._operators.append(f)
        self._type = output_type

    def sink_to(self, sink):
        ds = self
        from pyflink.datastream.connectors.base import SupportsPreprocessing
        if isinstance(sink, SupportsPreprocessing) and sink.get_transformer() is not None:
            ds = sink.get_transformer().apply(self)
        return ds
