fixed test failure

linkedin · Sep 29, 2024 · 70dd36d · 70dd36d
1 parent d99a2a6
commit 70dd36d
Show file tree

Hide file tree

Showing 4 changed files with 41 additions and 9 deletions.
diff --git a/...nice-common/src/main/java/com/linkedin/venice/pushstatushelper/PushStatusStoreReader.java b/...nice-common/src/main/java/com/linkedin/venice/pushstatushelper/PushStatusStoreReader.java
@@ -230,7 +230,10 @@ public long getHeartbeat(String storeName, String instanceName) {
   }
 
   public boolean isInstanceAlive(String storeName, String instanceName) {
-    long lastReportTimeStamp = getHeartbeat(storeName, instanceName);
+    return isInstanceAlive(getHeartbeat(storeName, instanceName));
+  }
+
+  boolean isInstanceAlive(long lastReportTimeStamp) {
     return System.currentTimeMillis() - lastReportTimeStamp <= TimeUnit.SECONDS
         .toMillis(heartbeatExpirationTimeInSeconds);
   }
@@ -240,11 +243,8 @@ public InstanceStatus getInstanceStatus(String storeName, String instanceName) {
     if (lastReportTimeStamp < 0) {
       return InstanceStatus.BOOTSTRAPPING;
     }
-    if (System.currentTimeMillis() - lastReportTimeStamp <= TimeUnit.SECONDS
-        .toMillis(heartbeatExpirationTimeInSeconds)) {
-      return InstanceStatus.ALIVE;
-    }
-    return InstanceStatus.DEAD;
+
+    return isInstanceAlive(lastReportTimeStamp) ? InstanceStatus.ALIVE : InstanceStatus.DEAD;
   }
 
   public Map<CharSequence, Integer> getSupposedlyOngoingIncrementalPushVersions(String storeName, int storeVersion) {

diff --git a/...-common/src/test/java/com/linkedin/venice/pushstatushelper/PushStatusStoreReaderTest.java b/...-common/src/test/java/com/linkedin/venice/pushstatushelper/PushStatusStoreReaderTest.java
@@ -6,13 +6,15 @@
 import static org.mockito.ArgumentMatchers.*;
 import static org.mockito.Mockito.any;
 import static org.mockito.Mockito.anySet;
+import static org.mockito.Mockito.doCallRealMethod;
 import static org.mockito.Mockito.doReturn;
 import static org.mockito.Mockito.eq;
 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;
+import static org.testng.Assert.assertEquals;
 import static org.testng.Assert.assertEqualsDeep;
 import static org.testng.Assert.assertNotEquals;
 
@@ -365,4 +367,22 @@ public void testNullResponseWhenVersionLevelKeyIsNotWritten()
     // Test that push status store reader will also return null instead of empty map in this case
     Assert.assertNull(storeReaderSpy.getVersionStatus(storeName, storeVersion));
   }
+
+  @Test
+  public void testGetInstanceStatus() {
+    PushStatusStoreReader mockReader = mock(PushStatusStoreReader.class);
+    doCallRealMethod().when(mockReader).getInstanceStatus(any(), any());
+
+    doReturn(-1l).when(mockReader).getHeartbeat("store_1", "instance_1");
+    assertEquals(
+        mockReader.getInstanceStatus("store_1", "instance_1"),
+        PushStatusStoreReader.InstanceStatus.BOOTSTRAPPING);
+
+    doReturn(1000l).when(mockReader).getHeartbeat("store_1", "instance_1");
+    doReturn(true).when(mockReader).isInstanceAlive(anyLong());
+    assertEquals(mockReader.getInstanceStatus("store_1", "instance_1"), PushStatusStoreReader.InstanceStatus.ALIVE);
+
+    doReturn(false).when(mockReader).isInstanceAlive(anyLong());
+    assertEquals(mockReader.getInstanceStatus("store_1", "instance_1"), PushStatusStoreReader.InstanceStatus.DEAD);
+  }
 }
