Remove 2nd case

sushantmane · Nov 23, 2024 · de38ec6 · de38ec6
1 parent 2842dd0
commit de38ec6
Show file tree

Hide file tree

Showing 2 changed files with 19 additions and 16 deletions.
diff --git a/...ent/src/main/java/com/linkedin/davinci/kafka/consumer/ActiveActiveStoreIngestionTask.java b/...ent/src/main/java/com/linkedin/davinci/kafka/consumer/ActiveActiveStoreIngestionTask.java
@@ -1111,18 +1111,18 @@ public void consumerSubscribe(PubSubTopicPartition pubSubTopicPartition, long st
     VeniceServerConfig serverConfig = getServerConfig();
     if (isDaVinciClient() || serverConfig.getKafkaClusterUrlToIdMap().containsKey(pubSubAddress)) {
       super.consumerSubscribe(pubSubTopicPartition, startOffset, pubSubAddress);
-    } else {
-      LOGGER.error(
-          "PubSub address: {} is not in the pubsub cluster map: {}. Cannot subscribe to topic-partition: {}",
-          pubSubAddress,
-          serverConfig.getKafkaClusterUrlToIdMap(),
-          pubSubTopicPartition);
-      throw new VeniceException(
-          String.format(
-              "PubSub address: %s is not in the pubsub cluster map. Cannot subscribe to topic-partition: %s",
-              pubSubAddress,
-              pubSubTopicPartition));
+      return;
     }
+    LOGGER.error(
+        "PubSub address: {} is not in the pubsub cluster map: {}. Cannot subscribe to topic-partition: {}",
+        pubSubAddress,
+        serverConfig.getKafkaClusterUrlToIdMap(),
+        pubSubTopicPartition);
+    throw new VeniceException(
+        String.format(
+            "PubSub address: %s is not in the pubsub cluster map. Cannot subscribe to topic-partition: %s",
+            pubSubAddress,
+            pubSubTopicPartition));
   }
 
   private long calculateRewindStartTime(PartitionConsumptionState partitionConsumptionState) {

diff --git a/...src/test/java/com/linkedin/davinci/kafka/consumer/ActiveActiveStoreIngestionTaskTest.java b/...src/test/java/com/linkedin/davinci/kafka/consumer/ActiveActiveStoreIngestionTaskTest.java
@@ -12,6 +12,7 @@
 import static org.mockito.ArgumentMatchers.eq;
 import static org.mockito.Mockito.doCallRealMethod;
 import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.spy;
 import static org.mockito.Mockito.times;
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
@@ -741,7 +742,7 @@ public void getKeyLevelLockMaxPoolSizeBasedOnServerConfigTest() {
   @Test
   public void testConsumerSubscribe() {
     // Case 1: AA store ingestion task with invalid pubsub address
-    ActiveActiveStoreIngestionTask ingestionTask = mock(ActiveActiveStoreIngestionTask.class);
+    ActiveActiveStoreIngestionTask ingestionTask = spy(ActiveActiveStoreIngestionTask.class);
     VeniceServerConfig mockServerConfig = mock(VeniceServerConfig.class);
 
     when(ingestionTask.getServerConfig()).thenReturn(mockServerConfig);
@@ -764,17 +765,19 @@ public void testConsumerSubscribe() {
     verify(ingestionTask, times(1)).consumerSubscribe(pubSubTopicPartition, 100L, "invalidPubSubAddress");
 
     // Case 2: DaVinci client
-    ActiveActiveStoreIngestionTask dvcIngestionTask = mock(ActiveActiveStoreIngestionTask.class);
+    ActiveActiveStoreIngestionTask dvcIngestionTask = spy(ActiveActiveStoreIngestionTask.class);
     doCallRealMethod().when(dvcIngestionTask).consumerSubscribe(any(), anyLong(), anyString());
     when(dvcIngestionTask.getServerConfig()).thenReturn(mockServerConfig);
     when(dvcIngestionTask.isDaVinciClient()).thenReturn(true);
     when(mockServerConfig.getKafkaClusterUrlToIdMap()).thenReturn(Object2IntMaps.emptyMap());
     try {
       dvcIngestionTask.consumerSubscribe(pubSubTopicPartition, 100L, "validPubSubAddress");
     } catch (Exception e) {
-      assertFalse(
-          e.getMessage().contains("is not in the pubsub cluster map"),
-          "Exception message should not contain the expected message but found: " + e.getMessage());
+      if (e.getMessage() != null) {
+        assertFalse(
+            e.getMessage().contains("is not in the pubsub cluster map"),
+            "Exception message should not contain the expected message but found: " + e.getMessage());
+      }
     }
     verify(dvcIngestionTask, times(1)).consumerSubscribe(pubSubTopicPartition, 100L, "validPubSubAddress");
   }