use new vscode testing api

package.json

@@ -8,7 +8,7 @@
   "description": "Vitest Runner for VSCode that actually work",
   "publisher": "rluvaton",
   "engines": {
-    "vscode": "^1.65.0"
+    "vscode": "^1.68.0"
   },
   "categories": [
     "Testing",
@@ -28,7 +28,9 @@
     "onLanguage:javascriptreact"
   ],
   "devDependencies": {
-    "@types/vscode": "^1.65.0",
+    "@types/get-installed-path": "^4.0.1",
+    "@types/vscode": "^1.81.0",
+    "@vscode/dts": "^0.4.0",
     "esbuild": "^0.14.27",
     "prettier": "^2.6.0",
     "typescript": "^4.6.2",
@@ -43,6 +45,7 @@
     "format:write": "yarn format:check --write"
   },
   "dependencies": {
-    "find-up": "^5.0.0"
+    "find-up": "^5.0.0",
+    "get-installed-path": "^4.0.8"
   }
 }

src/index.ts

@@ -1,16 +1,188 @@
-import type * as ts from 'typescript';
 import * as vscode from 'vscode';
-import { CodeLensProvider } from './codelens';
-import { getVscodeTypescriptPath } from './utils';
+import * as typescript from 'typescript';
+import { TestCase, testData, TestFile } from './test-tree/sample';
+import { TestTreeBuilder } from './test-tree/build';
 
