From b03f0249c0f71424385fadc92d3a9a2307c04c11 Mon Sep 17 00:00:00 2001
From: Pavel Feldman <pavel.feldman@gmail.com>
Date: Thu, 22 Feb 2024 14:33:17 -0800
Subject: [PATCH] chore: use test server as a singleton

---
 src/playwrightTest.ts       |  8 ++---
 src/testServerController.ts | 72 ++++++++++++++++++++++++++++++-------
 2 files changed, 64 insertions(+), 16 deletions(-)

diff --git a/src/playwrightTest.ts b/src/playwrightTest.ts
index f1a5a72d1..8388a9fad 100644
--- a/src/playwrightTest.ts
+++ b/src/playwrightTest.ts
@@ -258,11 +258,11 @@ export class PlaywrightTest {
     const env = await reporterServer.env({ selfDestruct: false });
     const reporter = reporterServer.reporterFile();
     if (mode === 'list')
-      testServer.list({ locations, reporter, env });
+      testServer.list({ configFile: config.configFile, locations, reporter, env });
     if (mode === 'run') {
-      testServer.test({ locations, reporter, env, options });
+      testServer.test({ configFile: config.configFile, locations, reporter, env, ...options });
       token.onCancellationRequested(() => {
-        testServer.stop();
+        testServer.stop({ configFile: config.configFile });
       });
       testServer.on('stdio', params => {
         if (params.type === 'stdout')
@@ -309,7 +309,7 @@ export class PlaywrightTest {
     const testServer = await this._testServerController.testServerFor(config);
     if (!testServer)
       return await this._findRelatedTestFilesCLI(config, files);
-    return await testServer.findRelatedTestFiles({ files });
+    return await testServer.findRelatedTestFiles({ configFile: config.configFile, files });
   }
 
   async debugTests(vscode: vscodeTypes.VSCode, config: TestConfig, projectNames: string[], testDirs: string[], settingsEnv: NodeJS.ProcessEnv, locations: string[] | null, listener: TestListener, parametrizedTestTitle: string | undefined, token: vscodeTypes.CancellationToken) {
diff --git a/src/testServerController.ts b/src/testServerController.ts
index 17aeb7efa..b4f32f403 100644
--- a/src/testServerController.ts
+++ b/src/testServerController.ts
@@ -17,11 +17,13 @@
 import { BackendClient, BackendServer } from './backend';
 import { ConfigFindRelatedTestFilesReport } from './listTests';
 import { TestConfig } from './playwrightTest';
+import type { TestError } from './reporter';
 import * as vscodeTypes from './vscodeTypes';
 
 export class TestServerController implements vscodeTypes.Disposable {
   private _vscode: vscodeTypes.VSCode;
-  private _testServers = new Map<TestConfig, TestServer>();
+  private _instancePromise: Promise<TestServer | null> | undefined;
+  private _instance: TestServer | null = null;
   private _envProvider: () => NodeJS.ProcessEnv;
 
   constructor(vscode: vscodeTypes.VSCode, envProvider: () => NodeJS.ProcessEnv) {
@@ -29,10 +31,14 @@ export class TestServerController implements vscodeTypes.Disposable {
     this._envProvider = envProvider;
   }
 
-  async testServerFor(config: TestConfig): Promise<TestServer | null> {
-    const existing = this._testServers.get(config);
-    if (existing)
-      return existing;
+  async testServerFor(config: TestConfig): Promise<TestServerInterface & TestServerEvents | null> {
+    if (this._instancePromise)
+      return this._instancePromise;
+    this._instancePromise = this._createTestServer(config);
+    return this._instancePromise;
+  }
+
+  private async _createTestServer(config: TestConfig): Promise<TestServer | null> {
     const args = [config.cli, 'test-server'];
     const testServerBackend = new BackendServer<TestServer>(this._vscode, {
       args,
@@ -47,9 +53,7 @@ export class TestServerController implements vscodeTypes.Disposable {
       dumpIO: false,
     });
     const testServer = await testServerBackend.start();
-    if (!testServer)
-      return null;
-    this._testServers.set(config, testServer);
+    this._instance = testServer;
     return testServer;
   }
 
@@ -58,13 +62,52 @@ export class TestServerController implements vscodeTypes.Disposable {
   }
 
   reset() {
-    for (const backend of this._testServers.values())
-      backend.close();
-    this._testServers.clear();
+    if (this._instancePromise)
+      this._instancePromise.then(server => server?.closeGracefully());
+    this._instancePromise = undefined;
+    this._instance = null;
   }
 }
 
-class TestServer extends BackendClient {
+interface TestServerInterface {
+  list(params: {
+    configFile: string;
+    locations: string[];
+    reporter: string;
+    env: NodeJS.ProcessEnv;
+  }): Promise<void>;
+
+  test(params: {
+    configFile: string;
+    locations: string[];
+    reporter: string;
+    env: NodeJS.ProcessEnv;
+    headed?: boolean;
+    oneWorker?: boolean;
+    trace?: 'on' | 'off';
+    projects?: string[];
+    grep?: string;
+    reuseContext?: boolean;
+    connectWsEndpoint?: string;
+  }): Promise<void>;
+
+  findRelatedTestFiles(params: {
+    configFile: string;
+    files: string[];
+  }): Promise<{ testFiles: string[]; errors?: TestError[]; }>;
+
+  stop(params: {
+    configFile: string;
+  }): Promise<void>;
+
+  closeGracefully(): Promise<void>;
+}
+
+interface TestServerEvents {
+  on(event: 'stdio', listener: (params: { type: 'stdout' | 'stderr', text?: string, buffer?: string }) => void): void;
+}
+
+class TestServer extends BackendClient implements TestServerInterface, TestServerEvents {
   override async initialize(): Promise<void> {
   }
 
@@ -83,4 +126,9 @@ class TestServer extends BackendClient {
   async stop() {
     await this.send('stop', {});
   }
+
+  async closeGracefully() {
+    await this.send('closeGracefully', {});
+    this.close();
+  }
 }