[BEAM-12308] change expected value in kakfa IT (#17740)

apache · May 24, 2022 · bb8dafb · bb8dafb
1 parent d9436c4
commit bb8dafb
Showing 1 changed file with 5 additions and 4 deletions.
diff --git a/sdks/python/apache_beam/io/external/xlang_kafkaio_it_test.py b/sdks/python/apache_beam/io/external/xlang_kafkaio_it_test.py
@@ -124,7 +124,7 @@ def test_kafkaio_populated_key(self):
           bootstrap_servers, kafka_topic, False)
 
       self.run_kafka_write(pipeline_creator)
-      self.run_kafka_read(pipeline_creator)
+      self.run_kafka_read(pipeline_creator, b'key')
 
   def test_kafkaio_null_key(self):
     kafka_topic = 'xlang_kafkaio_test_null_key_{}'.format(uuid.uuid4())
@@ -136,20 +136,21 @@ def test_kafkaio_null_key(self):
           bootstrap_servers, kafka_topic, True)
 
       self.run_kafka_write(pipeline_creator)
-      self.run_kafka_read(pipeline_creator)
+      self.run_kafka_read(pipeline_creator, None)
 
   def run_kafka_write(self, pipeline_creator):
     with TestPipeline() as pipeline:
       pipeline.not_use_test_runner_api = True
       pipeline_creator.build_write_pipeline(pipeline)
 
-  def run_kafka_read(self, pipeline_creator):
+  def run_kafka_read(self, pipeline_creator, expected_key):
     with TestPipeline() as pipeline:
       pipeline.not_use_test_runner_api = True
       result = pipeline_creator.build_read_pipeline(pipeline, NUM_RECORDS)
       assert_that(
           result,
-          equal_to([(b'', str(i).encode()) for i in range(NUM_RECORDS)]))
+          equal_to([(expected_key, str(i).encode())
+                    for i in range(NUM_RECORDS)]))
 
   def get_platform_localhost(self):
     if sys.platform == 'darwin':