From 5376d5be3e43f281c3700d716ef4a8846fef357e Mon Sep 17 00:00:00 2001
From: Hiroshi Ogawa <hi.ogawa.zz@gmail.com>
Date: Tue, 20 Feb 2024 22:05:16 +0900
Subject: [PATCH] fix(ui): auto reload coverage iframe after test run (#5242)

---
 packages/ui/client/composables/client/index.ts |  6 ++++++
 packages/vitest/src/api/setup.ts               | 10 ++++++++--
 packages/vitest/src/api/types.ts               |  1 +
 packages/vitest/src/node/core.ts               |  9 ++++++++-
 packages/ws-client/src/index.ts                |  3 +++
 5 files changed, 26 insertions(+), 3 deletions(-)

diff --git a/packages/ui/client/composables/client/index.ts b/packages/ui/client/composables/client/index.ts
index 2f0cfedaf79a..c484711bc532 100644
--- a/packages/ui/client/composables/client/index.ts
+++ b/packages/ui/client/composables/client/index.ts
@@ -29,6 +29,12 @@ export const client = (function createVitestClient() {
           testRunState.value = 'idle'
           unhandledErrors.value = (errors || []).map(parseError)
         },
+        onFinishedReportCoverage() {
+          // reload coverage iframe
+          const iframe = document.querySelector('iframe#vitest-ui-coverage')
+          if (iframe instanceof HTMLIFrameElement && iframe.contentWindow)
+            iframe.contentWindow.location.reload()
+        },
       },
     })
   }
diff --git a/packages/vitest/src/api/setup.ts b/packages/vitest/src/api/setup.ts
index a56dd53b15fc..fd5e15df4e39 100644
--- a/packages/vitest/src/api/setup.ts
+++ b/packages/vitest/src/api/setup.ts
@@ -162,7 +162,7 @@ export function setup(vitestOrWorkspace: Vitest | WorkspaceProject, _server?: Vi
       {
         post: msg => ws.send(msg),
         on: fn => ws.on('message', fn),
-        eventNames: ['onUserConsoleLog', 'onFinished', 'onCollected', 'onCancel'],
+        eventNames: ['onUserConsoleLog', 'onFinished', 'onFinishedReportCoverage', 'onCollected', 'onCancel'],
         serialize: (data: any) => stringify(data, stringifyReplace),
         deserialize: parse,
         onTimeoutError(functionName) {
@@ -183,7 +183,7 @@ export function setup(vitestOrWorkspace: Vitest | WorkspaceProject, _server?: Vi
   ctx.reporters.push(new WebSocketReporter(ctx, wss, clients))
 }
 
-class WebSocketReporter implements Reporter {
+export class WebSocketReporter implements Reporter {
   constructor(
     public ctx: Vitest,
     public wss: WebSocketServer,
@@ -226,6 +226,12 @@ class WebSocketReporter implements Reporter {
     })
   }
 
+  onFinishedReportCoverage() {
+    this.clients.forEach((client) => {
+      client.onFinishedReportCoverage?.()
+    })
+  }
+
   onUserConsoleLog(log: UserConsoleLog) {
     this.clients.forEach((client) => {
       client.onUserConsoleLog?.(log)
diff --git a/packages/vitest/src/api/types.ts b/packages/vitest/src/api/types.ts
index 333182a93ce5..b59c6dd151ff 100644
--- a/packages/vitest/src/api/types.ts
+++ b/packages/vitest/src/api/types.ts
@@ -39,4 +39,5 @@ export interface WebSocketHandlers {
 
 export interface WebSocketEvents extends Pick<Reporter, 'onCollected' | 'onFinished' | 'onTaskUpdate' | 'onUserConsoleLog' | 'onPathsCollected'> {
   onCancel: (reason: CancelReason) => void
+  onFinishedReportCoverage: () => void
 }
diff --git a/packages/vitest/src/node/core.ts b/packages/vitest/src/node/core.ts
index 8bd9f601bb7e..aa2bdfb458f9 100644
--- a/packages/vitest/src/node/core.ts
+++ b/packages/vitest/src/node/core.ts
@@ -16,6 +16,7 @@ import { getCoverageProvider } from '../integrations/coverage'
 import type { BrowserProvider } from '../types/browser'
 import { CONFIG_NAMES, configFiles, workspacesFiles as workspaceFiles } from '../constants'
 import { rootDir } from '../paths'
+import { WebSocketReporter } from '../api/setup'
 import { createPool } from './pool'
 import type { ProcessPool, WorkspaceSpec } from './pool'
 import { createBenchmarkReporters, createReporters } from './reporters/utils'
@@ -800,8 +801,14 @@ export class Vitest {
     if (!this.config.coverage.reportOnFailure && this.state.getCountOfFailedTests() > 0)
       return
 
-    if (this.coverageProvider)
+    if (this.coverageProvider) {
       await this.coverageProvider.reportCoverage({ allTestsRun })
+      // notify coverage iframe reload
+      for (const reporter of this.reporters) {
+        if (reporter instanceof WebSocketReporter)
+          reporter.onFinishedReportCoverage()
+      }
+    }
   }
 
   async close() {
diff --git a/packages/ws-client/src/index.ts b/packages/ws-client/src/index.ts
index 0118c87552b9..3069899bb392 100644
--- a/packages/ws-client/src/index.ts
+++ b/packages/ws-client/src/index.ts
@@ -70,6 +70,9 @@ export function createClient(url: string, options: VitestClientOptions = {}) {
     onFinished(files, errors) {
       handlers.onFinished?.(files, errors)
     },
+    onFinishedReportCoverage() {
+      handlers.onFinishedReportCoverage?.()
+    },
     onCancel(reason: CancelReason) {
       handlers.onCancel?.(reason)
     },