From 256e95c90119d82a26d2526052af40e604d5762e Mon Sep 17 00:00:00 2001
From: Andrew Steinborn <git@steinborn.me>
Date: Sun, 14 May 2023 02:31:36 -0400
Subject: [PATCH] Fix flaky test

Passes locally but not in GitHub Actions
---
 .../proxy/scheduler/VelocityScheduler.java        | 15 ++++++++++++++-
 .../proxy/scheduler/VelocitySchedulerTest.java    |  3 ++-
 2 files changed, 16 insertions(+), 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
index c679b4055b..35cfabb999 100644
--- 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
index 9602e70e1d..70cd20b100 100644
--- 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());
   }