Use native VS Code test UI in canary

extensions/ql-vscode/package-lock.json

@@ -93,7 +93,7 @@
         "@types/through2": "^2.0.36",
         "@types/tmp": "^0.1.0",
         "@types/unzipper": "~0.10.1",
-        "@types/vscode": "^1.59.0",
+        "@types/vscode": "^1.67.0",
         "@types/webpack": "^5.28.0",
         "@types/webpack-env": "^1.18.0",
         "@types/xml2js": "~0.4.4",
@@ -17194,9 +17194,9 @@
       }
     },
     "node_modules/@types/vscode": {
-      "version": "1.63.1",
+      "version": "1.77.0",
-      "resolved": "https://registry.npmjs.org/@types/vscode/-/vscode-1.63.1.tgz",
+      "resolved": "https://registry.npmjs.org/@types/vscode/-/vscode-1.77.0.tgz",
-      "integrity": "sha512-Z+ZqjRcnGfHP86dvx/BtSwWyZPKQ/LBdmAVImY82TphyjOw2KgTKcp7Nx92oNwCTsHzlshwexAG/WiY2JuUm3g==",
+      "integrity": "sha512-MWFN5R7a33n8eJZJmdVlifjig3LWUNRrPeO1xemIcZ0ae0TEQuRc7G2xV0LUX78RZFECY1plYBn+dP/Acc3L0Q==",
       "dev": true
     },
     "node_modules/@types/webpack": {
@@ -58329,9 +58329,9 @@
       }
     },
     "@types/vscode": {
-      "version": "1.63.1",
+      "version": "1.77.0",
-      "resolved": "https://registry.npmjs.org/@types/vscode/-/vscode-1.63.1.tgz",
+      "resolved": "https://registry.npmjs.org/@types/vscode/-/vscode-1.77.0.tgz",
-      "integrity": "sha512-Z+ZqjRcnGfHP86dvx/BtSwWyZPKQ/LBdmAVImY82TphyjOw2KgTKcp7Nx92oNwCTsHzlshwexAG/WiY2JuUm3g==",
+      "integrity": "sha512-MWFN5R7a33n8eJZJmdVlifjig3LWUNRrPeO1xemIcZ0ae0TEQuRc7G2xV0LUX78RZFECY1plYBn+dP/Acc3L0Q==",
       "dev": true
     },
     "@types/webpack": {

extensions/ql-vscode/package.json

@@ -973,6 +973,13 @@
           "when": "viewItem == testWithSource"
         }
       ],