diff --git a/...-common/src/test/java/com/linkedin/venice/pushstatushelper/PushStatusStoreWriterTest.java b/...-common/src/test/java/com/linkedin/venice/pushstatushelper/PushStatusStoreWriterTest.java
@@ -144,4 +144,13 @@ public void testWriteHeartbeat() {
     verify(veniceWriterMock)
         .update(eq(statusKey), eq(getHeartbeatRecord(heartbeat)), eq(valueSchemaId), eq(derivedSchemaId), eq(null));
   }
+
+  @Test
+  public void testWriteHeartbeatForBootstrappingInstance() {
+    PushStatusKey statusKey = PushStatusStoreUtils.getHeartbeatKey(instanceName);
+    pushStatusStoreWriter.writeHeartbeatForBootstrappingInstance(storeName);
+    verify(veniceWriterCacheMock).prepareVeniceWriter(storeName);
+    verify(veniceWriterMock)
+        .update(eq(statusKey), eq(getHeartbeatRecord(-1)), eq(valueSchemaId), eq(derivedSchemaId), eq(null));
+  }
 }
diff --git a/...ice-controller/src/test/java/com/linkedin/venice/pushmonitor/PushStatusCollectorTest.java b/...ice-controller/src/test/java/com/linkedin/venice/pushmonitor/PushStatusCollectorTest.java
@@ -109,7 +109,8 @@ public void testPushStatusCollector() {
         .thenReturn(startedInstancePushStatus, dvcTooManyDeadInstancesErrorInstancePushStatus);
     when(pushStatusStoreReader.getPartitionStatus(daVinciStoreName, 11, 0, Optional.empty()))
         .thenReturn(startedInstancePushStatus, dvcOtherErrorInstancePushStatus);
-    when(pushStatusStoreReader.isInstanceAlive(daVinciStoreName, "instance")).thenReturn(true);
+    when(pushStatusStoreReader.getInstanceStatus(daVinciStoreName, "instance"))
+        .thenReturn(PushStatusStoreReader.InstanceStatus.ALIVE);
     pushStatusCollector.subscribeTopic(daVinciStoreTopicV1, 1);
     Assert.assertFalse(pushStatusCollector.getTopicToPushStatusMap().containsKey(daVinciStoreTopicV1));
 
@@ -241,7 +242,8 @@ public void testPushStatusCollectorDaVinciStatusPollingRetry() {
         .thenReturn(Collections.emptyMap(), startedInstancePushStatus, dvcTooManyDeadInstancesErrorInstancePushStatus);
     when(pushStatusStoreReader.getPartitionStatus(daVinciStoreName, 6, 0, Optional.empty()))
         .thenReturn(Collections.emptyMap(), startedInstancePushStatus, dvcOtherErrorInstancePushStatus);
-    when(pushStatusStoreReader.isInstanceAlive(daVinciStoreName, "instance")).thenReturn(true);
+    when(pushStatusStoreReader.getInstanceStatus(daVinciStoreName, "instance"))
+        .thenReturn(PushStatusStoreReader.InstanceStatus.ALIVE);
     pushStatusCollector.subscribeTopic(daVinciStoreTopicV1, 1);
     Assert.assertFalse(pushStatusCollector.getTopicToPushStatusMap().containsKey(daVinciStoreTopicV1));
 
@@ -330,7 +332,8 @@ public void testPushStatusCollectorDaVinciStatusPollingRetryWhenEmptyResultUntil
 
     when(pushStatusStoreReader.getPartitionStatus(daVinciStoreName, 2, 0, Optional.empty()))
         .thenReturn(Collections.emptyMap());
-    when(pushStatusStoreReader.isInstanceAlive(daVinciStoreName, "instance")).thenReturn(true);
+    when(pushStatusStoreReader.getInstanceStatus(daVinciStoreName, "instance"))
+        .thenReturn(PushStatusStoreReader.InstanceStatus.ALIVE);
     pushStatusCollector.subscribeTopic(daVinciStoreTopicV1, 1);
     Assert.assertFalse(pushStatusCollector.getTopicToPushStatusMap().containsKey(daVinciStoreTopicV1));