Unit test modification

linkedin · Feb 7, 2025 · 680faec · 680faec
1 parent 206addf
commit 680faec
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 7 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
@@ -401,7 +401,7 @@ private Function<String, Boolean> functionToCheckWhetherStorageEngineShouldBeKep
 
   private synchronized void bootstrap() {
     List<AbstractStorageEngine> storageEngines =
-        storageService.getStorageEngineRepository().getAllLocalStorageEngines();
+        getStorageService().getStorageEngineRepository().getAllLocalStorageEngines();
     LOGGER.info("Starting bootstrap, storageEngines: {}", storageEngines);
     Map<String, Version> storeNameToBootstrapVersionMap = new HashMap<>();
     Map<String, List<Integer>> storeNameToPartitionListMap = new HashMap<>();
@@ -440,7 +440,7 @@ private synchronized void bootstrap() {
       if (!(storeNameToBootstrapVersionMap.containsKey(storeName)
           && (storeNameToBootstrapVersionMap.get(storeName).getNumber() < versionNumber))) {
         storeNameToBootstrapVersionMap.put(storeName, version);
-        storeNameToPartitionListMap.put(storeName, storageService.getUserPartitions(kafkaTopicName));
+        storeNameToPartitionListMap.put(storeName, getStorageService().getUserPartitions(kafkaTopicName));
       }
     }
 
@@ -455,12 +455,12 @@ private synchronized void bootstrap() {
          * In this case we will only need to close metadata partition, as it is supposed to be opened and managed by
          * forked ingestion process via following subscribe call.
          */
-        for (AbstractStorageEngine storageEngine: storageService.getStorageEngineRepository()
+        for (AbstractStorageEngine storageEngine: getStorageService().getStorageEngineRepository()
             .getAllLocalStorageEngines()) {
           storageEngine.closeMetadataPartition();
         }
       } else {
-        storageService.closeAllStorageEngines();
+        getStorageService().closeAllStorageEngines();
       }
     }
 
@@ -470,13 +470,13 @@ private synchronized void bootstrap() {
             metricsRepository,
             storageMetadataService,
             ingestionService,
-            storageService,
+            getStorageService(),
             blobTransferManager,
             this::getVeniceCurrentVersionNumber)
         : new DefaultIngestionBackend(
             storageMetadataService,
             ingestionService,
-            storageService,
+            getStorageService(),
             blobTransferManager,
             configLoader.getVeniceServerConfig());
     ingestionBackend.addIngestionNotifier(ingestionListener);
@@ -487,7 +487,7 @@ private synchronized void bootstrap() {
         List<Integer> partitions = storeNameToPartitionListMap.get(storeName);
         String versionTopic = version.kafkaTopicName();
         LOGGER.info("Bootstrapping partitions {} for {}", partitions, versionTopic);
-        AbstractStorageEngine storageEngine = storageService.getStorageEngine(versionTopic);
+        AbstractStorageEngine storageEngine = getStorageService().getStorageEngine(versionTopic);
         aggVersionedStorageEngineStats.setStorageEngine(versionTopic, storageEngine);
         StoreBackend storeBackend = getStoreOrThrow(storeName);
         storeBackend.subscribe(ComplementSet.newSet(partitions), Optional.of(version));

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
@@ -189,6 +189,7 @@ public void testBootstrappingSubscription()
     List<AbstractStorageEngine> localStorageEngines = new ArrayList<>();
     localStorageEngines.add(abstractStorageEngine);
 
+    when(backend.getStorageService()).thenReturn(mockStorageService);
     when(mockStorageService.getStorageEngineRepository()).thenReturn(mockStorageEngineRepository);
     when(mockStorageService.getStorageEngine(resourceName)).thenReturn(abstractStorageEngine);
     when(mockStorageEngineRepository.getAllLocalStorageEngines()).thenReturn(localStorageEngines);