From 1f9f06d490cc5f24b230e72932c1ab86a8ff4579 Mon Sep 17 00:00:00 2001
From: Vladimir <sleuths.slews0s@icloud.com>
Date: Wed, 9 Oct 2024 13:12:45 +0200
Subject: [PATCH] fix: don't hang when reloading during a test run (#496)

---
 src/api.ts           |  2 ++
 src/extension.ts     |  7 +++++
 src/runner/runner.ts | 65 ++++++++++++++++++++++++++++----------------
 3 files changed, 50 insertions(+), 24 deletions(-)

diff --git a/src/api.ts b/src/api.ts
index 6da26658..225ef9d6 100644
--- a/src/api.ts
+++ b/src/api.ts
@@ -177,6 +177,8 @@ export class VitestFolderAPI {
   }
 
   async cancelRun() {
+    if (this.process.closed)
+      return
     await this.meta.rpc.cancelRun()
   }
 
diff --git a/src/extension.ts b/src/extension.ts
index 0e51290c..30215cb0 100644
--- a/src/extension.ts
+++ b/src/extension.ts
@@ -35,6 +35,8 @@ class VitestExtension {
   private tagsManager: TagsManager
   private api: VitestAPI | undefined
 
+  private runners: TestRunner[] = []
+
   private disposables: vscode.Disposable[] = []
 
   constructor() {
@@ -64,6 +66,8 @@ class VitestExtension {
 
   private async _defineTestProfiles(showWarning: boolean, cancelToken?: vscode.CancellationToken) {
     this.testTree.reset([])
+    this.runners.forEach(runner => runner.dispose())
+    this.runners = []
 
     const vitest = await resolveVitestPackages(showWarning)
 
@@ -167,6 +171,7 @@ class VitestExtension {
         this.testTree,
         api,
       )
+      this.runners.push(runner)
 
       const prefix = api.prefix
       let runProfile = previousRunProfiles.get(`${api.id}:run`)
@@ -368,5 +373,7 @@ class VitestExtension {
     this.runProfiles.clear()
     this.disposables.forEach(d => d.dispose())
     this.disposables = []
+    this.runners.forEach(runner => runner.dispose())
+    this.runners = []
   }
 }
diff --git a/src/runner/runner.ts b/src/runner/runner.ts
index e4ad4c48..d06c6a34 100644
--- a/src/runner/runner.ts
+++ b/src/runner/runner.ts
@@ -22,6 +22,8 @@ export class TestRunner extends vscode.Disposable {
   private testRun: vscode.TestRun | undefined
   private testRunDefer: PromiseWithResolvers<void> | undefined
 
+  private disposables: vscode.Disposable[] = []
+
   constructor(
     private readonly controller: vscode.TestController,
     private readonly tree: TestTree,
@@ -33,10 +35,10 @@ export class TestRunner extends vscode.Disposable {
       this.endTestRun()
       this.nonContinuousRequest = undefined
       this.continuousRequests.clear()
-      if (!this.api.process.closed) {
-        this.api.cancelRun()
-      }
+      this.api.cancelRun()
       this._onRequestsExhausted.dispose()
+      this.disposables.forEach(d => d.dispose())
+      this.disposables = []
     })
 
     api.onWatcherRerun((files, _trigger, collecting) => {
@@ -155,17 +157,19 @@ export class TestRunner extends vscode.Disposable {
   private async watchContinuousTests(request: vscode.TestRunRequest, token: vscode.CancellationToken) {
     this.continuousRequests.add(request)
 
-    token.onCancellationRequested(() => {
-      log.verbose?.('Continuous test run for', join(request.include), 'was cancelled')
+    this.disposables.push(
+      token.onCancellationRequested(() => {
+        log.verbose?.('Continuous test run for', join(request.include), 'was cancelled')
 
-      this.continuousRequests.delete(request)
-      if (!this.continuousRequests.size) {
-        log.verbose?.('Stopped watching test files')
-        this._onRequestsExhausted.fire()
-        this.api.unwatchTests()
-        this.endTestRun()
-      }
-    })
+        this.continuousRequests.delete(request)
+        if (!this.continuousRequests.size) {
+          log.verbose?.('Stopped watching test files')
+          this._onRequestsExhausted.fire()
+          this.api.unwatchTests()
+          this.endTestRun()
+        }
+      }),
+    )
 
     if (!request.include?.length) {
       log.info('[RUNNER]', 'Watching all test files')
@@ -202,10 +206,12 @@ export class TestRunner extends vscode.Disposable {
       }
     })
 
-    token.onCancellationRequested(() => {
-      log.verbose?.('Coverage for', join(request.include), 'was manually stopped')
-      this.api.disableCoverage()
-    })
+    this.disposables.push(
+      token.onCancellationRequested(() => {
+        log.verbose?.('Coverage for', join(request.include), 'was manually stopped')
+        this.api.disableCoverage()
+      }),
+    )
 
     await this.runTests(request, token)
   }
@@ -218,14 +224,25 @@ export class TestRunner extends vscode.Disposable {
 
     this.nonContinuousRequest = request
 
-    token.onCancellationRequested(() => {
-      this.endTestRun()
-      this.nonContinuousRequest = undefined
-      this.api.cancelRun()
-      log.verbose?.('Test run was cancelled manually for', join(request.include))
-    })
+    this.disposables.push(
+      token.onCancellationRequested(() => {
+        this.endTestRun()
+        this.nonContinuousRequest = undefined
+        this.api.cancelRun()
+        log.verbose?.('Test run was cancelled manually for', join(request.include))
+      }),
+    )
 
-    await this.runTestItems(request)
+    try {
+      await this.runTestItems(request)
+    }
+    catch (err: any) {
+      // the rpc can be closed during the test run by clicking on reload
+      if (!err.message.startsWith('[birpc] rpc is closed')) {
+        log.error('Failed to run tests', err)
+      }
+      this.endTestRun()
+    }
 
     this.nonContinuousRequest = undefined
     this._onRequestsExhausted.fire()