Fix flaky test

Passes locally but not in GitHub Actions
PaperMC · May 14, 2023 · 256e95c · 256e95c
1 parent 832eea0
commit 256e95c
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 2 deletions.
diff --git a/proxy/src/main/java/com/velocitypowered/proxy/scheduler/VelocityScheduler.java b/proxy/src/main/java/com/velocitypowered/proxy/scheduler/VelocityScheduler.java
@@ -35,6 +35,7 @@
 import java.util.IdentityHashMap;
 import java.util.Optional;
 import java.util.Set;
+import java.util.concurrent.ExecutionException;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
 import java.util.concurrent.ScheduledExecutorService;
@@ -45,6 +46,7 @@
 import org.apache.logging.log4j.Logger;
 import org.checkerframework.checker.nullness.qual.NonNull;
 import org.checkerframework.checker.nullness.qual.Nullable;
+import org.jetbrains.annotations.VisibleForTesting;
 
 /**
  * The Velocity "scheduler", which is actually a thin wrapper around
@@ -198,7 +200,8 @@ public ScheduledTask schedule() {
     }
   }
 
-  private class VelocityTask implements Runnable, ScheduledTask {
+  @VisibleForTesting
+  class VelocityTask implements Runnable, ScheduledTask {
 
     private final PluginContainer container;
     private final Runnable runnable;
@@ -296,6 +299,16 @@ public void run() {
     private void onFinish() {
       tasksByPlugin.remove(plugin(), this);
     }
+
+    public void awaitCompletion() {
+      try {
+        future.get();
+      } catch (InterruptedException e) {
+        Thread.currentThread().interrupt();
+      } catch (ExecutionException e) {
+        throw new RuntimeException(e);
+      }
+    }
   }
 
   private static class Log {

diff --git a/proxy/src/test/java/com/velocitypowered/proxy/scheduler/VelocitySchedulerTest.java b/proxy/src/test/java/com/velocitypowered/proxy/scheduler/VelocitySchedulerTest.java
@@ -21,6 +21,7 @@
 
 import com.velocitypowered.api.scheduler.ScheduledTask;
 import com.velocitypowered.api.scheduler.TaskStatus;
+import com.velocitypowered.proxy.scheduler.VelocityScheduler.VelocityTask;
 import com.velocitypowered.proxy.testutil.FakePluginManager;
 import java.time.Duration;
 import java.util.concurrent.CountDownLatch;
@@ -39,6 +40,7 @@ void buildTask() throws Exception {
     ScheduledTask task = scheduler.buildTask(FakePluginManager.PLUGIN_A, latch::countDown)
         .schedule();
     latch.await();
+    ((VelocityTask) task).awaitCompletion();
     assertEquals(TaskStatus.FINISHED, task.status());
   }
 
@@ -50,7 +52,6 @@ void cancelWorks() throws Exception {
         .delay(100, TimeUnit.SECONDS)
         .schedule();
     task.cancel();
-    Thread.sleep(200);
     assertEquals(3, i.get());
     assertEquals(TaskStatus.CANCELLED, task.status());
   }