[incubator-uniffle] 02/03: [#408] test: fix memory check failure in ShuffleBufferManagerTest#bufferSizeTest (#657)

[zuston]

This is an automated email from the ASF dual-hosted git repository.

zuston pushed a commit to branch branch-0.7
in repository https://gitbox.apache.org/repos/asf/incubator-uniffle.git

commit 2b43922e5017f29fe5898e4e1d4bcae47bd76cfa
Author: Junfan Zhang <zus...@apache.org>
AuthorDate: Mon Feb 27 16:17:51 2023 +0800

    [#408] test: fix memory check failure in 
ShuffleBufferManagerTest#bufferSizeTest (#657)
    
    ### What changes were proposed in this pull request?
    
    Fix memory check failure due to the execution sequence in 
ShuffleBufferManagerTest#bufferSizeTest
    
    ### Why are the changes needed?
    
    Flaky test fix: #408
    
    ### Does this PR introduce _any_ user-facing change?
    
    No.
    
    ### How was this patch tested?
    
    Existing UTs
---
 .../org/apache/uniffle/server/buffer/ShuffleBufferManagerTest.java | 7 +++++++
 1 file changed, 7 insertions(+)

diff --git 
a/server/src/test/java/org/apache/uniffle/server/buffer/ShuffleBufferManagerTest.java
 
b/server/src/test/java/org/apache/uniffle/server/buffer/ShuffleBufferManagerTest.java
index 00d8467e..670e684e 100644
--- 
a/server/src/test/java/org/apache/uniffle/server/buffer/ShuffleBufferManagerTest.java
+++ 
b/server/src/test/java/org/apache/uniffle/server/buffer/ShuffleBufferManagerTest.java
@@ -18,11 +18,13 @@
 package org.apache.uniffle.server.buffer;
 
 import java.io.File;
+import java.time.Duration;
 import java.util.Arrays;
 import java.util.Map;
 import java.util.concurrent.atomic.AtomicLong;
 
 import com.google.common.collect.RangeMap;
+import org.awaitility.Awaitility;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.io.TempDir;
@@ -621,5 +623,10 @@ public class ShuffleBufferManagerTest extends 
BufferTestBase {
         fail("Flush data time out");
       }
     } while (committedCount < expectedBlockNum);
+
+    // Need to wait for `event.doCleanup` to be executed
+    // to ensure the correctness of subsequent checks of
+    // `shuffleBufferManager.getUsedMemory()` and 
`shuffleBufferManager.getInFlushSize()`.
+    Awaitility.await().atMost(Duration.ofSeconds(5)).until(() -> 
shuffleBufferManager.getUsedMemory() == 0);
   }
 }