[dvc] Config flag for subscribing on disk partitions automatically. U…

…pdated with integration test and unit test.
kristyelee · Feb 5, 2025 · 206addf · 206addf
1 parent 388d2af
commit 206addf
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 6 deletions.
diff --git a/clients/da-vinci-client/src/main/java/com/linkedin/davinci/DaVinciBackend.java b/clients/da-vinci-client/src/main/java/com/linkedin/davinci/DaVinciBackend.java
@@ -399,7 +399,7 @@ private Function<String, Boolean> functionToCheckWhetherStorageEngineShouldBeKep
     };
   }
 
-  final synchronized void bootstrap() {
+  private synchronized void bootstrap() {
     List<AbstractStorageEngine> storageEngines =
         storageService.getStorageEngineRepository().getAllLocalStorageEngines();
     LOGGER.info("Starting bootstrap, storageEngines: {}", storageEngines);

diff --git a/clients/da-vinci-client/src/test/java/com/linkedin/davinci/DaVinciBackendTest.java b/clients/da-vinci-client/src/test/java/com/linkedin/davinci/DaVinciBackendTest.java
@@ -9,7 +9,6 @@
 import static org.mockito.ArgumentMatchers.anyInt;
 import static org.mockito.ArgumentMatchers.anyString;
 import static org.mockito.Mockito.doAnswer;
-import static org.mockito.Mockito.doCallRealMethod;
 import static org.mockito.Mockito.doNothing;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.times;
@@ -38,6 +37,8 @@
 import com.linkedin.venice.utils.ComplementSet;
 import com.linkedin.venice.utils.VeniceProperties;
 import java.lang.reflect.Field;
+import java.lang.reflect.InvocationTargetException;
+import java.lang.reflect.Method;
 import java.util.ArrayList;
 import java.util.HashSet;
 import java.util.List;
@@ -168,7 +169,8 @@ public void testBootstrappingAwareCompletableFuture()
   }
 
   @Test
-  public void testBootstrappingSubscription() throws NoSuchFieldException, IllegalAccessException {
+  public void testBootstrappingSubscription()
+      throws IllegalAccessException, NoSuchFieldException, NoSuchMethodException, InvocationTargetException {
     DaVinciBackend backend = mock(DaVinciBackend.class);
     StorageService mockStorageService = mock(StorageService.class);
 
@@ -254,18 +256,20 @@ public Object answer(InvocationOnMock invocation) throws Throwable {
     ingestionServiceField.set(backend, storeIngestionService);
     doNothing().when(ingestionBackend).addIngestionNotifier(any());
 
+    Method bootstrapMethod = DaVinciBackend.class.getDeclaredMethod("bootstrap");
+    bootstrapMethod.setAccessible(true);
+
     // DA_VINCI_SUBSCRIBE_ON_DISK_PARTITIONS_AUTOMATICALLY == false
     when(mockCombinedProperties.getBoolean(anyString(), anyBoolean())).thenReturn(false);
-    doCallRealMethod().when(backend).bootstrap();
-    backend.bootstrap();
+    bootstrapMethod.invoke(backend);
 
     ComplementSet<Integer> subscription = mockStoreBackend.getSubscription();
     assertTrue(subscription.contains(0));
     assertTrue(subscription.contains(1));
     assertFalse(subscription.contains(2));
 
     when(mockCombinedProperties.getBoolean(anyString(), anyBoolean())).thenReturn(true);
-    backend.bootstrap();
+    bootstrapMethod.invoke(backend);
 
     // DA_VINCI_SUBSCRIBE_ON_DISK_PARTITIONS_AUTOMATICALLY == true
     subscription = mockStoreBackend.getSubscription();