From b24fb2535991df1f1053973e48c5a7e46b543135 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Joaqu=C3=ADn=20S=C3=A1nchez?= <joaquin.sanchez@fi2net.es>
Date: Mon, 4 Apr 2022 21:57:00 +0200
Subject: [PATCH] fix(logs): collect logs per test (#1088)

---
 packages/ui/client/auto-imports.d.ts          |  1 +
 packages/ui/client/components/FileDetails.vue |  7 +-
 packages/vitest/src/runtime/setup.ts          | 83 +++++++++++++------
 packages/vitest/src/types/general.ts          |  1 +
 4 files changed, 64 insertions(+), 28 deletions(-)

diff --git a/packages/ui/client/auto-imports.d.ts b/packages/ui/client/auto-imports.d.ts
index b3553d813dfc..6e5fda32acc3 100644
--- a/packages/ui/client/auto-imports.d.ts
+++ b/packages/ui/client/auto-imports.d.ts
@@ -86,6 +86,7 @@ declare global {
   const toRef: typeof import('vue')['toRef']
   const toRefs: typeof import('vue')['toRefs']
   const triggerRef: typeof import('vue')['triggerRef']
+  const tryOnBeforeMount: typeof import('@vueuse/core')['tryOnBeforeMount']
   const tryOnBeforeUnmount: typeof import('@vueuse/core')['tryOnBeforeUnmount']
   const tryOnMounted: typeof import('@vueuse/core')['tryOnMounted']
   const tryOnScopeDispose: typeof import('@vueuse/core')['tryOnScopeDispose']
diff --git a/packages/ui/client/components/FileDetails.vue b/packages/ui/client/components/FileDetails.vue
index beb8c11c082c..615e6d8273b0 100644
--- a/packages/ui/client/components/FileDetails.vue
+++ b/packages/ui/client/components/FileDetails.vue
@@ -29,6 +29,9 @@ const open = () => {
 const changeViewMode = (view: Params['view']) => {
   viewMode.value = view
 }
+const consoleCount = computed(() => {
+  return currentLogs.value?.reduce((s, { size }) => s + size, 0) ?? 0
+})
 </script>
 
 <template>
@@ -72,10 +75,10 @@ const changeViewMode = (view: Params['view']) => {
         </button>
         <button
           tab-button
-          :class="{ 'tab-button-active': viewMode === 'console', 'op20': viewMode !== 'console' && currentLogs?.length === 0 }"
+          :class="{ 'tab-button-active': viewMode === 'console', 'op20': viewMode !== 'console' && consoleCount === 0 }"
           @click="changeViewMode('console')"
         >
-          Console ({{ currentLogs?.length || 0 }})
+          Console ({{ consoleCount }})
         </button>
       </div>
     </div>
diff --git a/packages/vitest/src/runtime/setup.ts b/packages/vitest/src/runtime/setup.ts
index 8532e9bdc156..8de40ca3201a 100644
--- a/packages/vitest/src/runtime/setup.ts
+++ b/packages/vitest/src/runtime/setup.ts
@@ -38,16 +38,17 @@ function setupDefines(defines: Record<string, any>) {
 }
 
 export function setupConsoleLogSpy() {
-  const stdoutBuffer: any[] = []
-  const stderrBuffer: any[] = []
-  let stdoutTime = 0
-  let stderrTime = 0
-  let timer: any
+  const stdoutBuffer = new Map<string, any[]>()
+  const stderrBuffer = new Map<string, any[]>()
+  const timers = new Map<string, { stdoutTime: number; stderrTime: number; timer: any }>()
+  const unknownTestId = '__vitest__unknown_test__'
 
   // group sync console.log calls with macro task
-  function schedule(taskId?: string) {
-    clearTimeout(timer)
-    timer = setTimeout(() => {
+  function schedule(taskId: string) {
+    const timer = timers.get(taskId)!
+    const { stdoutTime, stderrTime } = timer
+    clearTimeout(timer.timer)
+    timer.timer = setTimeout(() => {
       if (stderrTime < stdoutTime) {
         sendStderr(taskId)
         sendStdout(taskId)
@@ -58,45 +59,75 @@ export function setupConsoleLogSpy() {
       }
     })
   }
-  function sendStdout(taskId?: string) {
-    if (stdoutBuffer.length) {
+  function sendStdout(taskId: string) {
+    const buffer = stdoutBuffer.get(taskId)
+    if (buffer) {
+      const timer = timers.get(taskId)!
       rpc().onUserConsoleLog({
         type: 'stdout',
-        content: stdoutBuffer.map(i => String(i)).join(''),
+        content: buffer.map(i => String(i)).join(''),
         taskId,
-        time: stdoutTime || RealDate.now(),
+        time: timer.stdoutTime || RealDate.now(),
+        size: buffer.length,
       })
+      stdoutBuffer.set(taskId, [])
+      timer.stdoutTime = 0
     }
-    stdoutBuffer.length = 0
-    stdoutTime = 0
   }
-  function sendStderr(taskId?: string) {
-    if (stderrBuffer.length) {
+  function sendStderr(taskId: string) {
+    const buffer = stderrBuffer.get(taskId)
+    if (buffer) {
+      const timer = timers.get(taskId)!
       rpc().onUserConsoleLog({
         type: 'stderr',
-        content: stderrBuffer.map(i => String(i)).join(''),
+        content: buffer.map(i => String(i)).join(''),
         taskId,
-        time: stderrTime || RealDate.now(),
+        time: timer.stderrTime || RealDate.now(),
+        size: buffer.length,
       })
+      stderrBuffer.set(taskId, [])
+      timer.stderrTime = 0
     }
-    stderrBuffer.length = 0
-    stderrTime = 0
   }
 
   const stdout = new Writable({
     write(data, encoding, callback) {
-      const id = getWorkerState()?.current?.id
-      stdoutTime = stdoutTime || RealDate.now()
-      stdoutBuffer.push(data)
+      const id = getWorkerState()?.current?.id ?? unknownTestId
+      let timer = timers.get(id)
+      if (timer) {
+        timer.stdoutTime = timer.stdoutTime || RealDate.now()
+      }
+      else {
+        timer = { stdoutTime: RealDate.now(), stderrTime: RealDate.now(), timer: 0 }
+        timers.set(id, timer)
+      }
+      let buffer = stdoutBuffer.get(id)
+      if (!buffer) {
+        buffer = []
+        stdoutBuffer.set(id, buffer)
+      }
+      buffer.push(data)
       schedule(id)
       callback()
     },
   })
   const stderr = new Writable({
     write(data, encoding, callback) {
-      const id = getWorkerState()?.current?.id
-      stderrTime = stderrTime || RealDate.now()
-      stderrBuffer.push(data)
+      const id = getWorkerState()?.current?.id ?? unknownTestId
+      let timer = timers.get(id)
+      if (timer) {
+        timer.stderrTime = timer.stderrTime || RealDate.now()
+      }
+      else {
+        timer = { stderrTime: RealDate.now(), stdoutTime: RealDate.now(), timer: 0 }
+        timers.set(id, timer)
+      }
+      let buffer = stderrBuffer.get(id)
+      if (!buffer) {
+        buffer = []
+        stderrBuffer.set(id, buffer)
+      }
+      buffer.push(data)
       schedule(id)
       callback()
     },
diff --git a/packages/vitest/src/types/general.ts b/packages/vitest/src/types/general.ts
index 08f65bf7481c..7005ed271b98 100644
--- a/packages/vitest/src/types/general.ts
+++ b/packages/vitest/src/types/general.ts
@@ -44,6 +44,7 @@ export interface UserConsoleLog {
   type: 'stdout' | 'stderr'
   taskId?: string
   time: number
+  size: number
 }
 
 export interface Position {