From 2e59055b6fbb1b5e12586a065716e564ed5451dc Mon Sep 17 00:00:00 2001
From: userquin <userquin@gmail.com>
Date: Tue, 9 Jul 2024 20:35:57 +0200
Subject: [PATCH] fix(ui): test files notified only when running

---
 .../ui/client/composables/client/index.ts     |  4 +--
 .../client/composables/explorer/collector.ts  | 32 +++++++++++++++++--
 .../ui/client/composables/explorer/filter.ts  |  3 +-
 .../ui/client/composables/explorer/utils.ts   |  6 ++++
 4 files changed, 39 insertions(+), 6 deletions(-)

diff --git a/packages/ui/client/composables/client/index.ts b/packages/ui/client/composables/client/index.ts
index 1ca85f8a618c..56ed67f154b5 100644
--- a/packages/ui/client/composables/client/index.ts
+++ b/packages/ui/client/composables/client/index.ts
@@ -9,7 +9,6 @@ import { parseError } from '../error'
 import { activeFileId } from '../params'
 import { createStaticClient } from './static'
 import { testRunState, unhandledErrors } from './state'
-import { uiFiles } from '~/composables/explorer/state'
 import { explorerTree } from '~/composables/explorer'
 import { isFileNode } from '~/composables/explorer/utils'
 
@@ -51,8 +50,7 @@ export const status = ref<WebSocketStatus>('CONNECTING')
 
 export const current = computed(() => {
   const currentFileId = activeFileId.value
-  const entry = uiFiles.value.find(file => file.id === currentFileId)!
-  return entry ? findById(entry.id) : undefined
+  return currentFileId ? findById(currentFileId) : undefined
 })
 export const currentLogs = computed(() => getTasks(current.value).map(i => i?.logs || []).flat() || [])
 
diff --git a/packages/ui/client/composables/explorer/collector.ts b/packages/ui/client/composables/explorer/collector.ts
index a5adca3196d8..3d8e433af94f 100644
--- a/packages/ui/client/composables/explorer/collector.ts
+++ b/packages/ui/client/composables/explorer/collector.ts
@@ -122,6 +122,16 @@ function updateRunningTodoTests() {
 }
 
 function traverseFiles(collect: boolean) {
+  // add missing files: now we have only files with running tests on the initial ws open event
+  const files = client.state.getFiles()
+  const currentFiles = explorerTree.nodes
+  const missingFiles = files.filter(f => !currentFiles.has(f.id))
+  for (let i = 0; i < missingFiles.length; i++) {
+    createOrUpdateFileNode(missingFiles[i], collect)
+    createOrUpdateEntry(missingFiles[i].tasks)
+  }
+
+  // update pending tasks
   const rootTasks = explorerTree.root.tasks
   // collect remote children
   for (let i = 0; i < rootTasks.length; i++) {
@@ -142,11 +152,29 @@ function traverseFiles(collect: boolean) {
 }
 
 function traverseReceivedFiles(collect: boolean) {
-  const rootTasks = explorerTree.root.tasks
   const updatedFiles = new Map(explorerTree.pendingTasks.entries())
   explorerTree.pendingTasks.clear()
-  const idMap = client.state.idMap
+
+  // add missing files: now we have only files with running tests on the initial ws open event
+  const currentFiles = explorerTree.nodes
+  const missingFiles = Array
+    .from(updatedFiles.keys())
+    .filter(id => !currentFiles.has(id))
+    .map(id => findById(id))
+    .filter(Boolean) as File[]
+
+  let newFile: File
+  for (let i = 0; i < missingFiles.length; i++) {
+    newFile = missingFiles[i]
+    createOrUpdateFileNode(newFile, false)
+    createOrUpdateEntry(newFile.tasks)
+    // remove the file from the updated files
+    updatedFiles.delete(newFile.id)
+  }
+
   // collect remote children
+  const idMap = client.state.idMap
+  const rootTasks = explorerTree.root.tasks
   for (let i = 0; i < rootTasks.length; i++) {
     const fileNode = rootTasks[i]
     const file = findById(fileNode.id)
diff --git a/packages/ui/client/composables/explorer/filter.ts b/packages/ui/client/composables/explorer/filter.ts
index 1af7a0886921..7cb25523acb7 100644
--- a/packages/ui/client/composables/explorer/filter.ts
+++ b/packages/ui/client/composables/explorer/filter.ts
@@ -5,6 +5,7 @@ import {
   isFileNode,
   isParentNode,
   isTestNode,
+  sortedRootTasks,
 } from '~/composables/explorer/utils'
 import { client, findById } from '~/composables/client'
 import { filteredFiles, uiEntries } from '~/composables/explorer/state'
@@ -35,7 +36,7 @@ export function* filterAll(
   search: string,
   filter: Filter,
 ) {
-  for (const node of explorerTree.root.tasks) {
+  for (const node of sortedRootTasks()) {
     yield * filterNode(node, search, filter)
   }
 }
diff --git a/packages/ui/client/composables/explorer/utils.ts b/packages/ui/client/composables/explorer/utils.ts
index 3e524b68c034..776b4806adc2 100644
--- a/packages/ui/client/composables/explorer/utils.ts
+++ b/packages/ui/client/composables/explorer/utils.ts
@@ -33,6 +33,12 @@ export function isParentNode(node: UITaskTreeNode): node is FileTreeNode | Suite
   return node.type === 'file' || node.type === 'suite'
 }
 
+export function sortedRootTasks(tasks = explorerTree.root.tasks) {
+  return tasks.sort((a, b) => {
+    return `${a.filepath}:${a.projectName}`.localeCompare(`${b.filepath}:${b.projectName}`)
+  })
+}
+
 export function createOrUpdateFileNode(
   file: File,
   collect = false,