Test rerun feature enhancement

Author: vishnuv688

packages/app/src/components/browser/snapshot.ts

@@ -9,11 +9,10 @@ import type { CommandLog } from '@wdio/devtools-service/types'
 
 import {
   mutationContext,
-  type TraceMutation,
   metadataContext,
-  type Metadata,
   commandContext
-} from '../../controller/DataManager.js'
+} from '../../controller/context.js'
+import type { Metadata } from '@wdio/devtools-service/types'
 
 import '~icons/mdi/world.js'
 import '../placeholder.js'

packages/app/src/components/sidebar/explorer.ts

@@ -7,9 +7,8 @@ import type { Metadata } from '@wdio/devtools-service/types'
 import { repeat } from 'lit/directives/repeat.js'
 import {
   suiteContext,
-  metadataContext,
-  isTestRunningContext
-} from '../../controller/DataManager.js'
+  metadataContext
+} from '../../controller/context.js'
 import type {
   TestEntry,
   RunCapabilities,
@@ -73,9 +72,6 @@ export class DevtoolsSidebarExplorer extends CollapseableEntry {
   @consume({ context: metadataContext, subscribe: true })
   metadata: Metadata | undefined = undefined
 
-  @consume({ context: isTestRunningContext, subscribe: true })
-  isTestRunning = false
-
   updated(changedProperties: Map<string | number | symbol, unknown>) {
     super.updated(changedProperties)
   }
@@ -491,10 +487,6 @@ export class DevtoolsSidebarExplorer extends CollapseableEntry {
             )
           : html`<div class="text-sm px-4 py-2">
               <p class="text-disabledForeground">No tests to display</p>
-              <p class="text-xs text-disabledForeground mt-2">
-                Debug: suites=${this.suites?.length || 0},
-                rootSuites=${uniqueSuites.length}, filtered=${suites.length}
-              </p>
             </div>`}
       </wdio-test-suite>
     `

packages/app/src/components/workbench.ts

@@ -7,7 +7,7 @@ import { DragController, Direction } from '../utils/DragController.js'
 import {
   consoleLogContext,
   networkRequestContext
-} from '../controller/DataManager.js'
+} from '../controller/context.js'
 
 import '~icons/mdi/arrow-collapse-down.js'
 import '~icons/mdi/arrow-collapse-up.js'

packages/app/src/components/workbench/actions.ts

@@ -3,12 +3,11 @@ import { html, css } from 'lit'
 import { customElement } from 'lit/decorators.js'
 import { consume } from '@lit/context'
 
+import type { CommandLog } from '@wdio/devtools-service/types'
 import {
   mutationContext,
-  type TraceMutation,
-  commandContext,
-  type CommandLog
-} from '../../controller/DataManager.js'
+  commandContext
+} from '../../controller/context.js'
 
 import '~icons/mdi/pencil.js'
 import '~icons/mdi/family-tree.js'

packages/app/src/components/workbench/console.ts

@@ -3,7 +3,7 @@ import { html, css, nothing } from 'lit'
 import { customElement } from 'lit/decorators.js'
 import { consume } from '@lit/context'
 
-import { consoleLogContext } from '../../controller/DataManager.js'
+import { consoleLogContext } from '../../controller/context.js'
 
 const LOG_ICONS: Record<ConsoleLogs['type'], string> = {
   log: '📄',

packages/app/src/components/workbench/metadata.ts

@@ -3,7 +3,8 @@ import { html, css } from 'lit'
 import { customElement } from 'lit/decorators.js'
 import { consume } from '@lit/context'
 
-import { metadataContext, type Metadata } from '../../controller/DataManager.js'
+import type { Metadata } from '@wdio/devtools-service/types'
+import { metadataContext } from '../../controller/context.js'
 
 import './list.js'
 import '../placeholder.js'

packages/app/src/components/workbench/network.ts

@@ -2,7 +2,7 @@ import { Element } from '@core/element'
 import { html, css, nothing } from 'lit'
 import { customElement, state } from 'lit/decorators.js'
 import { consume } from '@lit/context'
-import { networkRequestContext } from '../../controller/DataManager.js'
+import { networkRequestContext } from '../../controller/context.js'
 import { RESOURCE_TYPES } from '../../utils/network-constants.js'
 import {
   formatBytes,

packages/app/src/components/workbench/source.ts

@@ -8,7 +8,7 @@ import type { EditorViewConfig } from '@codemirror/view'
 import { javascript } from '@codemirror/lang-javascript'
 import { oneDark } from '@codemirror/theme-one-dark'
 
-import { sourceContext } from '../../controller/DataManager.js'
+import { sourceContext } from '../../controller/context.js'
 
 import '../placeholder.js'
 

packages/app/src/controller/DataManager.ts

@@ -1,73 +1,22 @@
 import { createContext, ContextProvider } from '@lit/context'
 import type { ReactiveController, ReactiveControllerHost } from 'lit'
-import type {
-  Metadata,
-  CommandLog,
-  TraceLog
-} from '@wdio/devtools-service/types'
-import type { SuiteStats, TestStats } from '@wdio/reporter'
+import type { Metadata, CommandLog, TraceLog } from '@wdio/devtools-service/types'
+
+import {
+  mutationContext,
+  logContext,
+  consoleLogContext,
+  networkRequestContext,
+  metadataContext,
+  commandContext,
+  sourceContext,
+  suiteContext
+} from './context.js'
+import type { TestStatsFragment, SuiteStatsFragment, SocketMessage } from './types.js'
 
 const CACHE_ID = 'wdio-trace-cache'
 
-type TestStatsFragment = Omit<Partial<TestStats>, 'uid'> & { uid: string }
-
-type SuiteStatsFragment = Omit<
-  Partial<SuiteStats>,
-  'uid' | 'tests' | 'suites'
-> & {
-  uid: string
-  tests?: TestStatsFragment[]
-  suites?: SuiteStatsFragment[]
-}
-
-export const mutationContext = createContext<TraceMutation[]>(
-  Symbol('mutationContext')
-)
-export const logContext = createContext<string[]>(Symbol('logContext'))
-export const consoleLogContext = createContext<ConsoleLogs[]>(
-  Symbol('consoleLogContext')
-)
-export const networkRequestContext = createContext<NetworkRequest[]>(
-  Symbol('networkRequestContext')
-)
-export const metadataContext = createContext<Metadata>(
-  Symbol('metadataContext')
-)
-export const commandContext = createContext<CommandLog[]>(
-  Symbol('commandContext')
-)
-export const sourceContext = createContext<Record<string, string>>(
-  Symbol('sourceContext')
-)
-export const suiteContext = createContext<Record<string, any>[]>(
-  Symbol('suiteContext')
-)
-
 const hasConnection = createContext<boolean>(Symbol('hasConnection'))
-export const isTestRunningContext = createContext<boolean>(
-  Symbol('isTestRunning')
-)
-
-interface SocketMessage<
-  T extends
-    | keyof TraceLog
-    | 'testStopped'
-    | 'clearExecutionData'
-    | 'replaceCommand' =
-    | keyof TraceLog
-    | 'testStopped'
-    | 'clearExecutionData'
-    | 'replaceCommand'
-> {
-  scope: T
-  data: T extends keyof TraceLog
-    ? TraceLog[T]
-    : T extends 'clearExecutionData'
-      ? { uid?: string }
-      : T extends 'replaceCommand'
-        ? { oldTimestamp: number; command: CommandLog }
-        : unknown
-}
 
 export class DataManagerController implements ReactiveController {
   #ws?: WebSocket
@@ -129,15 +78,13 @@ export class DataManagerController implements ReactiveController {
     return this.metadataContextProvider.value?.type
   }
 
-  // Public method to clear execution data when rerun is triggered
   clearExecutionData(uid?: string) {
     this.#resetExecutionData()
     if (uid) {
       this.#markTestAsRunning(uid)
     }
   }
 
-  // Private method to mark a test/suite as running immediately for UI feedback
   #markTestAsRunning(uid: string) {
     const suites = this.suitesContextProvider.value || []
 
@@ -191,12 +138,11 @@ export class DataManagerController implements ReactiveController {
 
           // Recursive helper to mark tests/suites as running
           const markAsRunning = (s: SuiteStatsFragment): SuiteStatsFragment => {
-            // If this is the target suite/test, mark it as running
             if (s.uid === uid) {
               return {
                 ...s,
                 start: new Date(),
-                end: undefined, // Clear end to mark as running
+                end: undefined,
                 tests:
                   (s.tests?.map((test) => ({
                     ...test,
@@ -207,7 +153,6 @@ export class DataManagerController implements ReactiveController {
               }
             }
 
-            // Check if any child test matches
             const updatedTests = (s.tests?.map((test) => {
               if (test.uid === uid) {
                 return {
@@ -219,7 +164,6 @@ export class DataManagerController implements ReactiveController {
               return test
             }) ?? []) as TestStatsFragment[]
 
-            // Recursively check nested suites
             const updatedNestedSuites = s.suites?.map(markAsRunning)
 
             return {
@@ -277,22 +221,19 @@ export class DataManagerController implements ReactiveController {
         return
       }
 
-      // Handle test stopped event
       if (scope === 'testStopped') {
         this.#handleTestStopped()
         this.#host.requestUpdate()
         return
       }
 
-      // Handle clear execution data event (when tests change)
       if (scope === 'clearExecutionData') {
         const clearData = data as { uid?: string }
         this.clearExecutionData(clearData.uid)
         this.#host.requestUpdate()
         return
       }
 
-      // Handle in-place command replacement (retry deduplication)
       if (scope === 'replaceCommand') {
         const { oldTimestamp, command } = data as {
           oldTimestamp: number
@@ -303,17 +244,17 @@ export class DataManagerController implements ReactiveController {
         return
       }
 
-      // Check for new run BEFORE processing suites data
       if (scope === 'suites') {
         const shouldReset = this.#shouldResetForNewRun(data)
         if (shouldReset) {
           this.#resetExecutionData()
         }
       }
 
-      // Route data to appropriate handler
       if (scope === 'mutations') {
         this.#handleMutationsUpdate(data as TraceMutation[])
+      } else if (scope === 'logs') {
+        this.#handleLogsUpdate(data as string[])
       } else if (scope === 'commands') {
         this.#handleCommandsUpdate(data as CommandLog[])
       } else if (scope === 'metadata') {
@@ -326,8 +267,6 @@ export class DataManagerController implements ReactiveController {
         this.#handleSourcesUpdate(data as Record<string, string>)
       } else if (scope === 'suites') {
         this.#handleSuitesUpdate(data)
-      } else {
-        this.#handleGenericUpdate(scope, data)
       }
 
       this.#host.requestUpdate()
@@ -376,7 +315,7 @@ export class DataManagerController implements ReactiveController {
               ec as Record<string, SuiteStatsFragment>
             )) {
               if (uid === Object.keys(chunk)[0]) {
-                existingEnd = (existing as any)?.end
+                existingEnd = existing?.end
                 break outer
               }
             }
@@ -432,7 +371,7 @@ export class DataManagerController implements ReactiveController {
                 return {
                   ...test,
                   end: new Date(),
-                  state: 'failed' as 'failed',
+                  state: 'failed',
                   error: {
                     message: 'Test execution stopped',
                     name: 'TestStoppedError'
@@ -445,8 +384,24 @@ export class DataManagerController implements ReactiveController {
             // Recursively update nested suites (for Cucumber scenarios)
             const updatedNestedSuites = s.suites?.map(updateSuite)
 
+            // Derive the suite's own state from its updated children so that
+            // STATE_MAP['running'] no longer produces a spinner after stop.
+            const allTests = [...(updatedTests || []), ...(updatedNestedSuites || [])]
+            const hasFailed = allTests.some((t) => t?.state === 'failed')
+            const hasRunning = allTests.some((t) => !t?.end)
+            const derivedState: SuiteStatsFragment['state'] = hasRunning
+              ? s.state
+              : hasFailed
+                ? 'failed'
+                : s.state === 'running'
+                  ? 'failed'
+                  : s.state
+
             return {
               ...s,
+              state: derivedState,
+              ...(!hasRunning && !s.end ? { end: new Date() } : {}),
+
               tests: updatedTests || [],
               suites: updatedNestedSuites || []
             }
@@ -566,27 +521,11 @@ export class DataManagerController implements ReactiveController {
     return date instanceof Date ? date.getTime() : date
   }
 
-  #handleGenericUpdate(scope: keyof TraceLog, data: any) {
-    const providerMap = {
-      mutations: this.mutationsContextProvider,
-      logs: this.logsContextProvider,
-      consoleLogs: this.consoleLogsContextProvider,
-      metadata: this.metadataContextProvider,
-      commands: this.commandsContextProvider,
-      sources: this.sourcesContextProvider,
-      suites: this.suitesContextProvider
-    } as const
-
-    const provider = providerMap[scope as keyof typeof providerMap]
-    if (provider) {
-      provider.setValue(data)
-    }
+  #handleLogsUpdate(data: string[]) {
+    this.logsContextProvider.setValue(data)
   }
 
   #mergeSuite(existing: SuiteStatsFragment, incoming: SuiteStatsFragment) {
-    // Note: Rerun detection and clearing is now handled in #handleSuitesUpdate
-    // before any merges happen, so data is cleared proactively
-
     // First merge tests and suites properly
     const mergedTests = this.#mergeTests(existing.tests, incoming.tests)
     const mergedSuites = this.#mergeChildSuites(
@@ -676,7 +615,4 @@ export class DataManagerController implements ReactiveController {
   }
 }
 
-/**
- * re-export types used for context
- */
-export { type Metadata, type CommandLog, type TraceLog, type TraceMutation }
+

packages/app/src/controller/context.ts

@@ -0,0 +1,11 @@
+import { createContext } from '@lit/context'
+import type { Metadata, CommandLog } from '@wdio/devtools-service/types'
+
+export const mutationContext = createContext<TraceMutation[]>(Symbol('mutationContext'))
+export const logContext = createContext<string[]>(Symbol('logContext'))
+export const consoleLogContext = createContext<ConsoleLogs[]>(Symbol('consoleLogContext'))
+export const networkRequestContext = createContext<NetworkRequest[]>(Symbol('networkRequestContext'))
+export const metadataContext = createContext<Metadata>(Symbol('metadataContext'))
+export const commandContext = createContext<CommandLog[]>(Symbol('commandContext'))
+export const sourceContext = createContext<Record<string, string>>(Symbol('sourceContext'))
+export const suiteContext = createContext<Record<string, any>[]>(Symbol('suiteContext'))