refactor: change the way to manage metadata

Changed to using WeakMap to manage custom data according to VSCode documentation.

Author: mato533

src/api/run.ts

@@ -1,11 +1,11 @@
 import { getGrep, getRange, getCucumberSpec, getSpec } from './utils.js'
-import { isConfig, isSpec, isTestcase, isWdioTestItem } from '../test/index.js'
+import { RepositoryManager } from '../test/index.js'
 import { log } from '../utils/logger.js'
 
 import type * as vscode from 'vscode'
 import type { RunTestOptions } from './types.js'
 import type { WdioExtensionWorkerInterface } from './types.js'
-import type { SpecFileTestItem, WdioConfigTestItem, TestcaseTestItem } from '../test/index.js'
+import type { TestItemMetadata, TestItemMetadataWithRepository } from '../test/index.js'
 
 type Listeners = {
     stdout: (data: string) => void
@@ -31,11 +31,12 @@ export class TestRunner {
      * Run a test based on the provided TestItem
      */
     public async run(test: vscode.TestItem) {
+        const metadata = RepositoryManager.getMetadata(test)
         // Validate test item
-        this.validateTestItem(test)
+        this.validateTestItem(metadata)
 
         // Create test execution options
-        const testOptions = this.createTestOptions(test)
+        const testOptions = this.createTestOptions(test, metadata)
 
         try {
             // Execute test and process results
@@ -53,35 +54,29 @@ export class TestRunner {
     /**
      * Validates that the provided TestItem is a valid WebdriverIO test
      */
-    private validateTestItem(
-        test: vscode.TestItem
-    ): asserts test is TestcaseTestItem | WdioConfigTestItem | SpecFileTestItem {
-        if (!isWdioTestItem(test)) {
-            throw new Error("The metadata for TestItem is not set. This is extension's bug.")
-        }
-
-        if (!isConfig(test) && !isSpec(test) && !isTestcase(test)) {
+    private validateTestItem(metadata: TestItemMetadata): asserts metadata is TestItemMetadataWithRepository {
+        if ('repository' in metadata && typeof metadata.repository !== 'object') {
             throw new Error('Workspace TestItem is not valid.')
         }
     }
 
     /**
      * Creates RunTestOptions based on the test type and framework
      */
-    private createTestOptions(test: TestcaseTestItem | WdioConfigTestItem | SpecFileTestItem): RunTestOptions {
-        const isCucumberFramework = this.isCucumberFramework(test)
+    private createTestOptions(test: vscode.TestItem, metadata: TestItemMetadataWithRepository): RunTestOptions {
+        const isCucumberFramework = this.isCucumberFramework(metadata)
 
         // Get appropriate specs based on the test framework and type
-        const specs = this.determineSpecs(test, isCucumberFramework)
+        const specs = this.determineSpecs(test, isCucumberFramework, metadata)
 
         // Get grep pattern for mocha-like frameworks when testing individual test cases
-        const grep = !isCucumberFramework && isTestcase(test) ? getGrep(test) : undefined
+        const grep = !isCucumberFramework && metadata.isTestcase ? getGrep(test) : undefined
 
         // Get line range information for test file
-        const range = !isCucumberFramework && isTestcase(test) ? getRange(test) : undefined
+        const range = !isCucumberFramework && metadata.isTestcase ? getRange(test) : undefined
 
         return {
-            configPath: test.metadata.repository.wdioConfigPath,
+            configPath: metadata.repository.wdioConfigPath,
             specs,
             grep,
             range,
@@ -91,26 +86,23 @@ export class TestRunner {
     /**
      * Determines if the test is using Cucumber framework
      */
-    private isCucumberFramework(test: TestcaseTestItem | WdioConfigTestItem | SpecFileTestItem): boolean {
-        return test.metadata.repository.framework === 'cucumber'
+    private isCucumberFramework(metadata: TestItemMetadataWithRepository): boolean {
+        return metadata.repository.framework === 'cucumber'
     }
 
     /**
      * Determines the specs to run based on test type and framework
      */
     private determineSpecs(
-        test: TestcaseTestItem | WdioConfigTestItem | SpecFileTestItem,
-        isCucumberFramework: boolean
+        test: vscode.TestItem,
+        isCucumberFramework: boolean,
+        metadata: TestItemMetadataWithRepository
     ): string[] | undefined {
-        if (isCucumberFramework && isTestcase(test)) {
-            return getCucumberSpec(test)
-        }
-
-        if (isSpec(test) || isTestcase(test)) {
-            return getSpec(test)
+        if (isCucumberFramework) {
+            return getCucumberSpec(test, metadata)
         }
 
-        return undefined
+        return getSpec(test)
     }
 
     /**

src/api/utils.ts

@@ -1,11 +1,11 @@
 import { WebSocketServer } from 'ws'
 
 import { LOG_LEVEL, TEST_ID_SEPARATOR } from '../constants.js'
+import { RepositoryManager, type TestItemMetadataWithRepository } from '../test/index.js'
 import { log } from '../utils/logger.js'
 
 import type { Server } from 'node:http'
 import type * as vscode from 'vscode'
-import type { TestcaseTestItem } from '../test/index.js'
 import type { NumericLogLevel } from '../types.js'
 
 export async function loggingFn(_logLevel: NumericLogLevel, message: string) {
@@ -52,15 +52,16 @@ export function getRange(test: vscode.TestItem) {
     return isEmptyRange ? undefined : test.range
 }
 
-export function getCucumberSpec(testItem: TestcaseTestItem) {
+export function getCucumberSpec(testItem: vscode.TestItem, metadata: TestItemMetadataWithRepository) {
     const baseSpec = getSpec(testItem)
     if (!baseSpec) {
         return undefined
     }
-    if (testItem.metadata.type === 'rule') {
+    if (metadata.type === 'rule') {
         const specs = []
         for (const [_, childItem] of testItem.children) {
-            if ((childItem as TestcaseTestItem).metadata.type === 'scenario') {
+            const childeMetadata = RepositoryManager.getMetadata(childItem)
+            if (childeMetadata.type === 'scenario') {
                 const start = childItem.range?.start.line || 0
                 const end = childItem.range?.end.line || 0
                 if (start > 0 && end > 0) {
@@ -75,7 +76,7 @@ export function getCucumberSpec(testItem: TestcaseTestItem) {
         }
     }
 
-    if (testItem.metadata.type === 'scenario') {
+    if (metadata.type === 'scenario') {
         const specs = []
         const start = testItem.range?.start.line || 0
         const end = testItem.range?.end.line || 0

src/test/index.ts

@@ -1,5 +1,4 @@
 export { RepositoryManager } from './manager.js'
-export { isWdioTestItem, isConfig, isSpec, isTestcase, isWorkspace } from './utils.js'
 export { TestReporter } from './reporter.js'
 export { TestfileWatcher } from './watcher.js'
 export { convertUriToPath } from './converter.js'

src/test/manager.ts

@@ -8,10 +8,33 @@ import { createRunProfile } from './utils.js'
 import { TEST_ID_SEPARATOR } from '../constants.js'
 import { log } from '../utils/logger.js'
 
-import type { WdioConfigTestItem, WorkspaceTestItem } from './types.js'
+import type { TestItemMetadata } from './types.js'
 import type { ServerManager } from '../api/manager.js'
 import type { ExtensionConfigManager } from '../config/index.js'
 
+class MetadataRepository {
+    private static testMetadataRepository = new WeakMap<vscode.TestItem, TestItemMetadata>()
+    public static getMetadata(testItem: vscode.TestItem) {
+        const metadata = this.testMetadataRepository.get(testItem)
+        if (!metadata) {
+            throw new Error("The metadata for TestItem is not set. This is extension's bug.")
+        }
+        return metadata
+    }
+
+    public static getRepository(testItem: vscode.TestItem) {
+        const metadata = this.getMetadata(testItem)
+        if (!metadata.repository) {
+            throw new Error("The metadata for TestItem is not set. This is extension's bug.")
+        }
+        return metadata.repository
+    }
+
+    public static setMetadata(testItem: vscode.TestItem, metadata: TestItemMetadata) {
+        this.testMetadataRepository.set(testItem, metadata)
+    }
+}
+
 /**
  * workspace                                        -- managed by this class
  *   - configuration file (e.g. wdio.conf.js)       -- managed by this class
@@ -22,18 +45,19 @@ import type { ExtensionConfigManager } from '../config/index.js'
 
 const LOADING_TEST_ITEM_ID = '_resolving'
 
-export class RepositoryManager implements vscode.Disposable {
+export class RepositoryManager extends MetadataRepository implements vscode.Disposable {
     private readonly _repos = new Set<TestRepository>()
     private _loadingTestItem: vscode.TestItem
-    private _workspaceTestItems: WorkspaceTestItem[] = []
-    private _wdioConfigTestItems: WdioConfigTestItem[] = []
+    private _workspaceTestItems: vscode.TestItem[] = []
+    private _wdioConfigTestItems: vscode.TestItem[] = []
     private isCreatedDefaultProfile = false
 
     constructor(
         public readonly controller: vscode.TestController,
         public readonly configManager: ExtensionConfigManager,
         private readonly serverManager: ServerManager
     ) {
+        super()
         this._loadingTestItem = this.controller.createTestItem(LOADING_TEST_ITEM_ID, 'Resolving WebdriverIO Tests...')
         this._loadingTestItem.sortText = '.0' // show at first line
         this._loadingTestItem.busy = true
@@ -82,7 +106,8 @@ export class RepositoryManager implements vscode.Disposable {
         for (const workspaceTestItem of affectedWorkspaceItems) {
             const configTestItem = await this.createWdioConfigTestItem(workspaceTestItem, wdioConfigPath)
 
-            await configTestItem.metadata.repository.discoverAllTests()
+            const repo = RepositoryManager.getRepository(configTestItem)
+            await repo.discoverAllTests()
             if (!this.configManager.isMultiWorkspace) {
                 this.controller.items.add(configTestItem)
             }
@@ -96,14 +121,12 @@ export class RepositoryManager implements vscode.Disposable {
         log.debug(`Remove the config file from ${affectedWorkspaceItems.length} workspace(s)`)
         const configUri = convertPathToUri(wdioConfigPath)
         for (const workspaceItem of affectedWorkspaceItems) {
-            const config = workspaceItem.children.get(
-                this.generateConfigTestItemId(workspaceItem, configUri)
-            ) as WdioConfigTestItem
+            const config = workspaceItem.children.get(this.generateConfigTestItemId(workspaceItem, configUri))
             if (!config) {
                 continue
             }
             log.debug(`Remove the TestItem: ${config.id}`)
-            const targetRepo = config.metadata.repository
+            const targetRepo = RepositoryManager.getRepository(config)
             targetRepo.dispose()
             this._repos.delete(targetRepo)
 
@@ -140,22 +163,24 @@ export class RepositoryManager implements vscode.Disposable {
             `workspace:${workspaceFolder.uri.fsPath}`,
             workspaceFolder.name,
             workspaceFolder.uri
-        ) as WorkspaceTestItem
-        workspaceItem['metadata'] = {
+        )
+        RepositoryManager.setMetadata(workspaceItem, {
+            uri: workspaceFolder.uri,
             isWorkspace: true,
             isConfigFile: false,
             isSpecFile: false,
-        }
+            isTestcase: false,
+        })
         return workspaceItem
     }
 
-    private async createWdioConfigTestItem(workspaceTestItem: WorkspaceTestItem, wdioConfigPath: string) {
-        const configUri = convertPathToUri(wdioConfigPath)
+    private async createWdioConfigTestItem(workspaceTestItem: vscode.TestItem, wdioConfigPath: string) {
+        const uri = convertPathToUri(wdioConfigPath)
         const configItem = this.controller.createTestItem(
-            this.generateConfigTestItemId(workspaceTestItem, configUri),
+            this.generateConfigTestItemId(workspaceTestItem, uri),
             basename(wdioConfigPath),
-            configUri
-        ) as WdioConfigTestItem
+            uri
+        )
 
         workspaceTestItem.children.add(configItem)
         this._wdioConfigTestItems.push(configItem)
@@ -165,17 +190,20 @@ export class RepositoryManager implements vscode.Disposable {
         this._repos.add(repo)
 
         configItem.description = relative(workspaceTestItem.uri!.fsPath, dirname(wdioConfigPath))
-        configItem['metadata'] = {
+
+        RepositoryManager.setMetadata(configItem, {
+            uri,
             isWorkspace: false,
             isConfigFile: true,
             isSpecFile: false,
+            isTestcase: false,
             repository: repo,
             runProfiles: createRunProfile.call(this, configItem, !this.isCreatedDefaultProfile),
-        }
+        })
         return configItem
     }
 
-    private generateConfigTestItemId(workspaceTestItem: WorkspaceTestItem, configUri: vscode.Uri) {
+    private generateConfigTestItemId(workspaceTestItem: vscode.TestItem, configUri: vscode.Uri) {
         return [workspaceTestItem.id, `config:${configUri.fsPath}`].join(TEST_ID_SEPARATOR)
     }
 

src/test/reporter.ts

@@ -76,7 +76,13 @@ export class TestReporter {
      */
     private processHierarchicalSuite(suite: TestSuite, parentItem: vscode.TestItem): void {
         // Find the suite TestItem in the repository
-        const suiteItem = this._repository.searchSuite(suite.name, parentItem)
+        let suiteItem: vscode.TestItem | undefined
+
+        parentItem.children.forEach((child) => {
+            if (child.label === suite.name) {
+                suiteItem = child
+            }
+        })
 
         if (!suiteItem) {
             log.debug(`Suite TestItem not found for suite: ${suite.name}`)