chore: track Playwright instances in PlaywrightRegistry (#1495)

playwright/src/main/java/com/microsoft/playwright/junit/impl/PlaywrightExtension.java

@@ -2,54 +2,54 @@
 
 import com.microsoft.playwright.Playwright;
 import com.microsoft.playwright.junit.Options;
-import org.junit.jupiter.api.extension.*;
+import org.junit.jupiter.api.extension.ExtensionContext;
+import org.junit.jupiter.api.extension.ParameterContext;
+import org.junit.jupiter.api.extension.ParameterResolutionException;
+import org.junit.jupiter.api.extension.ParameterResolver;
 
 import java.util.ArrayList;
 import java.util.Collections;
 import java.util.List;
-import java.util.concurrent.locks.Lock;
-import java.util.concurrent.locks.ReentrantLock;
 
 import static com.microsoft.playwright.junit.impl.ExtensionUtils.isParameterSupported;
 import static com.microsoft.playwright.junit.impl.ExtensionUtils.setTestIdAttribute;
 
-public class PlaywrightExtension implements ParameterResolver, BeforeAllCallback, ExtensionContext.Store.CloseableResource {
-  private static final ThreadLocal<Playwright> threadLocalPlaywright;
-  private static final List<Playwright> playwrightList;
-  private final static Lock beforeLock;
-  private static boolean isTestRunStarted;
+public class PlaywrightExtension implements ParameterResolver {
+  private static final ThreadLocal<Playwright> threadLocalPlaywright = new ThreadLocal<>();
   private static final ExtensionContext.Namespace namespace = ExtensionContext.Namespace.create(PlaywrightExtension.class);
 
-  static {
-    isTestRunStarted = false;
-    beforeLock = new ReentrantLock();
-    threadLocalPlaywright = new ThreadLocal<>();
-    playwrightList = Collections.synchronizedList(new ArrayList<>());
-  }
+  // There should be at most one instance of PlaywrgihtRegistry per test run, it keeps
+  // track of all created Playwright instances and calls `close()` on each of them after
+  // the tests finished.
+  static class PlaywrgihtRegistry implements ExtensionContext.Store.CloseableResource {
+    private final List<Playwright> playwrightList = Collections.synchronizedList(new ArrayList<>());
 
-  // Before a Test class starts, we register a closeable resource
-  // We use a lock + boolean to ensure this only gets run once for the entire test run
-  @Override
-  public void beforeAll(final ExtensionContext context) {
-    try {
-      beforeLock.lock();
-      if (!isTestRunStarted) {
-        isTestRunStarted = true;
-        context.getRoot().getStore(namespace).put(context.getUniqueId(), this);
+    static synchronized PlaywrgihtRegistry getOrCreateFor(ExtensionContext extensionContext) {
+      ExtensionContext.Store rootStore = extensionContext.getRoot().getStore(namespace);
+      PlaywrgihtRegistry instance = (PlaywrgihtRegistry) rootStore.get(PlaywrgihtRegistry.class);
+      if (instance == null) {
+        instance = new PlaywrgihtRegistry();
+        rootStore.put(PlaywrgihtRegistry.class, instance);
       }
-    } finally {
-      beforeLock.unlock();
+      return instance;
     }
-  }
 
-  // This is a workaround for JUnit's lack of an "AfterTestRun" hook
-  // This will be called once after all tests have completed.
-  @Override
-  public void close() throws Throwable {
-    for (Playwright playwright : playwrightList) {
-      playwright.close();
+    Playwright createPlaywright(Playwright.CreateOptions options) {
+      Playwright playwright = Playwright.create(options);
+      playwrightList.add(playwright);
+      return playwright;
+    }
+
+
+    // This is a workaround for JUnit's lack of an "AfterTestRun" hook
+    // This will be called once after all tests have completed.
+    @Override
+    public void close() throws Throwable {
+      for (Playwright playwright : playwrightList) {
+        playwright.close();
+      }
+      playwrightList.clear();
     }
-    playwrightList.clear();
   }
 
   @Override
@@ -70,9 +70,10 @@ static Playwright getOrCreatePlaywright(ExtensionContext extensionContext) {
     }
 
     Options options = OptionsExtension.getOptions(extensionContext);
-    playwright = Playwright.create(options.playwrightCreateOptions);
+    PlaywrgihtRegistry registry = PlaywrgihtRegistry.getOrCreateFor(extensionContext);
+    playwright = registry.createPlaywright(options.playwrightCreateOptions);
     threadLocalPlaywright.set(playwright);
-    playwrightList.add(playwright);
+
     setTestIdAttribute(playwright, options);
     return playwright;
   }