More fixes

linkedin · Jan 31, 2024 · 6aecf47 · 6aecf47
1 parent 2a054fd
commit 6aecf47
Showing 1 changed file with 8 additions and 0 deletions.
diff --git a/...i-client/src/test/java/com/linkedin/davinci/helix/VeniceLeaderFollowerStateModelTest.java b/...i-client/src/test/java/com/linkedin/davinci/helix/VeniceLeaderFollowerStateModelTest.java
@@ -11,9 +11,13 @@
 import com.linkedin.davinci.config.VeniceStoreVersionConfig;
 import com.linkedin.davinci.stats.ingestion.heartbeat.HeartbeatMonitoringService;
 import com.linkedin.venice.meta.Store;
+import com.linkedin.venice.meta.Version;
+import com.linkedin.venice.meta.VersionImpl;
 import io.tehuti.metrics.MetricsRepository;
 import java.util.HashSet;
+import java.util.Optional;
 import java.util.concurrent.CompletableFuture;
+import org.mockito.Mockito;
 import org.testng.annotations.Test;
 
 
@@ -41,6 +45,8 @@ protected StateModelIngestionProgressNotifier getNotifier() {
   @Test
   public void testOnBecomeFollowerFromOffline() throws Exception {
     // if the resource is not the current serving version, latch is not placed.
+    Version version = new VersionImpl("mockStore.getName()", 2, "");
+    when(mockStore.getVersion(Mockito.anyInt())).thenReturn(Optional.of(version));
     when(mockStore.getCurrentVersion()).thenReturn(2);
     testStateModel.onBecomeStandbyFromOffline(mockMessage, mockContext);
     verify(mockNotifier, never()).waitConsumptionCompleted(
@@ -76,6 +82,8 @@ public void testGracefulDropForCurrentVersionResource() {
 
   @Test
   public void testRemoveCVStateWhenBecomeOfflineFromStandby() {
+    Version version = new VersionImpl("mockStore.getName()", 2, "");
+    when(mockStore.getVersion(Mockito.anyInt())).thenReturn(Optional.of(version));
     when(mockStore.getCurrentVersion()).thenReturn(2);
     when(mockIngestionBackend.stopConsumption(any(VeniceStoreVersionConfig.class), eq(testPartition)))
         .thenReturn(CompletableFuture.completedFuture(null));