implement test specifier configuration

package.json

@@ -147,6 +147,52 @@
               }
             }
           },
+          "nodejs-testing.testSpecifiers": {
+            "type": "array",
+            "markdownDescription": "_Advanced_: A list of specifiers that indicate test function to search for:\nIt defaults to:\n\n```json\n[\n  {\n    \"import\": \"node:test\",\n    \"name\": [\"default\", \"it\", \"test\", \"describe\", \"suite\"]\n  }\n]\n```\n\nBut in case your test function is wrapped, you can specify it with a relative import:\nNOTE: relative imports must be prefixed with ./\n\n```json\n[\n  {\n    \"import\": \"./test/utils.js\",\n    \"name\": \"test\"\n  }\n]\n```\n",
+            "default": [
+              {
+                "import": "node:test",
+                "name": [
+                  "default",
+                  "it",
+                  "test",
+                  "describe",
+                  "suite"
+                ]
+              }
+            ],
+            "items": {
+              "type": "object",
+              "default": {
+                "import": "node:test",
+                "name": [
+                  "default",
+                  "it",
+                  "test",
+                  "describe",
+                  "suite"
+                ]
+              },
+              "required": [
+                "import",
+                "name"
+              ],
+              "properties": {
+                "import": {
+                  "type": "string",
+                  "markdownDescription": "A package specifier (i.e. node:test) or workspace-relative path beginning with ./ (like ./test/utils.js) that indicates where the 'test' function can be imported from in your codebase"
+                },
+                "name": {
+                  "type": "array",
+                  "markdownDescription": "A list of functions that are imported from `import` that should be treated as test functions",
+                  "items": {
+                    "type": "string"
+                  }
+                }
+              }
+            }
+          },
           "nodejs-testing.envFile": {
             "type": "string",
             "markdownDescription": "Absolute path to a file containing environment variable definitions.\n\nNote: template parameters like ${workspaceFolder} will be resolved.",

src/controller.ts

@@ -5,12 +5,13 @@ import picomatch from "picomatch";
 import * as vscode from "vscode";
 import { coverageContext } from "./coverage";
 import { DisposableStore, MutableDisposable } from "./disposable";
+import { ExtensionConfig } from "./extension-config";
 import { last } from "./iterable";
 import { ICreateOpts, ItemType, getContainingItemsForFile, testMetadata } from "./metadata";
 import { IParsedNode, parseSource } from "./parsing";
 import { RunHandler, TestRunner } from "./runner";
 import { ISourceMapMaintainer, SourceMapStore } from "./source-map-store";
-import { ExtensionConfig } from './extension-config';
+import type { TestFunctionSpecifierConfig } from "./test-function-specifier-config";
 
 const diagnosticCollection = vscode.languages.createDiagnosticCollection("nodejs-testing-dupes");
 
@@ -61,6 +62,11 @@ export class Controller {
     }
   >();
 
+  /**
+   * The configuration which defines which functions should be treated as tests
+   */
+  private readonly testSpecifiers: TestFunctionSpecifierConfig[];
+
   /** Change emtiter used for testing, to pick up when the file watcher detects a chagne */
   public readonly onDidChange = this.didChangeEmitter.event;
   /** Handler for a normal test run */
@@ -78,7 +84,9 @@ export class Controller {
     include: string[],
     exclude: string[],
     extensionConfigs: ExtensionConfig[],
+    testSpecifiers: TestFunctionSpecifierConfig[],
   ) {
+    this.testSpecifiers = testSpecifiers;
     this.disposable.add(ctrl);
     this.disposable.add(runner);
     const extensions = extensionConfigs.flatMap((x) => x.extensions);
@@ -157,15 +165,18 @@ export class Controller {
     }
   }
 
+  /**
+   * Re-check this file for tests, and add them to the UI.
+   * Assumes that the URI has already passed `this.includeTest`
+   *
+   * @param folder the workspace folder this test file belongs to
+   * @param uri the URI of the file in question to reparse and check for tests
+   * @param contents the file contents of uri - to be used as an optimization
+   */
   private async _syncFile(uri: vscode.Uri, contents?: string) {
+    const folder = vscode.workspace.getWorkspaceFolder(uri);
     contents ??= await fs.readFile(uri.fsPath, "utf8");
 
-    // cheap test for relevancy:
-    if (!contents.includes("node:test")) {
-      this.deleteFileTests(uri.toString());
-      return;
-    }
-
     // avoid re-parsing if the contents are the same (e.g. if a file is edited
     // and then saved.)
     const previous = this.testsInFiles.get(uri.toString());
@@ -174,7 +185,7 @@ export class Controller {
       return;
     }
 
-    const tree = parseSource(contents);
+    const tree = parseSource(contents, folder, uri, this.testSpecifiers);
     if (!tree.length) {
       this.deleteFileTests(uri.toString());
       return;
@@ -224,7 +235,7 @@ export class Controller {
       return item;
     };
 
-    // We assume that all tests inside a top-leve describe/test are from the same
+    // We assume that all tests inside a top-level describe/test are from the same
     // source file. This is probably a good assumption. Likewise we assume that a single
     // a single describe/test is not split between different files.
     const newTestsInFile = new Map<string, vscode.TestItem>();
@@ -295,8 +306,8 @@ export class Controller {
       new vscode.RelativePattern(this.wf, `**/*`),
     ));
 
-    watcher.onDidCreate((uri) => this.includeTest(uri.fsPath) && this._syncFile(uri));
-    watcher.onDidChange((uri) => this.includeTest(uri.fsPath) && this._syncFile(uri));
+    watcher.onDidCreate((uri) => this.syncFile(uri));
+    watcher.onDidChange((uri) => this.syncFile(uri));
     watcher.onDidDelete((uri) => {
       const prefix = uri.toString();
       for (const key of this.testsInFiles.keys()) {

src/extension.ts

@@ -3,6 +3,7 @@ import { ConfigValue } from "./configValue";
 import { Controller } from "./controller";
 import { TestRunner } from "./runner";
 import { SourceMapStore } from "./source-map-store";
+import { defaultTestFunctionSpecifiers } from "./test-function-specifier-config";
 
 export async function activate(context: vscode.ExtensionContext) {
   const smStore = new SourceMapStore();
@@ -15,6 +16,8 @@ export async function activate(context: vscode.ExtensionContext) {
     },
   ]);
 
+  const testSpecifiers = new ConfigValue("testSpecifiers", defaultTestFunctionSpecifiers);
+
   const ctrls = new Map<vscode.WorkspaceFolder, Controller>();
   const refreshFolders = () => {
     for (const ctrl of ctrls.values()) {
@@ -48,6 +51,7 @@ export async function activate(context: vscode.ExtensionContext) {
             includePattern.value,
             excludePatterns.value,
             extensions.value,
+            testSpecifiers.value,
           ),
         );
       }

src/parsing.ts

@@ -1,7 +1,21 @@
 import type { Options } from "acorn";
 import { parse } from "acorn-loose";
 import * as evk from "eslint-visitor-keys";
-import { CallExpression, Expression, Node, SourceLocation, Super } from "estree";
+import {
+  CallExpression,
+  Expression,
+  Node,
+  SourceLocation,
+  Super,
+  type ImportDeclaration,
+} from "estree";
+import * as Path from "node:path";
+import * as vscode from "vscode";
+import {
+  defaultTestFunctionSpecifiers,
+  type TestFunctionSpecifierConfig,
+} from "./test-function-specifier-config";
+import { assertUnreachable } from "./utils";
 
 const enum C {
   ImportDeclaration = "ImportDeclaration",
@@ -76,30 +90,94 @@ export interface IParsedNode {
   children: IParsedNode[];
 }
 
-export const parseSource = (text: string) => {
+/**
+ * Look for test function imports in this AST node
+ *
+ * @param folder The workspace folder this file belongs to, used for relative path references to a custom test function
+ * @param fileUri The URI of the file we are extracting from
+ * @param testFunctions the tests function imports to check for
+ * @param node the ImportDelcaration to look for test imports
+ * @returns
+ */
+function importDeclarationExtractTests(
+  folder: vscode.WorkspaceFolder | undefined,
+  fileUri: vscode.Uri | undefined,
+  testFunctions: TestFunctionSpecifierConfig[],
+  node: ImportDeclaration,
+): ExtractTest[] {
+  const idTests: ExtractTest[] = [];
+  if (typeof node.source.value !== "string") {
+    return [];
+  }
+
+  let importValue = node.source.value;
+  if (node.source.value.startsWith("./") || node.source.value.startsWith("../")) {
+    if (!folder || !fileUri) {
+      console.warn(`Trying to match custom test function without specifying a folder or fileUri`);
+      return [];
+    }
+
+    // This is a relative import, we need to adjust the import value for matching purposes
+    const importRelativeToRoot = Path.relative(
+      folder.uri.fsPath,
+      Path.resolve(Path.dirname(fileUri.fsPath), node.source.value),
+    );
+
+    importValue = `./${importRelativeToRoot}`;
+  }
+
+  for (const specifier of testFunctions) {
+    if (specifier.import !== importValue) {
+      continue;
+    }
+
+    // Next check to see if the functions imported are tests functions
+    const validNames = new Set(
+      typeof specifier.name === "string" ? [specifier.name] : specifier.name,
+    );
+
+    for (const spec of node.specifiers) {
+      const specType = spec.type;
+      if (specType === C.ImportDefaultSpecifier || specType === C.ImportNamespaceSpecifier) {
+        if (validNames.has("default")) {
+          idTests.push(matchNamespaced(spec.local.name));
+        }
+      } else if (specType === C.ImportSpecifier) {
+        if (spec.imported.type === C.Identifier) {
+          if (validNames.has(spec.imported.name)) {
+            idTests.push(matchIdentified(spec.imported.name, spec.local.name));
+          }
+        }
+      } else {
+        assertUnreachable(specType, `${specType} was unhandled`);
+      }
+    }
+  }
+
+  return idTests;
+}
+
+export const parseSource = (
+  text: string,
+  folder?: vscode.WorkspaceFolder,
+  fileUri?: vscode.Uri,
+  testFunctions?: TestFunctionSpecifierConfig[],
+): IParsedNode[] => {
   const ast = parse(text, acornOptions);
+  const testMatchers = testFunctions ?? defaultTestFunctionSpecifiers;
 
   const idTests: ExtractTest[] = [];
 
+  // Since tests can be nested inside of each other, for example a test suite.
+  // We keep track of the test declarations in a tree.
   const stack: { node: Node; r: IParsedNode }[] = [];
   stack.push({ node: undefined, r: { children: [] } } as any);
 
   traverse(ast as Node, {
     enter(node) {
-      if (node.type === C.ImportDeclaration && node.source.value === C.NodeTest) {
-        for (const spec of node.specifiers) {
-          switch (spec.type) {
-            case C.ImportNamespaceSpecifier:
-            case C.ImportDefaultSpecifier:
-              idTests.push(matchNamespaced(spec.local.name));
-              break;
-            case C.ImportSpecifier:
-              if (spec.imported.type === C.Identifier) {
-                idTests.push(matchIdentified(spec.imported.name, spec.local.name));
-              }
-              break;
-          }
-        }
+      if (node.type === C.ImportDeclaration) {
+        const matchers = importDeclarationExtractTests(folder, fileUri, testMatchers, node);
+        idTests.push(...matchers);
       } else if (
         node.type === C.VariableDeclarator &&
         node.init?.type === C.CallExpression &&
@@ -137,14 +215,23 @@ export const parseSource = (text: string) => {
               fn,
               name,
             };
+
+            // We have encountered a test function, record it in the tree.
             stack[stack.length - 1].r.children.push(child);
+
+            // This test function is potentially a "parent" for subtests, so
+            // keep it as the "current leaf" of the stack, so future sub-tests
+            // can be associated with it
             stack.push({ node, r: child });
             break;
           }
         }
       }
     },
     leave(node) {
+      // We are exiting a node that was potentially a test function.  If it was,
+      // we need to pop it of the stack, since there are no more subtests to be
+      // associated with it.
       if (stack[stack.length - 1].node === node) {
         stack.pop();
       }

src/test-function-specifier-config.ts

@@ -0,0 +1,20 @@
+/**
+ * A declarative way to target a function call either imported from a package,
+ * like node:test or from another file in the project
+ */
+export interface TestFunctionSpecifierConfig {
+  /** The names of the functions that should be included in the test runner view */
+  name: string[] | string;
+
+  /**
+   * The import location where thoes functions were imported from. If the import
+   * starts with `./` it will be treated as a file import relative to the root
+   * of the workspace, otherwise it refers to a package, like node:test or
+   * vitest
+   */
+  import: string;
+}
+
+export const defaultTestFunctionSpecifiers: TestFunctionSpecifierConfig[] = [
+  { import: "node:test", name: ["default", "it", "test", "describe", "suite"] },
+];

src/utils.ts

@@ -0,0 +1,11 @@
+/**
+ * A utility function to ensure that typescript unions are exhaustively checked.
+ * This function will fail at compile time if a previously exhaustive check is
+ * no longer exhaustive (in case a new value is added)
+ *
+ * @param _ the value that should have already been exhaustivly checked
+ * @param message The error message to throw in case this code is reached during runtime
+ */
+export function assertUnreachable(_: never, message: string): never {
+  throw new Error(message);
+}

testCases/simple/workspace/inRootFolderWrapped.test.mjs

@@ -0,0 +1,6 @@
+import { test } from "./test/utils.js";
+import { strictEqual } from "node:assert";
+
+test("using wrapped test function from the workspace folder", () => {
+  strictEqual(1 + 1, 2);
+});

testCases/simple/workspace/otherFolder/otherFolderWrapped.test.mjs

@@ -0,0 +1,6 @@
+import { test } from "../test/utils.js";
+import { strictEqual } from "node:assert";
+
+test("using wrapped test function from the workspace/otherFolder folder", () => {
+  strictEqual(1 + 1, 2);
+});

testCases/simple/workspace/test/inTestFolderWrapped.mjs

@@ -0,0 +1,6 @@
+import { test } from "./utils.js";
+import { strictEqual } from "node:assert";
+
+test("using wrapped test function from the workspace/test folder", () => {
+  strictEqual(1 + 1, 2);
+});

testCases/simple/workspace/test/utils.js

@@ -0,0 +1,5 @@
+const { test: nodeTest } = require("node:test");
+
+exports.test = function test(name, fn) {
+  return nodeTest(name, fn);
+};