JUnit5 java FX test

JavaFXExtension.java

import javafx.application.Platform;
import org.junit.jupiter.api.extension.*;

import java.lang.reflect.Method;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.atomic.AtomicReference;

public class JavaFXExtension implements Extension, TestInstancePreConstructCallback, InvocationInterceptor {

  private boolean initialized = false;

  @Override
  public synchronized void preConstructTestInstance(TestInstanceFactoryContext testInstanceFactoryContext, ExtensionContext extensionContext) throws Exception {
    if (!initialized) {
      System.out.println("Starting JavaFX");
      var startLatch = new CountDownLatch(1);
      Platform.startup(startLatch::countDown);
      startLatch.await();
      initialized = true;
      System.out.println("Started JavaFX");
    }
  }

  @Override
  public void interceptTestMethod(Invocation<Void> invocation, ReflectiveInvocationContext<Method> invocationContext, ExtensionContext extensionContext) throws Throwable {
    var testLatch = new CountDownLatch(1);
    if (invocationContext.getExecutable().isAnnotationPresent(RunInFX.class)) {
      var thrown = new AtomicReference<Throwable>();
      Platform.runLater(() -> {
        try {
          invocation.proceed();
        } catch (Throwable throwable) {
          thrown.set(throwable);
        } finally {
          testLatch.countDown();
        }
      });
      testLatch.await();
      if (thrown.get() != null) {
        throw thrown.get();
      }
    } else {
      invocation.proceed();
    }
  }


}

RunInFX.java

import java.lang.annotation.Retention;
import java.lang.annotation.Target;

@Target({java.lang.annotation.ElementType.METHOD})
@Retention(java.lang.annotation.RetentionPolicy.RUNTIME)
public @interface RunInFX {
}

SampleTest.java

import javafx.concurrent.Task;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.extension.ExtendWith;

@ExtendWith(JavaFXExtension.class)
public class SampleTest {

  @Test
  @RunInFX
  void test() {
    new Task<>() {
      @Override
      protected Object call() throws Exception {
        return "Task creation would fail if not in FX Thread";
      }
    };
  }

  @Test
  void testNotInFXThread() {
    new Task<>() {
      @Override
      protected Object call() throws Exception {
        return "This fails actually";
      }
    };
  }

}