Restore previously loaded test cases (#52)

This change will persist the prior list of test cases as part of the
workspace state, and restore this data when the test explorer
relaunches. This avoids a wait for the full BSP build target request
before test cases become visible:
- Update the test case store to cache the build targets result and each
target sources result. When a refresh runs, the latest version of these
get stored.
- Before registering the handlers, there is now an additional step to
pull these stored test cases and re-add them back to the test explorer.
- These get processed through the usual test case creation logic, so
that a test explorer tree created from a restored result or a new one
from the build server will behave the same way.
- The refresh run always clears out any cached values, so these are only
used when restoring the window, and an actual refresh run is always
guaranteed to contain the latest results from the build server.

Author: mnoah1

src/test-explorer/resolver.ts

@@ -39,10 +39,86 @@ export class TestResolver implements OnModuleInit, vscode.Disposable {
 
   onModuleInit() {
     this.ctx.subscriptions.push(this)
+    this.restorePriorTests()
+      .catch(e => {
+        this.outputChannel.appendLine(`Error restoring prior tests: ${e}`)
+      })
+      .finally(() => {
+        this.registerHandlers()
+      })
+  }
+
+  registerHandlers() {
     this.store.testController.resolveHandler = this.resolveHandler.bind(this)
     this.store.testController.refreshHandler = this.refreshHandler.bind(this)
   }
 
+  private async restorePriorTests() {
+    await this.resolveRoot()
+    const cachedBuildTargetsResult = this.store.getCachedBuildTargetsResult()
+
+    // Restore the root test item and build targets.
+    let addedTargets = false
+    if (cachedBuildTargetsResult) {
+      this.outputChannel.appendLine('Restoring prior test cases')
+      const root = this.store.testController.items.get('root')
+      if (root) {
+        try {
+          this.outputChannel.appendLine('Restoring Source Files for Root')
+          await this.processWorkspaceBuildTargetsResult(
+            root,
+            cachedBuildTargetsResult
+          )
+          root.canResolveChildren = false
+          addedTargets = true
+        } catch (e) {
+          this.outputChannel.appendLine(
+            `Error restoring source files for root: ${e}`
+          )
+        }
+      }
+    }
+
+    // Restore cached source files for targets that have them.
+    const promises: Promise<void>[] = []
+    for (const [key, testItem] of this.store.targetIdentifiers.entries()) {
+      let cachedSourceFiles: bsp.SourcesResult | undefined
+      try {
+        const target: bsp.BuildTargetIdentifier = JSON.parse(key)
+        const sourceParams: bsp.SourcesParams = {
+          targets: [target],
+        }
+        cachedSourceFiles = this.store.getCachedSourcesResult(sourceParams)
+      } catch (e) {
+        this.outputChannel.appendLine(
+          `Invalid key '${key}' when restoring source files: ${e}`
+        )
+      }
+
+      if (cachedSourceFiles) {
+        this.outputChannel.appendLine(
+          `Restoring source files for target: ${key}`
+        )
+        try {
+          testItem.canResolveChildren = false
+          const promise = this.processTargetSourcesResult(
+            testItem,
+            cachedSourceFiles
+          )
+          promises.push(promise)
+        } catch (e) {
+          this.outputChannel.appendLine(
+            `Error restoring source files for target: ${key}`
+          )
+        }
+      }
+    }
+    await Promise.all(promises)
+
+    // Kick off discovery of test cases in other open files that may not have been discovered yet.
+    if (addedTargets) await this.resolveOpenSourceFiles()
+  }
+
   dispose() {}
 
   private async resolveHandler(
@@ -74,7 +150,6 @@ export class TestResolver implements OnModuleInit, vscode.Disposable {
         if (parentTest === undefined) {
           try {
             await this.buildServer.getConnection()
-            this.resolveRoot()
           } catch (e) {
             this.outputChannel.appendLine(
               'Test explorer disabled due to lack of available build server.'
@@ -102,6 +177,7 @@ export class TestResolver implements OnModuleInit, vscode.Disposable {
                 'Fetching test targets from build server ([progress](command:bazelbsp.showServerOutput))',
             })
             await this.resolveTargets(parentTest, combinedToken)
+            await this.resolveOpenSourceFiles()
             break
           case TestItemType.BazelTarget:
             progress.report({
@@ -124,6 +200,7 @@ export class TestResolver implements OnModuleInit, vscode.Disposable {
    * @param token Cancellation token tied to the refresh button on the VS Code UI.
    */
   private async refreshHandler(token: vscode.CancellationToken) {
+    this.store.clearCache()
     const promises: Promise<void>[] = []
     this.store.testController.items.forEach(async item => {
       promises.push(this.resolveHandler(item, token))
@@ -173,6 +250,7 @@ export class TestResolver implements OnModuleInit, vscode.Disposable {
         bsp.WorkspaceBuildTargets.type,
         cancellationToken
       )
+      this.store.cacheBuildTargetsResult(result)
     } catch (e) {
       if (e.code === CANCEL_ERROR_CODE) {
         updateDescription(
@@ -184,7 +262,13 @@ export class TestResolver implements OnModuleInit, vscode.Disposable {
       updateDescription(parentTest, 'Error: unable to fetch targets')
       throw e
     }
+    await this.processWorkspaceBuildTargetsResult(parentTest, result)
+  }
 
+  private async processWorkspaceBuildTargetsResult(
+    parentTest: vscode.TestItem,
+    result: bsp.WorkspaceBuildTargetsResult
+  ) {
     updateDescription(parentTest, 'Loading: processing target results')
 
     // Process the returned targets, create new test items, and store their metadata.
@@ -243,8 +327,6 @@ export class TestResolver implements OnModuleInit, vscode.Disposable {
     // Replace all children with the newly returned test cases.
     this.condenseTestItems(parentTest)
     updateDescription(parentTest)
-
-    await this.resolveOpenSourceFiles()
   }
 
   /**
@@ -357,6 +439,17 @@ export class TestResolver implements OnModuleInit, vscode.Disposable {
       params,
       cancellationToken
     )
+    this.store.cacheSourcesResult(params, result)
+    await this.processTargetSourcesResult(parentTest, result)
+  }
+
+  private async processTargetSourcesResult(
+    parentTest: vscode.TestItem,
+    result: bsp.SourcesResult,
+    cancellationToken?: vscode.CancellationToken
+  ) {
+    const parentTarget = this.store.testCaseMetadata.get(parentTest)?.target
+    if (!parentTarget) return
 
     const directories = new Map<string, vscode.TestItem>()
     parentTest.children.replace([])

src/test-explorer/store.ts

@@ -1,4 +1,5 @@
 import * as vscode from 'vscode'
+import * as bsp from '../bsp/bsp'
 import {Inject, OnModuleInit} from '@nestjs/common'
 
 import {
@@ -8,6 +9,9 @@ import {
 import {TestCaseInfo} from '../test-info/test-info'
 import {BuildTargetIdentifier} from 'src/bsp/bsp'
 
+const CONTEXT_KEY_TARGETS_RESULT = 'testExplorerBuildTargetsResult'
+const CONTEXT_KEY_SOURCES_RESULT = 'testExplorerSourcesResult'
+
 export class TestCaseStore implements OnModuleInit, vscode.Disposable {
   @Inject(EXTENSION_CONTEXT_TOKEN) private readonly ctx: vscode.ExtensionContext
   @Inject(TEST_CONTROLLER_TOKEN) readonly testController: vscode.TestController
@@ -17,7 +21,7 @@ export class TestCaseStore implements OnModuleInit, vscode.Disposable {
   // Watcher to update a test item's children.  Key corresponds to the test item ID.
   testItemWatchers: Map<string, vscode.FileSystemWatcher>
   knownFiles: Set<string>
-  private targetIdentifiers: Map<string, vscode.TestItem>
+  targetIdentifiers: Map<string, vscode.TestItem>
 
   constructor() {
     this.testCaseMetadata = new WeakMap<vscode.TestItem, TestCaseInfo>()
@@ -83,4 +87,41 @@ export class TestCaseStore implements OnModuleInit, vscode.Disposable {
   clearTargetIdentifiers() {
     this.targetIdentifiers.clear()
   }
+
+  cacheBuildTargetsResult(result: bsp.WorkspaceBuildTargetsResult) {
+    this.ctx.workspaceState.update(CONTEXT_KEY_TARGETS_RESULT, result)
+  }
+
+  getCachedBuildTargetsResult(): bsp.WorkspaceBuildTargetsResult | undefined {
+    return this.ctx.workspaceState.get<bsp.WorkspaceBuildTargetsResult>(
+      CONTEXT_KEY_TARGETS_RESULT
+    )
+  }
+
+  cacheSourcesResult(params: bsp.SourcesParams, result: bsp.SourcesResult) {
+    const key = JSON.stringify(params)
+    const allResults =
+      this.ctx.workspaceState.get<Record<string, bsp.SourcesResult>>(
+        CONTEXT_KEY_SOURCES_RESULT
+      ) || {}
+
+    // Add or update the result for this params key
+    allResults[key] = result
+    this.ctx.workspaceState.update(CONTEXT_KEY_SOURCES_RESULT, allResults)
+  }
+
+  getCachedSourcesResult(
+    params: bsp.SourcesParams
+  ): bsp.SourcesResult | undefined {
+    const key = JSON.stringify(params)
+    const allResults = this.ctx.workspaceState.get<
+      Record<string, bsp.SourcesResult>
+    >(CONTEXT_KEY_SOURCES_RESULT)
+    return allResults?.[key]
+  }
+
+  clearCache() {
+    this.ctx.workspaceState.update(CONTEXT_KEY_TARGETS_RESULT, undefined)
+    this.ctx.workspaceState.update(CONTEXT_KEY_SOURCES_RESULT, undefined)
+  }
 }