feat: implement basic related code support

package.json

@@ -3,7 +3,8 @@
   "displayName": "VS Code Selfhost Test Provider",
   "description": "Test provider for the VS Code project",
   "enabledApiProposals": [
-    "testObserver"
+    "testObserver",
+    "testRelatedCode"
   ],
   "version": "0.3.9",
   "publisher": "ms-vscode",

src/extension.ts

@@ -4,6 +4,7 @@
 
 import * as vscode from 'vscode';
 import { FailingDeepStrictEqualAssertFixer } from './failingDeepStrictEqualAssertFixer';
+import { updateRelatedCodeForImplementation } from './relatedCode';
 import { scanTestOutput } from './testOutputScanner';
 import { guessWorkspaceFolder, itemData, TestCase, TestFile } from './testTree';
 import { BrowserTestRunner, PlatformTestRunner, VSCodeTestRunner } from './vscodeTestRunner';
@@ -37,32 +38,34 @@ export async function activate(context: vscode.ExtensionContext) {
   };
 
   let runQueue = Promise.resolve();
-  const createRunHandler = (
-    runnerCtor: { new (folder: vscode.WorkspaceFolder): VSCodeTestRunner },
-    debug: boolean,
-    args: string[] = []
-  ) => async (req: vscode.TestRunRequest, cancellationToken: vscode.CancellationToken) => {
-    const folder = await guessWorkspaceFolder();
-    if (!folder) {
-      return;
-    }
+  const createRunHandler =
+    (
+      runnerCtor: { new (folder: vscode.WorkspaceFolder): VSCodeTestRunner },
+      debug: boolean,
+      args: string[] = []
+    ) =>
+    async (req: vscode.TestRunRequest, cancellationToken: vscode.CancellationToken) => {
+      const folder = await guessWorkspaceFolder();
+      if (!folder) {
+        return;
+      }
 
-    const runner = new runnerCtor(folder);
-    const map = await getPendingTestMap(ctrl, req.include ?? gatherTestItems(ctrl.items));
-    const task = ctrl.createTestRun(req);
-    for (const test of map.values()) {
-      task.enqueued(test);
-    }
+      const runner = new runnerCtor(folder);
+      const map = await getPendingTestMap(ctrl, req.include ?? gatherTestItems(ctrl.items));
+      const task = ctrl.createTestRun(req);
+      for (const test of map.values()) {
+        task.enqueued(test);
+      }
 
-    return (runQueue = runQueue.then(async () => {
-      await scanTestOutput(
-        map,
-        task,
-        debug ? await runner.debug(args, req.include) : await runner.run(args, req.include),
-        cancellationToken
-      );
-    }));
-  };
+      return (runQueue = runQueue.then(async () => {
+        await scanTestOutput(
+          map,
+          task,
+          debug ? await runner.debug(args, req.include) : await runner.run(args, req.include),
+          cancellationToken
+        );
+      }));
+    };
 
   ctrl.createRunProfile(
     'Run in Electron',
@@ -99,6 +102,8 @@ export async function activate(context: vscode.ExtensionContext) {
     const data = node && itemData.get(node);
     if (data instanceof TestFile) {
       data.updateFromContents(ctrl, e.getText(), node!);
+    } else {
+      updateRelatedCodeForImplementation(e.uri, ctrl.items, e.getText());
     }
   }
 

src/relatedCode.ts

@@ -0,0 +1,91 @@
+/*---------------------------------------------------------
+ * Copyright (C) Microsoft Corporation. All rights reserved.
+ *--------------------------------------------------------*/
+
+import * as vscode from 'vscode';
+import { getContentFromFilesystem, itemData, TestFile } from './testTree';
+
+type ChildIterator = { forEach(fn: (t: vscode.TestItem) => void): void };
+
+/**
+ * Updates the "related code" of tests in the file.
+ */
+export const updatedRelatedCodeForTestFile = async (file: TestFile, testItems: ChildIterator) => {
+  if (!file.relatedCodeFile) {
+    return;
+  }
+
+  const content = await getContentFromFilesystem(file.relatedCodeFile);
+  if (!content) {
+    return;
+  }
+
+  return updatedRelatedCodeFromContents(file.relatedCodeFile, testItems, content);
+};
+
+/**
+ * Updates the "related code" of tests that have implementations in
+ * the file. Assumes the tests are already discovered, doesn't
+ * eagerly do discovery.
+ */
+export const updateRelatedCodeForImplementation = async (
+  file: vscode.Uri,
+  children: vscode.TestItemCollection,
+  contents: string
+) => {
+  children.forEach(child => {
+    const item = itemData.get(child);
+    if (item instanceof TestFile && item.relatedCodeFile?.path === file.path) {
+      updatedRelatedCodeFromContents(file, child.children, contents);
+    }
+  });
+};
+
+const isDocCommentLike = /^\s* \* /;
+
+/**
+ * Updates related code in each TestItem from the contents of the
+ * implementation code. Assumes that each test involves testing a
+ * method, function, or class and is labeled correspondingly. This
+ * may not hold true for all tests, but is true for some tests.
+ */
+const updatedRelatedCodeFromContents = (
+  uri: vscode.Uri,
+  testItems: ChildIterator,
+  relatedCodeFileContents: string
+) => {
+  const lines = relatedCodeFileContents.split(/\r?\n/g);
+  const addRelatedCode = (testItem: vscode.TestItem) => {
+    let found = false;
+    for (let line = 0; line < lines.length; line++) {
+      const contents = lines[line];
+      if (contents.startsWith('//') || isDocCommentLike.test(contents)) {
+        continue;
+      }
+
+      const index = contents.indexOf(testItem.label);
+      if (index === -1) {
+        continue;
+      }
+
+      testItem.relatedCode = [
+        new vscode.Location(
+          uri,
+          new vscode.Range(
+            new vscode.Position(line, index),
+            new vscode.Position(line, index + testItem.label.length)
+          )
+        ),
+      ];
+      found = true;
+      break;
+    }
+
+    if (!found) {
+      testItem.relatedCode = undefined;
+    }
+
+    testItem.children.forEach(addRelatedCode);
+  };
+  testItems.forEach(addRelatedCode);
+};

src/sourceUtils.ts

@@ -13,7 +13,19 @@ export const enum Action {
   Recurse,
 }
 
+export class ImportDeclaration {
+  public get basename() {
+    return this.text.split('/').pop()!;
+  }
+
+  constructor(public readonly text: string) {}
+}
+
 export const extractTestFromNode = (src: ts.SourceFile, node: ts.Node, parent: VSCodeTest) => {
+  if (ts.isImportDeclaration(node) && ts.isStringLiteralLike(node.moduleSpecifier)) {
+    return new ImportDeclaration(node.moduleSpecifier.text);
+  }
+
   if (!ts.isCallExpression(node)) {
     return Action.Recurse;
   }

src/testTree.ts

@@ -6,7 +6,8 @@ import { join, relative } from 'path';
 import * as ts from 'typescript';
 import { TextDecoder } from 'util';
 import * as vscode from 'vscode';
-import { Action, extractTestFromNode } from './sourceUtils';
+import { updatedRelatedCodeForTestFile } from './relatedCode';
+import { Action, extractTestFromNode, ImportDeclaration } from './sourceUtils';
 
 const textDecoder = new TextDecoder('utf-8');
 
@@ -50,11 +51,15 @@ export const getContentFromFilesystem: ContentGetter = async uri => {
 
 export class TestFile {
   public hasBeenRead = false;
+  public relatedCodeFile: vscode.Uri | undefined;
+  private readonly basename: string;
 
   constructor(
     public readonly uri: vscode.Uri,
     public readonly workspaceFolder: vscode.WorkspaceFolder
-  ) {}
+  ) {
+    this.basename = uri.path.split('/').pop()!;
+  }
 
   public getId() {
     return this.uri.toString().toLowerCase();
@@ -68,7 +73,7 @@ export class TestFile {
     try {
       const content = await getContentFromFilesystem(item.uri!);
       item.error = undefined;
-      this.updateFromContents(controller, content, item);
+      await this.updateFromContents(controller, content, item);
     } catch (e) {
       item.error = (e as Error).stack;
     }
@@ -77,7 +82,7 @@ export class TestFile {
   /**
    * Refreshes all tests in this file, `sourceReader` provided by the root.
    */
-  public updateFromContents(
+  public async updateFromContents(
     controller: vscode.TestController,
     content: string,
     file: vscode.TestItem
@@ -106,6 +111,19 @@ export class TestFile {
           return;
         }
 
+        if (childData instanceof ImportDeclaration) {
+          // yes, we make a lof of assumptions here. But it's all heuristic.
+          // We could get higher accuracy by using the type checker, but that
+          // would be much, much slower (and more complicated,
+          // especially for edit files)
+          if (childData.text.startsWith('vs/') && this.basename.includes(childData.basename)) {
+            this.relatedCodeFile = vscode.Uri.file(
+              join(this.workspaceFolder.uri.fsPath, 'src', `${childData.text}.ts`)
+            );
+          }
+          return;
+        }
+
         const id = `${file.uri}/${childData.fullName}`.toLowerCase();
 
         // Skip duplicated tests. They won't run correctly with the way
@@ -127,6 +145,7 @@ export class TestFile {
       };
 
       ts.forEachChild(ast, traverse);
+      await updatedRelatedCodeForTestFile(this, parents[0].children);
       file.error = undefined;
       file.children.replace(parents[0].children);
       this.hasBeenRead = true;

tsconfig.json

@@ -1,9 +1,9 @@
 {
   "compilerOptions": {
     "module": "esnext",
-    "target": "ES2019",
+    "target": "ES2020",
     "outDir": "out",
-    "lib": ["ES2019"],
+    "lib": ["ES2020"],
     "sourceMap": true,
     "rootDir": "src",
     "moduleResolution": "node",