-export function activate(context: vscode.ExtensionContext) {
-    const tsPath = getVscodeTypescriptPath(vscode.env.appRoot);
-    const typescript = require(tsPath) as typeof ts;
+export async function activate(context: vscode.ExtensionContext) {
+    const ctrl = vscode.tests.createTestController('vitestTestController', 'Vitest');
+    context.subscriptions.push(ctrl);
+
+    const fileChangedEmitter = new vscode.EventEmitter<vscode.Uri>();
+    const runHandler = (isDebug: boolean, request: vscode.TestRunRequest, cancellation: vscode.CancellationToken) => {
+        if (!request.continuous) {
+            return startTestRun(request, isDebug);
+        }
+
+        const l = fileChangedEmitter.event(uri => startTestRun(
+            new vscode.TestRunRequest(
+                [getOrCreateFile(ctrl, uri).file],
+                undefined,
+                request.profile,
+                true
+            ),
+            isDebug
+        ));
+        cancellation.onCancellationRequested(() => l.dispose());
+    };
+
+    const startTestRun = (request: vscode.TestRunRequest, isDebug: boolean) => {
+        const queue: { test: vscode.TestItem; data: TestCase }[] = [];
+        const run = ctrl.createTestRun(request);
+        // map of file uris to statements on each line:
+
+        const discoverTests = async (tests: Iterable<vscode.TestItem>) => {
+            for (const test of tests) {
+                if (request.exclude?.includes(test)) {
+                    continue;
+                }
+
+                const data = testData.get(test);
+                if (data instanceof TestCase) {
+                    run.enqueued(test);
+                    queue.push({ test, data });
+                } else {
+                    if (data instanceof TestFile && !data.didResolve) {
+                        await data.updateFromDisk(ctrl, test);
+                    }
+
+                    await discoverTests(gatherTestItems(test.children));
+                }
+            }
+        };
+
+        const runTestQueue = async () => {
+            for (const { test, data } of queue) {
+                run.appendOutput(`Running ${test.id}\r\n`);
+                if (run.token.isCancellationRequested) {
+                    run.skipped(test);
+                } else {
+                    run.started(test);
+                    await data.run(test, run, isDebug);
+                }
+
+                run.appendOutput(`Completed ${test.id}\r\n`);
+            }
+
+            run.end();
+        };
+
+        discoverTests(request.include ?? gatherTestItems(ctrl.items)).then(runTestQueue);
+    };
+
+    ctrl.refreshHandler = async () => {
+        await Promise.all(getWorkspaceTestPatterns().map(({ pattern }) => findInitialFiles(ctrl, pattern)));
+    };
+
+    ctrl.createRunProfile('Run Tests', vscode.TestRunProfileKind.Run, (...args) => runHandler(false, ...args), true, undefined, true);
+    ctrl.createRunProfile('Debug Tests', vscode.TestRunProfileKind.Debug, (...args) => runHandler(true, ...args), false, undefined, false);
+
+    ctrl.resolveHandler = async item => {
+        if (!item) {
+            context.subscriptions.push(...startWatchingWorkspace(ctrl, fileChangedEmitter));
+            return;
+        }
+
+        const data = testData.get(item);
+        if (data instanceof TestFile) {
+            await data.updateFromDisk(ctrl, item);
+        }
+    };
+
+    function updateNodeForDocument(e: vscode.TextDocument) {
+        if (e.uri.scheme !== 'file') {
+            return;
+        }
+
+        const ac = new AbortController();
+
+        // TODO - in the future we should check if file match the test pattern in the vitest config
+        TestTreeBuilder.build(typescript, e.getText(), ac.signal, {
+            onSuite() {
+                ac.abort();
+            },
+            onTest() {
+                ac.abort();
+            }
+        });
+
+        // No test file found
+        if(!ac.signal.aborted) {
+            return;
+        }
+
+        const { file, data } = getOrCreateFile(ctrl, e.uri);
+        data.updateFromContents(ctrl, e, e.getText(), file);
+    }
+
+    for (const document of vscode.workspace.textDocuments) {
+        updateNodeForDocument(document);
+    }
 
     context.subscriptions.push(
-        vscode.languages.registerCodeLensProvider(
-            ['typescript', 'javascript', 'typescriptreact', 'javascriptreact'],
-            new CodeLensProvider(typescript)
-        )
+        vscode.workspace.onDidOpenTextDocument(updateNodeForDocument),
+        vscode.workspace.onDidChangeTextDocument(e => updateNodeForDocument(e.document)),
     );
 }
+
+function getOrCreateFile(controller: vscode.TestController, uri: vscode.Uri) {
+    const existing = controller.items.get(uri.toString());
+    if (existing) {
+        return { file: existing, data: testData.get(existing) as TestFile };
+    }
+
+    const file = controller.createTestItem(uri.toString(), uri.path.split('/').pop()!, uri);
+    controller.items.add(file);
+
+    const data = new TestFile();
+    testData.set(file, data);
+
+    file.canResolveChildren = true;
+    return { file, data };
+}
+
+function gatherTestItems(collection: vscode.TestItemCollection) {
+    const items: vscode.TestItem[] = [];
+    collection.forEach(item => items.push(item));
+    return items;
+}
+
+function getWorkspaceTestPatterns() {
+    if (!vscode.workspace.workspaceFolders) {
+        return [];
+    }
+
+    return vscode.workspace.workspaceFolders.map(workspaceFolder => ({
+        workspaceFolder,
+        // TODO - ADD SUPPORT FOR js AND spec AND tsx
+        pattern: new vscode.RelativePattern(workspaceFolder, '**/*.test.ts'),
+    }));
+}
+
+async function findInitialFiles(controller: vscode.TestController, pattern: vscode.GlobPattern) {
+    for (const file of await vscode.workspace.findFiles(pattern)) {
+        getOrCreateFile(controller, file);
+    }
+}
+
+function startWatchingWorkspace(controller: vscode.TestController, fileChangedEmitter: vscode.EventEmitter<vscode.Uri>) {
+    return getWorkspaceTestPatterns().map(({ workspaceFolder, pattern }) => {
+        const watcher = vscode.workspace.createFileSystemWatcher(pattern);
+
+        watcher.onDidCreate(uri => {
+            getOrCreateFile(controller, uri);
+            fileChangedEmitter.fire(uri);
+        });
+        watcher.onDidChange(async uri => {
+            const { file, data } = getOrCreateFile(controller, uri);
+            if (data.didResolve) {
+                await data.updateFromDisk(controller, file);
+            }
+            fileChangedEmitter.fire(uri);
+        });
+        watcher.onDidDelete(uri => controller.items.delete(uri.toString()));
+
+        findInitialFiles(controller, pattern);
+
+        return watcher;
+    });
+}

src/reporter.ts

@@ -0,0 +1,48 @@
+import {Awaitable, Reporter, UserConsoleLog, Vitest} from "vitest";
+import {File, TaskResultPack} from "@vitest/runner";
+
+export default class CustomReporter implements Reporter {
+    onInit(ctx: Vitest): void {
+        console.log('onInit', ctx)
+    }
+
+    onPathsCollected(paths?: string[]) {
+        console.log('onPathsCollected', paths)
+    };
+
+    onCollected(files?: File[]) {
+        console.log('onCollected', files)
+    }
+
+    onFinished(files?: File[], errors?: unknown[]) {
+        console.log('onFinished', files, errors)
+    }
+
+    onTaskUpdate(packs: TaskResultPack[]) {
+        console.log('onTaskUpdate', packs)
+    }
+
+    onTestRemoved(trigger?: string) {
+        console.log('onTestRemoved', trigger)
+    }
+
+    onWatcherStart(files?: File[], errors?: unknown[]) {
+        console.log('onWatcherStart', files, errors)
+    }
+
+    onWatcherRerun(files: string[], trigger?: string) {
+        console.log('onWatcherRerun', files, trigger)
+    }
+
+    onServerRestart(reason?: string) {
+        console.log('onServerRestart', reason)
+    }
+
+    onUserConsoleLog(log: UserConsoleLog) {
+        console.log('onUserConsoleLog', log)
+    }
+
+    onProcessTimeout() {
+        console.log('onProcessTimeout')
+    }
+}