+      "testing/item/context": [
+        {
+          "command": "codeQLTests.acceptOutput",
+          "group": "qltest@1",
+          "when": "controllerId == codeql && testId =~ /^test /"
+        }
+      ],
       "explorer/context": [
         {
           "command": "codeQL.setCurrentDatabase",
@@ -1523,7 +1530,7 @@
     "@types/through2": "^2.0.36",
     "@types/tmp": "^0.1.0",
     "@types/unzipper": "~0.10.1",
-    "@types/vscode": "^1.59.0",
+    "@types/vscode": "^1.67.0",
     "@types/webpack": "^5.28.0",
     "@types/webpack-env": "^1.18.0",
     "@types/xml2js": "~0.4.4",

extensions/ql-vscode/src/cli.ts

@@ -23,7 +23,7 @@ import {
   getErrorStack,
 } from "./pure/helpers-pure";
 import { QueryMetadata, SortDirection } from "./pure/interface-types";
-import { Logger, ProgressReporter } from "./common";
+import { BaseLogger, Logger, ProgressReporter } from "./common";
 import { CompilationMessage } from "./pure/legacy-messages";
 import { sarifParser } from "./sarif-parser";
 import { walkDirectory } from "./helpers";
@@ -134,6 +134,7 @@ export interface TestCompleted {
   compilationMs: number;
   evaluationMs: number;
   expected: string;
+  actual?: string;
   diff: string[] | undefined;
   failureDescription?: string;
   failureStage?: string;
@@ -424,7 +425,7 @@ export class CodeQLCliServer implements Disposable {
     command: string[],
     commandArgs: string[],
     cancellationToken?: CancellationToken,
-    logger?: Logger,
+    logger?: BaseLogger,
   ): AsyncGenerator<string, void, unknown> {
     // Add format argument first, in case commandArgs contains positional parameters.
     const args = [...command, "--format", "jsonz", ...commandArgs];
@@ -453,9 +454,13 @@ export class CodeQLCliServer implements Disposable {
       ])) {
         yield event;
       }
-
-      await childPromise;
     } finally {
+      try {
+        await childPromise;
+      } catch (_e) {
+        // We need to await this to avoid an unhandled rejection, but we want to propagate the
+        // original exception.
+      }
       if (cancellationRegistration !== undefined) {
         cancellationRegistration.dispose();
       }
@@ -482,7 +487,7 @@ export class CodeQLCliServer implements Disposable {
       logger,
     }: {
       cancellationToken?: CancellationToken;
-      logger?: Logger;
+      logger?: BaseLogger;
     } = {},
   ): AsyncGenerator<EventType, void, unknown> {
     for await (const event of this.runAsyncCodeQlCliCommandInternal(
@@ -761,7 +766,7 @@ export class CodeQLCliServer implements Disposable {
       logger,
     }: {
       cancellationToken?: CancellationToken;
-      logger?: Logger;
+      logger?: BaseLogger;
     },
   ): AsyncGenerator<TestCompleted, void, unknown> {
     const subcommandArgs = this.cliConfig.additionalTestArguments.concat([
@@ -1623,7 +1628,7 @@ const lineEndings = ["\r\n", "\r", "\n"];
  * @param stream The stream to log.
  * @param logger The logger that will consume the stream output.
  */
-async function logStream(stream: Readable, logger: Logger): Promise<void> {
+async function logStream(stream: Readable, logger: BaseLogger): Promise<void> {
   for await (const line of splitStreamAtSeparators(stream, lineEndings)) {
     // Await the result of log here in order to ensure the logs are written in the correct order.
     await logger.log(line);

extensions/ql-vscode/src/extension.ts

@@ -30,6 +30,7 @@ import { CodeQLCliServer } from "./cli";
 import {
   CliConfigListener,
   DistributionConfigListener,
+  isCanary,
   joinOrderWarningThreshold,
   QueryHistoryConfigListener,
   QueryServerConfigListener,
@@ -113,7 +114,6 @@ import {
   BaseCommands,
   PreActivationCommands,
   QueryServerCommands,
-  TestUICommands,
 } from "./common/commands";
 import { LocalQueries } from "./local-queries";
 import { getAstCfgCommands } from "./ast-cfg-commands";
@@ -121,6 +121,9 @@ import { getQueryEditorCommands } from "./query-editor";
 import { App } from "./common/app";
 import { registerCommandWithErrorHandling } from "./common/vscode/commands";
 import { DataExtensionsEditorModule } from "./data-extensions-editor/data-extensions-editor-module";
+import { TestManager } from "./test-manager";
+import { TestRunner } from "./test-runner";
+import { TestManagerBase } from "./test-manager-base";
 
 /**
  * extension.ts
@@ -876,25 +879,34 @@ async function activateWithInstalledDistribution(
     );
 
   void extLogger.log("Initializing QLTest interface.");
-  const testExplorerExtension = extensions.getExtension<TestHub>(
+
-    testExplorerExtensionId,
+  const testRunner = new TestRunner(dbm, cliServer);
-  );
+  ctx.subscriptions.push(testRunner);
-  let testUiCommands: Partial<TestUICommands> = {};
+
-  if (testExplorerExtension) {
+  let testManager: TestManagerBase | undefined = undefined;
-    const testHub = testExplorerExtension.exports;
+  if (isCanary()) {
-    const testAdapterFactory = new QLTestAdapterFactory(
+    testManager = new TestManager(app, testRunner, cliServer);
-      testHub,
+    ctx.subscriptions.push(testManager);
-      cliServer,
+  } else {
-      dbm,
+    const testExplorerExtension = extensions.getExtension<TestHub>(
+      testExplorerExtensionId,
     );
-    ctx.subscriptions.push(testAdapterFactory);
+    if (testExplorerExtension) {
+      const testHub = testExplorerExtension.exports;
+      const testAdapterFactory = new QLTestAdapterFactory(
+        testHub,
+        testRunner,
+        cliServer,
+      );
+      ctx.subscriptions.push(testAdapterFactory);
 
-    const testUIService = new TestUIService(app, testHub);
+      testManager = new TestUIService(app, testHub);
-    ctx.subscriptions.push(testUIService);
+      ctx.subscriptions.push(testManager);
-
+    }
-    testUiCommands = testUIService.getCommands();
   }
 
+  const testUiCommands = testManager?.getCommands() ?? {};
+
   const astViewer = new AstViewer();
   const astTemplateProvider = new TemplatePrintAstProvider(
     cliServer,

extensions/ql-vscode/src/test-adapter.ts

@@ -1,4 +1,3 @@
-import { access } from "fs-extra";
 import { dirname, extname } from "path";
 import * as vscode from "vscode";
 import {
@@ -20,23 +19,11 @@ import {
   QLTestDirectory,
   QLTestDiscovery,
 } from "./qltest-discovery";
-import {
+import { Event, EventEmitter, CancellationTokenSource } from "vscode";
-  Event,
-  EventEmitter,
-  CancellationTokenSource,
-  CancellationToken,
-} from "vscode";
 import { DisposableObject } from "./pure/disposable-object";
-import { CodeQLCliServer } from "./cli";
+import { CodeQLCliServer, TestCompleted } from "./cli";
-import {
-  getOnDiskWorkspaceFolders,
-  showAndLogExceptionWithTelemetry,
-  showAndLogWarningMessage,
-} from "./helpers";
 import { testLogger } from "./common";
-import { DatabaseItem, DatabaseManager } from "./local-databases";
+import { TestRunner } from "./test-runner";
-import { asError, getErrorMessage } from "./pure/helpers-pure";
-import { redactableError } from "./pure/errors";
 
 /**
  * Get the full path of the `.expected` file for the specified QL test.
@@ -77,8 +64,8 @@ function getTestOutputFile(testPath: string, extension: string): string {
 export class QLTestAdapterFactory extends DisposableObject {
   constructor(
     testHub: TestHub,
+    testRunner: TestRunner,
     cliServer: CodeQLCliServer,
-    databaseManager: DatabaseManager,
   ) {
     super();
 
@@ -87,7 +74,7 @@ export class QLTestAdapterFactory extends DisposableObject {
       new TestAdapterRegistrar(
         testHub,
         (workspaceFolder) =>
-          new QLTestAdapter(workspaceFolder, cliServer, databaseManager),
+          new QLTestAdapter(workspaceFolder, testRunner, cliServer),
       ),
     );
   }
@@ -120,8 +107,8 @@ export class QLTestAdapter extends DisposableObject implements TestAdapter {
 
   constructor(
     public readonly workspaceFolder: vscode.WorkspaceFolder,
-    private readonly cliServer: CodeQLCliServer,
+    private readonly testRunner: TestRunner,
-    private readonly databaseManager: DatabaseManager,
+    cliServer: CodeQLCliServer,
   ) {
     super();
 
@@ -232,110 +219,14 @@ export class QLTestAdapter extends DisposableObject implements TestAdapter {
       tests,
     } as TestRunStartedEvent);
 
-    const currentDatabaseUri =
+    await this.testRunner.run(tests, testLogger, token, (event) =>
-      this.databaseManager.currentDatabaseItem?.databaseUri;
+      this.processTestEvent(event),
-    const databasesUnderTest: DatabaseItem[] = [];
-    for (const database of this.databaseManager.databaseItems) {
-      for (const test of tests) {
-        if (await database.isAffectedByTest(test)) {
-          databasesUnderTest.push(database);
-          break;
-        }
-      }
-    }
-
-    await this.removeDatabasesBeforeTests(databasesUnderTest, token);
-    try {
-      await this.runTests(tests, token);
-    } catch (e) {
-      // CodeQL testing can throw exception even in normal scenarios. For example, if the test run
-      // produces no output (which is normal), the testing command would throw an exception on
-      // unexpected EOF during json parsing. So nothing needs to be done here - all the relevant
-      // error information (if any) should have already been written to the test logger.
-    }
-    await this.reopenDatabasesAfterTests(
-      databasesUnderTest,
-      currentDatabaseUri,
-      token,
     );
 
     this._testStates.fire({ type: "finished" } as TestRunFinishedEvent);
     this.clearTask();
   }
 
-  private async removeDatabasesBeforeTests(
-    databasesUnderTest: DatabaseItem[],
-    token: vscode.CancellationToken,
-  ): Promise<void> {
-    for (const database of databasesUnderTest) {
-      try {
-        await this.databaseManager.removeDatabaseItem(
-          (_) => {
-            /* no progress reporting */
-          },
-          token,
-          database,
-        );
-      } catch (e) {
-        // This method is invoked from Test Explorer UI, and testing indicates that Test
-        // Explorer UI swallows any thrown exception without reporting it to the user.
-        // So we need to display the error message ourselves and then rethrow.
-        void showAndLogExceptionWithTelemetry(
-          redactableError(asError(e))`Cannot remove database ${
-            database.name
-          }: ${getErrorMessage(e)}`,
-        );
-        throw e;
-      }
-    }
-  }
-
-  private async reopenDatabasesAfterTests(
-    databasesUnderTest: DatabaseItem[],
-    currentDatabaseUri: vscode.Uri | undefined,
-    token: vscode.CancellationToken,
-  ): Promise<void> {
-    for (const closedDatabase of databasesUnderTest) {
-      const uri = closedDatabase.databaseUri;
-      if (await this.isFileAccessible(uri)) {
-        try {
-          const reopenedDatabase = await this.databaseManager.openDatabase(
-            (_) => {
-              /* no progress reporting */
-            },
-            token,
-            uri,
-          );
-          await this.databaseManager.renameDatabaseItem(
-            reopenedDatabase,
-            closedDatabase.name,
-          );
-          if (currentDatabaseUri?.toString() === uri.toString()) {
-            await this.databaseManager.setCurrentDatabaseItem(
-              reopenedDatabase,
-              true,
-            );
-          }
-        } catch (e) {
-          // This method is invoked from Test Explorer UI, and testing indicates that Test
-          // Explorer UI swallows any thrown exception without reporting it to the user.
-          // So we need to display the error message ourselves and then rethrow.
-          void showAndLogWarningMessage(`Cannot reopen database ${uri}: ${e}`);
-          throw e;
-        }
-      }
-    }
-  }
-
-  private async isFileAccessible(uri: vscode.Uri): Promise<boolean> {
-    try {
-      await access(uri.fsPath);
-      return true;
-    } catch {
-      return false;
-    }
-  }
-
   private clearTask(): void {
     if (this.runningTask !== undefined) {
       const runningTask = this.runningTask;
@@ -352,49 +243,40 @@ export class QLTestAdapter extends DisposableObject implements TestAdapter {
     }
   }
 
-  private async runTests(
+  private async processTestEvent(event: TestCompleted): Promise<void> {
-    tests: string[],
+    const state = event.pass
-    cancellationToken: CancellationToken,
+      ? "passed"
-  ): Promise<void> {
+      : event.messages?.length
-    const workspacePaths = getOnDiskWorkspaceFolders();
+      ? "errored"
-    for await (const event of this.cliServer.runTests(tests, workspacePaths, {
+      : "failed";
-      cancellationToken,
+    let message: string | undefined;
-      logger: testLogger,
+    if (event.failureDescription || event.diff?.length) {
-    })) {
+      message =
-      const state = event.pass
+        event.failureStage === "RESULT"
-        ? "passed"
+          ? [
-        : event.messages?.length
+              "",
-        ? "errored"
+              `${state}: ${event.test}`,
-        : "failed";
+              event.failureDescription || event.diff?.join("\n"),
-      let message: string | undefined;
+              "",
-      if (event.failureDescription || event.diff?.length) {
+            ].join("\n")
-        message =
+          : [
-          event.failureStage === "RESULT"
+              "",
-            ? [
+              `${event.failureStage?.toLowerCase()} error: ${event.test}`,
-                "",
+              event.failureDescription ||
-                `${state}: ${event.test}`,
+                `${event.messages[0].severity}: ${event.messages[0].message}`,
-                event.failureDescription || event.diff?.join("\n"),
+              "",
-                "",
+            ].join("\n");
-              ].join("\n")
+      void testLogger.log(message);
-            : [
-                "",
-                `${event.failureStage?.toLowerCase()} error: ${event.test}`,
-                event.failureDescription ||
-                  `${event.messages[0].severity}: ${event.messages[0].message}`,
-                "",
-              ].join("\n");
-        void testLogger.log(message);
-      }
-      this._testStates.fire({
-        type: "test",
-        state,
-        test: event.test,
-        message,
-        decorations: event.messages?.map((msg) => ({
-          line: msg.position.line,
-          message: msg.message,
-        })),
-      });
     }
+    this._testStates.fire({
+      type: "test",
+      state,
+      test: event.test,
+      message,
+      decorations: event.messages?.map((msg) => ({
+        line: msg.position.line,
+        message: msg.message,
+      })),
+    });
   }
 }

extensions/ql-vscode/src/test-manager-base.ts

@@ -0,0 +1,76 @@
+import { copy, createFile, lstat, pathExists } from "fs-extra";
+import { TestUICommands } from "./common/commands";
+import { DisposableObject } from "./pure/disposable-object";
+import { getActualFile, getExpectedFile } from "./test-adapter";
+import { TestItem, TextDocumentShowOptions, Uri, window } from "vscode";
+import { showAndLogWarningMessage } from "./helpers";
+import { basename } from "path";
+import { App } from "./common/app";
+import { TestTreeNode } from "./test-tree-node";
+
+export type TestNode = TestTreeNode | TestItem;
+
+/**
+ * Base class for both the legacy and new test services. Implements commands that are common to
+ * both.
+ */
+export abstract class TestManagerBase extends DisposableObject {
+  protected constructor(private readonly app: App) {
+    super();
+  }
+
+  public getCommands(): TestUICommands {
+    return {
+      "codeQLTests.showOutputDifferences":
+        this.showOutputDifferences.bind(this),
+      "codeQLTests.acceptOutput": this.acceptOutput.bind(this),
+    };
+  }
+
+  /** Override to compute the path of the test file from the selected node. */
+  protected abstract getTestPath(node: TestNode): string;
+
+  private async acceptOutput(node: TestNode): Promise<void> {
+    const testPath = this.getTestPath(node);
+    const stat = await lstat(testPath);
+    if (stat.isFile()) {
+      const expectedPath = getExpectedFile(testPath);
+      const actualPath = getActualFile(testPath);
+      await copy(actualPath, expectedPath, { overwrite: true });
+    }
+  }
+
+  private async showOutputDifferences(node: TestNode): Promise<void> {
+    const testId = this.getTestPath(node);
+    const stat = await lstat(testId);
+    if (stat.isFile()) {
+      const expectedPath = getExpectedFile(testId);
+      const expectedUri = Uri.file(expectedPath);
+      const actualPath = getActualFile(testId);
+      const options: TextDocumentShowOptions = {
+        preserveFocus: true,
+        preview: true,
+      };
+
+      if (!(await pathExists(expectedPath))) {
+        void showAndLogWarningMessage(
+          `'${basename(expectedPath)}' does not exist. Creating an empty file.`,
+        );
+        await createFile(expectedPath);
+      }
+
+      if (await pathExists(actualPath)) {
+        const actualUri = Uri.file(actualPath);
+        await this.app.commands.execute(
+          "vscode.diff",
+          expectedUri,
+          actualUri,
+          `Expected vs. Actual for ${basename(testId)}`,
+          options,
+        );
+      } else {
+        await window.showTextDocument(expectedUri, options);
+      }
+    }
+  }
+}

extensions/ql-vscode/src/test-manager.ts

@@ -0,0 +1,371 @@
+import { readFile } from "fs-extra";
+import {
+  CancellationToken,
+  Location,
+  Range,
+  TestController,
+  TestItem,
+  TestMessage,
+  TestRun,
+  TestRunProfileKind,
+  TestRunRequest,
+  Uri,
+  WorkspaceFolder,
+  WorkspaceFoldersChangeEvent,
+  tests,
+  workspace,
+} from "vscode";
+import { DisposableObject } from "./pure/disposable-object";
+import {
+  QLTestDirectory,
+  QLTestDiscovery,
+  QLTestFile,
+  QLTestNode,
+} from "./qltest-discovery";
+import { CodeQLCliServer } from "./cli";
+import { getErrorMessage } from "./pure/helpers-pure";
+import { BaseLogger, LogOptions } from "./common";
+import { TestRunner } from "./test-runner";
+import { TestManagerBase } from "./test-manager-base";
+import { App } from "./common/app";
+
+/**
+ * Returns the complete text content of the specified file. If there is an error reading the file,
+ * an error message is added to `testMessages` and this function returns undefined.
+ */
+async function tryReadFileContents(
+  path: string,
+  testMessages: TestMessage[],
+): Promise<string | undefined> {
+  try {
+    return await readFile(path, { encoding: "utf-8" });
+  } catch (e) {
+    testMessages.push(
+      new TestMessage(
+        `Error reading from file '${path}': ${getErrorMessage(e)}`,
+      ),
+    );
+    return undefined;
+  }
+}
+
+function forEachTest(testItem: TestItem, op: (test: TestItem) => void): void {
+  if (testItem.children.size > 0) {
+    // This is a directory, so recurse into the children.
+    for (const [, child] of testItem.children) {
+      forEachTest(child, op);
+    }
+  } else {
+    // This is a leaf node, so it's a test.
+    op(testItem);
+  }
+}
+
+/**
+ * Implementation of `BaseLogger` that logs to the output of a `TestRun`.
+ */
+class TestRunLogger implements BaseLogger {
+  public constructor(private readonly testRun: TestRun) {}
+
+  public async log(message: string, options?: LogOptions): Promise<void> {
+    // "\r\n" because that's what the test terminal wants.
+    const lineEnding = options?.trailingNewline === false ? "" : "\r\n";
+    this.testRun.appendOutput(message + lineEnding);
+  }
+}
+
+/**
+ * Handles test dicovery for a specific workspace folder, and reports back to `TestManager`.
+ */
+class WorkspaceFolderHandler extends DisposableObject {
+  private readonly testDiscovery: QLTestDiscovery;
+
+  public constructor(
+    private readonly workspaceFolder: WorkspaceFolder,
+    private readonly testUI: TestManager,
+    cliServer: CodeQLCliServer,
+  ) {
+    super();
+
+    this.testDiscovery = new QLTestDiscovery(workspaceFolder, cliServer);
+    this.push(
+      this.testDiscovery.onDidChangeTests(this.handleDidChangeTests, this),
+    );
+    this.testDiscovery.refresh();
+  }
+
+  private handleDidChangeTests(): void {
+    const testDirectory = this.testDiscovery.testDirectory;
+
+    this.testUI.updateTestsForWorkspaceFolder(
+      this.workspaceFolder,
+      testDirectory,
+    );
+  }
+}
+
+/**
+ * Service that populates the VS Code "Test Explorer" panel for CodeQL, and handles running and
+ * debugging of tests.
+ */
+export class TestManager extends TestManagerBase {
+  private readonly testController: TestController = tests.createTestController(
+    "codeql",
+    "Fancy CodeQL Tests",
+  );
+
+  /**
+   * Maps from each workspace folder being tracked to the `WorkspaceFolderHandler` responsible for
+   * tracking it.
+   */
+  private readonly workspaceFolderHandlers = new Map<
+    WorkspaceFolder,
+    WorkspaceFolderHandler
+  >();
+
+  public constructor(
+    app: App,
+    private readonly testRunner: TestRunner,
+    private readonly cliServer: CodeQLCliServer,
+  ) {
+    super(app);
+
+    this.testController.createRunProfile(
+      "Run",
+      TestRunProfileKind.Run,
+      this.run.bind(this),
+    );
+
+    // Start by tracking whatever folders are currently in the workspace.
+    this.startTrackingWorkspaceFolders(workspace.workspaceFolders ?? []);
+
+    // Listen for changes to the set of workspace folders.
+    workspace.onDidChangeWorkspaceFolders(
+      this.handleDidChangeWorkspaceFolders,
+      this,
+    );
+  }
+
+  public dispose(): void {
+    this.workspaceFolderHandlers.clear(); // These will be disposed in the `super.dispose()` call.
+    super.dispose();
+  }
+
+  protected getTestPath(node: TestItem): string {
+    if (node.uri === undefined || node.uri.scheme !== "file") {
+      throw new Error("Selected test is not a CodeQL test.");
+    }
+    return node.uri!.fsPath;
+  }
+
+  /** Start tracking tests in the specified workspace folders. */
+  private startTrackingWorkspaceFolders(
+    workspaceFolders: readonly WorkspaceFolder[],
+  ): void {
+    for (const workspaceFolder of workspaceFolders) {
+      const workspaceFolderHandler = new WorkspaceFolderHandler(
+        workspaceFolder,
+        this,
+        this.cliServer,
+      );
+      this.track(workspaceFolderHandler);
+      this.workspaceFolderHandlers.set(workspaceFolder, workspaceFolderHandler);
+    }
+  }
+
+  /** Stop tracking tests in the specified workspace folders. */
+  private stopTrackingWorkspaceFolders(
+    workspaceFolders: readonly WorkspaceFolder[],
+  ): void {
+    for (const workspaceFolder of workspaceFolders) {
+      const workspaceFolderHandler =
+        this.workspaceFolderHandlers.get(workspaceFolder);
+      if (workspaceFolderHandler !== undefined) {
+        // Delete the root item for this workspace folder, if any.
+        this.testController.items.delete(workspaceFolder.uri.toString());
+        this.disposeAndStopTracking(workspaceFolderHandler);
+        this.workspaceFolderHandlers.delete(workspaceFolder);
+      }
+    }
+  }
+
+  private handleDidChangeWorkspaceFolders(
+    e: WorkspaceFoldersChangeEvent,
+  ): void {
+    this.startTrackingWorkspaceFolders(e.added);
+    this.stopTrackingWorkspaceFolders(e.removed);
+  }
+
+  /**
+   * Update the test controller when we discover changes to the tests in the workspace folder.
+   */
+  public updateTestsForWorkspaceFolder(
+    workspaceFolder: WorkspaceFolder,
+    testDirectory: QLTestDirectory | undefined,
+  ): void {
+    if (testDirectory !== undefined) {
+      // Adding an item with the same ID as an existing item will replace it, which is exactly what
+      // we want.
+      // Test discovery returns a root `QLTestDirectory` representing the workspace folder itself,
+      // named after the `WorkspaceFolder` object's `name` property. We can map this directly to a
+      // `TestItem`.
+      this.testController.items.add(
+        this.createTestItemTree(testDirectory, true),
+      );
+    } else {
+      // No tests, so delete any existing item.
+      this.testController.items.delete(workspaceFolder.uri.toString());
+    }
+  }
+
+  /**
+   * Creates a tree of `TestItem`s from the root `QlTestNode` provided by test discovery.
+   */
+  private createTestItemTree(node: QLTestNode, isRoot: boolean): TestItem {
+    // Prefix the ID to identify it as a directory or a test
+    const itemType = node instanceof QLTestDirectory ? "dir" : "test";
+    const testItem = this.testController.createTestItem(
+      // For the root of a workspace folder, use the full path as the ID. Otherwise, use the node's
+      // name as the ID, since it's shorter but still unique.
+      `${itemType} ${isRoot ? node.path : node.name}`,
+      node.name,
+      Uri.file(node.path),
+    );
+
+    for (const childNode of node.children) {
+      const childItem = this.createTestItemTree(childNode, false);
+      if (childNode instanceof QLTestFile) {
+        childItem.range = new Range(0, 0, 0, 0);
+      }
+      testItem.children.add(childItem);
+    }
+
+    return testItem;
+  }
+
+  /**
+   * Run the tests specified by the `TestRunRequest` parameter.
+   */
+  private async run(
+    request: TestRunRequest,
+    token: CancellationToken,
+  ): Promise<void> {
+    const testsToRun = this.computeTestsToRun(request);
+    const testRun = this.testController.createTestRun(request, undefined, true);
+    try {
+      const tests: string[] = [];
+      testsToRun.forEach((t) => {
+        testRun.enqueued(t);
+        tests.push(t.uri!.fsPath);
+      });
+
+      const logger = new TestRunLogger(testRun);
+
+      await this.testRunner.run(tests, logger, token, async (event) => {
+        // Pass the test path from the event through `Uri` and back via `fsPath` so that it matches
+        // the canonicalization of the URI that we used to create the `TestItem`.
+        const testItem = testsToRun.get(Uri.file(event.test).fsPath);
+        if (testItem === undefined) {
+          throw new Error(
+            `Unexpected result from unknown test '${event.test}'.`,
+          );
+        }
+
+        const duration = event.compilationMs + event.evaluationMs;
+        if (event.pass) {
+          testRun.passed(testItem, duration);
+        } else {
+          // Construct a list of `TestMessage`s to report for the failure.
+          const testMessages: TestMessage[] = [];
+          if (event.failureDescription !== undefined) {
+            testMessages.push(new TestMessage(event.failureDescription));
+          }
+          if (event.diff?.length && event.actual !== undefined) {
+            // Actual results differ from expected results. Read both sets of results and create a
+            // diff to put in the message.
+            const expected = await tryReadFileContents(
+              event.expected,
+              testMessages,
+            );
+            const actual = await tryReadFileContents(
+              event.actual,
+              testMessages,
+            );
+            if (expected !== undefined && actual !== undefined) {
+              testMessages.push(
+                TestMessage.diff(
+                  "Actual output differs from expected",
+                  expected,
+                  actual,
+                ),
+              );
+            }
+          }
+          if (event.messages?.length > 0) {
+            // The test didn't make it far enough to produce results. Transform any error messages
+            // into `TestMessage`s and report the test as "errored".
+            const testMessages = event.messages.map((m) => {
+              const location = new Location(
+                Uri.file(m.position.fileName),
+                new Range(
+                  m.position.line - 1,
+                  m.position.column - 1,
+                  m.position.endLine - 1,
+                  m.position.endColumn - 1,
+                ),
+              );
+              const testMessage = new TestMessage(m.message);
+              testMessage.location = location;
+              return testMessage;
+            });
+            testRun.errored(testItem, testMessages, duration);
+          } else {
+            // Results didn't match expectations. Report the test as "failed".
+            if (testMessages.length === 0) {
+              // If we managed to get here without creating any `TestMessage`s, create a default one
+              // here. Any failed test needs at least one message.
+              testMessages.push(new TestMessage("Test failed"));
+            }
+            testRun.failed(testItem, testMessages, duration);
+          }
+        }
+      });
+    } finally {
+      testRun.end();
+    }
+  }
+
+  /**
+   * Computes the set of tests to run as specified in the `TestRunRequest` object.
+   */
+  private computeTestsToRun(request: TestRunRequest): Map<string, TestItem> {
+    const testsToRun = new Map<string, TestItem>();
+    if (request.include !== undefined) {
+      // Include these tests, recursively expanding test directories into their list of contained
+      // tests.
+      for (const includedTestItem of request.include) {
+        forEachTest(includedTestItem, (testItem) =>
+          testsToRun.set(testItem.uri!.fsPath, testItem),
+        );
+      }
+    } else {
+      // Include all of the tests.
+      for (const [, includedTestItem] of this.testController.items) {
+        forEachTest(includedTestItem, (testItem) =>
+          testsToRun.set(testItem.uri!.fsPath, testItem),
+        );
+      }
+    }
+    if (request.exclude !== undefined) {
+      // Exclude the specified tests from the set we've computed so far, again recursively expanding
+      // test directories into their list of contained tests.
+      for (const excludedTestItem of request.exclude) {
+        forEachTest(excludedTestItem, (testItem) =>
+          testsToRun.delete(testItem.uri!.fsPath),
+        );
+      }
+    }
+
+    return testsToRun;
+  }
+}

extensions/ql-vscode/src/test-runner.ts

@@ -0,0 +1,136 @@
+import { CancellationToken, Uri } from "vscode";
+import { CodeQLCliServer, TestCompleted } from "./cli";
+import { DatabaseItem, DatabaseManager } from "./local-databases";
+import {
+  getOnDiskWorkspaceFolders,
+  showAndLogExceptionWithTelemetry,
+  showAndLogWarningMessage,
+} from "./helpers";
+import { asError, getErrorMessage } from "./pure/helpers-pure";
+import { redactableError } from "./pure/errors";
+import { access } from "fs-extra";
+import { BaseLogger } from "./common";
+import { DisposableObject } from "./pure/disposable-object";
+
+async function isFileAccessible(uri: Uri): Promise<boolean> {
+  try {
+    await access(uri.fsPath);
+    return true;
+  } catch {
+    return false;
+  }
+}
+
+export class TestRunner extends DisposableObject {
+  public constructor(
+    private readonly databaseManager: DatabaseManager,
+    private readonly cliServer: CodeQLCliServer,
+  ) {
+    super();
+  }
+
+  public async run(
+    tests: string[],
+    logger: BaseLogger,
+    token: CancellationToken,
+    eventHandler: (event: TestCompleted) => Promise<void>,
+  ): Promise<void> {
+    const currentDatabaseUri =
+      this.databaseManager.currentDatabaseItem?.databaseUri;
+    const databasesUnderTest: DatabaseItem[] = [];
+    for (const database of this.databaseManager.databaseItems) {
+      for (const test of tests) {
+        if (await database.isAffectedByTest(test)) {
+          databasesUnderTest.push(database);
+          break;
+        }
+      }
+    }
+
+    await this.removeDatabasesBeforeTests(databasesUnderTest, token);
+    try {
+      const workspacePaths = getOnDiskWorkspaceFolders();
+      for await (const event of this.cliServer.runTests(tests, workspacePaths, {
+        cancellationToken: token,
+        logger,
+      })) {
+        await eventHandler(event);
+      }
+    } catch (e) {
+      // CodeQL testing can throw exception even in normal scenarios. For example, if the test run
+      // produces no output (which is normal), the testing command would throw an exception on
+      // unexpected EOF during json parsing. So nothing needs to be done here - all the relevant
+      // error information (if any) should have already been written to the test logger.
+    } finally {
+      await this.reopenDatabasesAfterTests(
+        databasesUnderTest,
+        currentDatabaseUri,
+        token,
+      );
+    }
+  }
+
+  private async removeDatabasesBeforeTests(
+    databasesUnderTest: DatabaseItem[],
+    token: CancellationToken,
+  ): Promise<void> {
+    for (const database of databasesUnderTest) {
+      try {
+        await this.databaseManager.removeDatabaseItem(
+          (_) => {
+            /* no progress reporting */
+          },
+          token,
+          database,
+        );
+      } catch (e) {
+        // This method is invoked from Test Explorer UI, and testing indicates that Test
+        // Explorer UI swallows any thrown exception without reporting it to the user.
+        // So we need to display the error message ourselves and then rethrow.
+        void showAndLogExceptionWithTelemetry(
+          redactableError(asError(e))`Cannot remove database ${
+            database.name
+          }: ${getErrorMessage(e)}`,
+        );
+        throw e;
+      }
+    }
+  }
+
+  private async reopenDatabasesAfterTests(
+    databasesUnderTest: DatabaseItem[],
+    currentDatabaseUri: Uri | undefined,
+    token: CancellationToken,
+  ): Promise<void> {
+    for (const closedDatabase of databasesUnderTest) {
+      const uri = closedDatabase.databaseUri;
+      if (await isFileAccessible(uri)) {
+        try {
+          const reopenedDatabase = await this.databaseManager.openDatabase(
+            (_) => {
+              /* no progress reporting */
+            },
+            token,
+            uri,
+          );
+          await this.databaseManager.renameDatabaseItem(
+            reopenedDatabase,
+            closedDatabase.name,
+          );
+          if (currentDatabaseUri?.toString() === uri.toString()) {
+            await this.databaseManager.setCurrentDatabaseItem(
+              reopenedDatabase,
+              true,
+            );
+          }
+        } catch (e) {
+          // This method is invoked from Test Explorer UI, and testing indicates that Test
+          // Explorer UI swallows any thrown exception without reporting it to the user.
+          // So we need to display the error message ourselves and then rethrow.
+          void showAndLogWarningMessage(`Cannot reopen database ${uri}: ${e}`);
+          throw e;
+        }
+      }
+    }
+  }
+}

extensions/ql-vscode/src/test-ui.ts

@@ -1,6 +1,3 @@
-import { lstat, copy, pathExists, createFile } from "fs-extra";
-import { basename } from "path";
-import { Uri, TextDocumentShowOptions, window } from "vscode";
 import {
   TestHub,
   TestController,
@@ -10,13 +7,11 @@ import {
   TestEvent,
   TestSuiteEvent,
 } from "vscode-test-adapter-api";
-
-import { showAndLogWarningMessage } from "./helpers";
 import { TestTreeNode } from "./test-tree-node";
 import { DisposableObject } from "./pure/disposable-object";
-import { QLTestAdapter, getExpectedFile, getActualFile } from "./test-adapter";
+import { QLTestAdapter } from "./test-adapter";
-import { TestUICommands } from "./common/commands";
 import { App } from "./common/app";
+import { TestManagerBase } from "./test-manager-base";
 
 type VSCodeTestEvent =
   | TestRunStartedEvent
@@ -42,23 +37,15 @@ class QLTestListener extends DisposableObject {
 /**
  * Service that implements all UI and commands for QL tests.
  */
-export class TestUIService extends DisposableObject implements TestController {
+export class TestUIService extends TestManagerBase implements TestController {
   private readonly listeners: Map<TestAdapter, QLTestListener> = new Map();
 
-  constructor(private readonly app: App, private readonly testHub: TestHub) {
+  public constructor(app: App, private readonly testHub: TestHub) {
-    super();
+    super(app);
 
     testHub.registerTestController(this);
   }
 
-  public getCommands(): TestUICommands {
-    return {
-      "codeQLTests.showOutputDifferences":
-        this.showOutputDifferences.bind(this),
-      "codeQLTests.acceptOutput": this.acceptOutput.bind(this),
-    };
-  }
-
   public dispose(): void {
     this.testHub.unregisterTestController(this);
 
@@ -75,47 +62,7 @@ export class TestUIService extends DisposableObject implements TestController {
     }
   }
 
-  private async acceptOutput(node: TestTreeNode): Promise<void> {
+  protected getTestPath(node: TestTreeNode): string {
-    const testId = node.info.id;
+    return node.info.id;
-    const stat = await lstat(testId);
-    if (stat.isFile()) {
-      const expectedPath = getExpectedFile(testId);
-      const actualPath = getActualFile(testId);
-      await copy(actualPath, expectedPath, { overwrite: true });
-    }
-  }
-
-  private async showOutputDifferences(node: TestTreeNode): Promise<void> {
-    const testId = node.info.id;
-    const stat = await lstat(testId);
-    if (stat.isFile()) {
-      const expectedPath = getExpectedFile(testId);
-      const expectedUri = Uri.file(expectedPath);
-      const actualPath = getActualFile(testId);
-      const options: TextDocumentShowOptions = {
-        preserveFocus: true,
-        preview: true,
-      };
-
-      if (!(await pathExists(expectedPath))) {
-        void showAndLogWarningMessage(
-          `'${basename(expectedPath)}' does not exist. Creating an empty file.`,
-        );
-        await createFile(expectedPath);
-      }
-
-      if (await pathExists(actualPath)) {
-        const actualUri = Uri.file(actualPath);
-        await this.app.commands.execute(
-          "vscode.diff",
-          expectedUri,
-          actualUri,
-          `Expected vs. Actual for ${basename(testId)}`,
-          options,
-        );
-      } else {
-        await window.showTextDocument(expectedUri, options);
-      }
-    }
   }
 }