Finish fixing dynamic.source tests

pinterest · Dec 12, 2024 · 5c95746 · 5c95746
1 parent d325fb4
commit 5c95746
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 11 deletions.
diff --git a/...nterest/flink/connector/psc/dynamic/source/enumerator/DynamicPscSourceEnumeratorTest.java b/...nterest/flink/connector/psc/dynamic/source/enumerator/DynamicPscSourceEnumeratorTest.java
@@ -837,17 +837,17 @@ private void verifyAllSplitsHaveBeenAssigned(
                 clusterToTopicPartition.values().stream()
                         .flatMap(Collection::stream)
                         .collect(Collectors.toSet());
-        for (Set<String> topics : clusterTopicMap.values()) {
-            for (String topic : topics) {
-                Set<TopicUriPartition> expectedTopicPartitions = new HashSet<>();
+        Set<TopicUriPartition> expectedTopicPartitions = new HashSet<>();
+        for (Map.Entry<String, Set<String>> entry : clusterTopicMap.entrySet()) {
+            for (String topic : entry.getValue()) {
                 for (int i = 0; i < NUM_SPLITS_PER_CLUSTER; i++) {
-                    expectedTopicPartitions.add(new TopicUriPartition(PscTestEnvironmentWithKafkaAsPubSub.PSC_TEST_CLUSTER0_URI_PREFIX + topic, i));
+                    expectedTopicPartitions.add(new TopicUriPartition(kafkaStream.getClusterMetadataMap().get(entry.getKey()).getClusterUriString() + topic, i));
                 }
-                assertThat(assignedTopicPartitionSet)
-                        .as("splits must contain all topics and 2 partitions per topic")
-                        .containsExactlyInAnyOrderElementsOf(expectedTopicPartitions);
             }
         }
+        assertThat(assignedTopicPartitionSet)
+                .as("splits must contain all topics and 2 partitions per topic")
+                .containsExactlyInAnyOrderElementsOf(expectedTopicPartitions);
     }
 
     private Map<Integer, Set<DynamicPscSourceSplit>> getReaderAssignments(

diff --git a/...a/com/pinterest/flink/connector/psc/dynamic/source/reader/DynamicPscSourceReaderTest.java b/...a/com/pinterest/flink/connector/psc/dynamic/source/reader/DynamicPscSourceReaderTest.java
@@ -25,6 +25,7 @@
 import com.pinterest.flink.connector.psc.source.reader.deserializer.PscRecordDeserializationSchema;
 import com.pinterest.flink.connector.psc.source.split.PscTopicUriPartitionSplit;
 import com.pinterest.flink.streaming.connectors.psc.DynamicPscSourceTestHelperWithKafkaAsPubSub;
+import com.pinterest.flink.streaming.connectors.psc.PscTestBaseWithKafkaAsPubSub;
 import com.pinterest.flink.streaming.connectors.psc.PscTestEnvironmentWithKafkaAsPubSub;
 import com.pinterest.psc.common.TopicUriPartition;
 import com.pinterest.psc.config.PscConfiguration;
@@ -60,7 +61,6 @@
  */
 public class DynamicPscSourceReaderTest extends SourceReaderTestBase<DynamicPscSourceSplit> {
     private static final String TOPIC = "DynamicPscSourceReaderTest";
-    private static final String TOPIC_URI = PscTestEnvironmentWithKafkaAsPubSub.PSC_TEST_CLUSTER0_URI_PREFIX + TOPIC;
 
     // we are testing two clusters and SourceReaderTestBase expects there to be a total of 10 splits
     private static final int NUM_SPLITS_PER_CLUSTER = 5;
@@ -316,16 +316,19 @@ protected DynamicPscSourceSplit getSplit(
 
         String kafkaClusterId;
         int splitIdForCluster = splitId % NUM_SPLITS_PER_CLUSTER;
+        int clusterIdx;
         if (splitId < NUM_SPLITS_PER_CLUSTER) {
-            kafkaClusterId = "pubsub-cluster-0";
+            clusterIdx = 0;
         } else {
-            kafkaClusterId = "pubsub-cluster-1";
+            clusterIdx = 1;
         }
 
+        kafkaClusterId = "pubsub-cluster-" + clusterIdx;
+
         return new DynamicPscSourceSplit(
                 kafkaClusterId,
                 new PscTopicUriPartitionSplit(
-                        new TopicUriPartition(TOPIC_URI, splitIdForCluster), 0L, stoppingOffset));
+                        new TopicUriPartition(PscTestBaseWithKafkaAsPubSub.clusterUris.get(clusterIdx) + TOPIC, splitIdForCluster), 0L, stoppingOffset));
     }
 
     @Override