update ut

apache · Feb 18, 2025 · ec31df6 · ec31df6
1 parent 215ce5e
commit ec31df6
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 1 deletion.
diff --git a/worker/src/main/java/org/apache/celeborn/service/deploy/worker/memory/MemoryManager.java b/worker/src/main/java/org/apache/celeborn/service/deploy/worker/memory/MemoryManager.java
@@ -636,6 +636,11 @@ public interface ReadBufferTargetChangeListener {
     void onChange(long newMemoryTarget);
   }
 
+  @VisibleForTesting
+  public static void clearInstance() {
+    _INSTANCE = null;
+  }
+
   public enum ServingState {
     NONE_PAUSED,
     PUSH_AND_REPLICATE_PAUSED,

diff --git a/...pache/celeborn/service/deploy/worker/storage/memory/MemoryPartitionFilesSorterSuiteJ.java b/...pache/celeborn/service/deploy/worker/storage/memory/MemoryPartitionFilesSorterSuiteJ.java
@@ -114,9 +114,10 @@ public long[] prepare(int mapCount) {
     conf.set(CelebornConf.WORKER_READBUFFER_ALLOCATIONWAIT().key(), "10ms");
 
     StorageManager storageManager = Mockito.mock(StorageManager.class);
+    Mockito.when(storageManager.storageBufferAllocator()).thenReturn(allocator);
+    MemoryManager.clearInstance();
     MemoryManager.initialize(conf, storageManager, null);
     partitionDataWriter = Mockito.mock(PartitionDataWriter.class);
-    Mockito.when(storageManager.storageBufferAllocator()).thenReturn(allocator);
     when(partitionDataWriter.getMemoryFileInfo()).thenAnswer(i -> fileInfo);
 
     return partitionSize;