unify logic

Author: mxschmitt

playwright/src/test/java/com/microsoft/playwright/Server.java

@@ -20,11 +20,10 @@
 
 import java.io.*;
 import java.net.InetSocketAddress;
-import java.nio.file.Files;
-import java.nio.file.Path;
 import java.util.*;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.Future;
+import java.util.function.Function;
 import java.util.zip.GZIPOutputStream;
 
 import static com.microsoft.playwright.Utils.copy;
@@ -42,7 +41,7 @@ public class Server implements HttpHandler {
   private final Map<String, String> csp = Collections.synchronizedMap(new HashMap<>());
   private final Map<String, HttpHandler> routes = Collections.synchronizedMap(new HashMap<>());
   private final Set<String> gzipRoutes = Collections.synchronizedSet(new HashSet<>());
-  private Path staticFilesDirectory;
+  private Function<String, InputStream> resourceProvider;
 
   private static class Auth {
     public final String user;
@@ -78,6 +77,7 @@ private Server(int port, boolean https) throws IOException {
     server.createContext("/", this);
     server.setExecutor(null); // creates a default executor
     server.start();
+    resourceProvider = path -> Server.class.getClassLoader().getResourceAsStream("resources" + path);
   }
 
   public void stop() {
@@ -96,8 +96,8 @@ void enableGzip(String path) {
     gzipRoutes.add(path);
   }
 
-  void setStaticFilesDirectory(Path staticFilesDirectory) {
-    this.staticFilesDirectory = staticFilesDirectory;
+  void setResourceProvider(Function<String, InputStream> resourceProvider) {
+    this.resourceProvider = resourceProvider;
   }
 
   static class Request {
@@ -194,30 +194,16 @@ public void handle(HttpExchange exchange) throws IOException {
       path = "/index.html";
     }
 
-    // If static files directory is set, serve from filesystem first
-    if (staticFilesDirectory != null) {
-      Path filePath = staticFilesDirectory.resolve(path.substring(1)); // Remove leading /
-      if (Files.exists(filePath) && !Files.isDirectory(filePath)) {
-        exchange.getResponseHeaders().add("Content-Type", mimeType(filePath.toFile()));
-        exchange.sendResponseHeaders(200, Files.size(filePath));
-        Files.copy(filePath, exchange.getResponseBody());
-        exchange.getResponseBody().close();
-        return;
-      }
-    }
-
-    // Fallback: Resources from "src/test/resources/" are copied to "resources/" directory in the jar.
-    String resourcePath = "resources" + path;
-    InputStream resource = getClass().getClassLoader().getResourceAsStream(resourcePath);
+    InputStream resource = resourceProvider.apply(path);
     if (resource == null) {
       exchange.getResponseHeaders().add("Content-Type", "text/plain");
       exchange.sendResponseHeaders(404, 0);
       try (Writer writer = new OutputStreamWriter(exchange.getResponseBody())) {
-        writer.write("File not found: " + resourcePath);
+        writer.write("File not found: " + path);
       }
       return;
     }
-    exchange.getResponseHeaders().add("Content-Type", mimeType(new File(resourcePath)));
+    exchange.getResponseHeaders().add("Content-Type", mimeType(new File(path)));
     ByteArrayOutputStream body = new ByteArrayOutputStream();
     OutputStream output = body;
     if (gzipRoutes.contains(path)) {

playwright/src/test/java/com/microsoft/playwright/TestTracing.java

@@ -368,7 +368,17 @@ private void showTraceViewer(Path tracePath, TraceViewerConsumer callback) throw
     Path driverDir = Driver.ensureDriverInstalled(java.util.Collections.emptyMap(), true).driverDir();
     Path traceViewerPath = driverDir.resolve("package").resolve("lib").resolve("vite").resolve("traceViewer");
     Server traceServer = Server.createHttp(Utils.nextFreePort());
-    traceServer.setStaticFilesDirectory(traceViewerPath);
+    traceServer.setResourceProvider(path -> {
+      Path filePath = traceViewerPath.resolve(path.substring(1));
+      if (Files.exists(filePath) && !Files.isDirectory(filePath)) {
+        try {
+          return Files.newInputStream(filePath);
+        } catch (IOException e) {
+          return null;
+        }
+      }
+      return null;
+    });
     traceServer.setRoute("/trace.zip", exchange -> {
       exchange.getResponseHeaders().add("Content-Type", "application/zip");
       exchange.sendResponseHeaders(200, Files.size(tracePath));