From a89e61b149a96907a0f2f96c295c0a5624e2216f Mon Sep 17 00:00:00 2001
From: Zixuan Chen <remch183@outlook.com>
Date: Wed, 18 May 2022 01:03:13 +0800
Subject: [PATCH] fix: dispose state after turning off watch mode

---
 src/pure/watch/ws-client.ts |  5 +++
 src/watch.ts                | 87 +++++++++++++++++--------------------
 2 files changed, 46 insertions(+), 46 deletions(-)

diff --git a/src/pure/watch/ws-client.ts b/src/pure/watch/ws-client.ts
index fd12ef23..0d7edb2f 100644
--- a/src/pure/watch/ws-client.ts
+++ b/src/pure/watch/ws-client.ts
@@ -100,6 +100,7 @@ export interface VitestClient {
   rpc: BirpcReturn<WebSocketHandlers>
   waitForConnection(): Promise<void>
   reconnect(): Promise<void>
+  dispose(): void
 }
 
 export function createClient(url: string, options: VitestClientOptions = {}) {
@@ -118,6 +119,10 @@ export function createClient(url: string, options: VitestClientOptions = {}) {
     state: new StateManager(),
     waitForConnection,
     reconnect: () => reconnect(true),
+    dispose: () => {
+      tries = 0
+      ctx.ws.close()
+    },
   }) as VitestClient
 
   ctx.state.filesMap = reactive(ctx.state.filesMap)
diff --git a/src/watch.ts b/src/watch.ts
index 9fb84efe..e8ec1874 100644
--- a/src/watch.ts
+++ b/src/watch.ts
@@ -57,8 +57,6 @@ export class TestWatcher extends Disposable {
   ) {
     super(() => {
       this.dispose()
-      this.vitestState?.client.ws.close()
-      this.vitestState = undefined
     })
   }
 
@@ -99,52 +97,47 @@ export class TestWatcher extends Disposable {
       console.log('VITEST WATCH PROCESS EXIT')
     })
 
-    if (this.vitestState) {
-      this.vitestState.client.reconnect()
-    }
-    else {
-      this.vitestState = buildWatchClient({
-        handlers: {
-          onTaskUpdate: (packs) => {
-            try {
-              if (!this.vitestState)
-                return
-
-              this.isRunning.value = true
-              const idMap = this.vitestState.client.state.idMap
-              const fileSet = new Set<File>()
-              for (const [id] of packs) {
-                const task = idMap.get(id)
-                if (!task)
-                  continue
-
-                task.file && fileSet.add(task.file)
-              }
-
-              this.onUpdated(Array.from(fileSet), false)
-            }
-            catch (e) {
-              console.error(e)
-            }
-          },
-          onFinished: (files) => {
-            try {
-              this.isRunning.value = false
-              this.onUpdated(files, true)
-              if (!this.run)
-                return
-
-              this.run.end()
-              this.run = undefined
-              this.updateStatus()
+    this.vitestState = buildWatchClient({
+      handlers: {
+        onTaskUpdate: (packs) => {
+          try {
+            if (!this.vitestState)
+              return
+
+            this.isRunning.value = true
+            const idMap = this.vitestState.client.state.idMap
+            const fileSet = new Set<File>()
+            for (const [id] of packs) {
+              const task = idMap.get(id)
+              if (!task)
+                continue
+
+              task.file && fileSet.add(task.file)
             }
-            catch (e) {
-              console.error(e)
-            }
-          },
+
+            this.onUpdated(Array.from(fileSet), false)
+          }
+          catch (e) {
+            console.error(e)
+          }
         },
-      })
-    }
+        onFinished: (files) => {
+          try {
+            this.isRunning.value = false
+            this.onUpdated(files, true)
+            if (!this.run)
+              return
+
+            this.run.end()
+            this.run = undefined
+            this.updateStatus()
+          }
+          catch (e) {
+            console.error(e)
+          }
+        },
+      },
+    })
 
     effect(() => {
       this.onFileUpdated(this.vitestState!.files.value)
@@ -361,7 +354,9 @@ export class TestWatcher extends Disposable {
   public dispose() {
     console.log('Stop watch mode')
     this.isWatching.value = false
+    this.vitestState?.client.dispose()
     this.process?.kill()
     this.process = undefined
+    this.vitestState = undefined
   }
 }