From 840e02f13e007b740079ef84b6bdf030624a0751 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Joaqu=C3=ADn=20S=C3=A1nchez?= <userquin@gmail.com>
Date: Fri, 12 Jul 2024 17:25:44 +0200
Subject: [PATCH] feat(ui): show all suites/tests when parent matches (#6106)

---
 .../ui/client/composables/explorer/filter.ts  | 63 ++++++++++++++++---
 .../ui/client/composables/explorer/types.ts   |  3 +
 .../ui/client/composables/explorer/utils.ts   |  6 +-
 3 files changed, 63 insertions(+), 9 deletions(-)

diff --git a/packages/ui/client/composables/explorer/filter.ts b/packages/ui/client/composables/explorer/filter.ts
index 7cb25523acb7..3894b522e8ca 100644
--- a/packages/ui/client/composables/explorer/filter.ts
+++ b/packages/ui/client/composables/explorer/filter.ts
@@ -48,14 +48,44 @@ export function* filterNode(
 ) {
   const treeNodes = new Set<string>()
 
+  const parentsMap = new Map<string, boolean>()
   const list: FilterResult[] = []
 
-  for (const entry of visitNode(
-    node,
-    treeNodes,
-    n => matcher(n, search, filter),
-  )) {
-    list.push(entry)
+  let fileId: string | undefined
+
+  if (filter.onlyTests) {
+    for (const [match, child] of visitNode(
+      node,
+      treeNodes,
+      n => matcher(n, search, filter),
+    )) {
+      list.push([match, child])
+    }
+  }
+  else {
+    for (const [match, child] of visitNode(
+      node,
+      treeNodes,
+      n => matcher(n, search, filter),
+    )) {
+      if (isParentNode(child)) {
+        parentsMap.set(child.id, match)
+        if (isFileNode(child)) {
+          match && (fileId = child.id)
+          list.push([match, child])
+        }
+        else {
+          list.push([match || parentsMap.get(child.parentId) === true, child])
+        }
+      }
+      else {
+        list.push([match || parentsMap.get(child.parentId) === true, child])
+      }
+    }
+    // when expanding a non-file node
+    if (!fileId && !isFileNode(node) && 'fileId' in node) {
+      fileId = node.fileId as string
+    }
   }
 
   const filesToShow = new Set<string>()
@@ -65,6 +95,7 @@ export function* filterNode(
     filter.onlyTests,
     treeNodes,
     filesToShow,
+    fileId,
   )].reverse()
 
   // We show only the files and parents whose parent is expanded.
@@ -129,10 +160,28 @@ function* filterParents(
   collapseParents: boolean,
   treeNodes: Set<string>,
   filesToShow: Set<string>,
+  nodeId?: string,
 ) {
   for (let i = list.length - 1; i >= 0; i--) {
     const [match, child] = list[i]
-    if (isParentNode(child)) {
+    const isParent = isParentNode(child)
+    if (!collapseParents && nodeId && treeNodes.has(nodeId) && 'fileId' in child && child.fileId === nodeId) {
+      if (isParent) {
+        treeNodes.add(child.id)
+      }
+      let parent = explorerTree.nodes.get(child.parentId)
+      while (parent) {
+        treeNodes.add(parent.id)
+        if (isFileNode(parent)) {
+          filesToShow.add(parent.id)
+        }
+        parent = explorerTree.nodes.get(parent.parentId)
+      }
+      yield child
+      continue
+    }
+
+    if (isParent) {
       const node = expandCollapseNode(
         match,
         child,
diff --git a/packages/ui/client/composables/explorer/types.ts b/packages/ui/client/composables/explorer/types.ts
index 163452b767c5..744dbce87b64 100644
--- a/packages/ui/client/composables/explorer/types.ts
+++ b/packages/ui/client/composables/explorer/types.ts
@@ -38,10 +38,12 @@ export interface UITaskTreeNode extends TaskTreeNode {
 }
 
 export interface TestTreeNode extends UITaskTreeNode {
+  fileId: string
   type: 'test'
 }
 
 export interface CustomTestTreeNode extends UITaskTreeNode {
+  fileId: string
   type: 'custom'
 }
 
@@ -51,6 +53,7 @@ export interface ParentTreeNode extends UITaskTreeNode {
 }
 
 export interface SuiteTreeNode extends ParentTreeNode {
+  fileId: string
   type: 'suite'
   typecheck?: boolean
 }
diff --git a/packages/ui/client/composables/explorer/utils.ts b/packages/ui/client/composables/explorer/utils.ts
index 776b4806adc2..9efd1193b61c 100644
--- a/packages/ui/client/composables/explorer/utils.ts
+++ b/packages/ui/client/composables/explorer/utils.ts
@@ -115,7 +115,7 @@ export function createOrUpdateNodeTask(id: string) {
   }
 
   const task = client.state.idMap.get(id)
-  // if no children just return
+  // if it is not a test just return
   if (!task || !isAtomTest(task)) {
     return
   }
@@ -149,6 +149,7 @@ export function createOrUpdateNode(
       if (isAtomTest(task)) {
         taskNode = {
           id: task.id,
+          fileId: task.file.id,
           parentId,
           name: task.name,
           mode: task.mode,
@@ -158,11 +159,12 @@ export function createOrUpdateNode(
           indent: node.indent + 1,
           duration: task.result?.duration,
           state: task.result?.state,
-        }
+        } as TestTreeNode | CustomTestTreeNode
       }
       else {
         taskNode = {
           id: task.id,
+          fileId: task.file.id,
           parentId,
           name: task.name,
           mode: task.mode,