From 5428bdc9d207292f6914d6dd1cd050554c06577a Mon Sep 17 00:00:00 2001
From: Connor Peet <connor@peet.io>
Date: Fri, 9 Jul 2021 16:19:26 -0700
Subject: [PATCH] test coverage apis

---
 test-provider-sample/src/extension.ts | 42 ++++++++++++++++++++++++++-
 test-provider-sample/src/testTree.ts  |  2 +-
 2 files changed, 42 insertions(+), 2 deletions(-)

diff --git a/test-provider-sample/src/extension.ts b/test-provider-sample/src/extension.ts
index 4b95b247..3649f36e 100644
--- a/test-provider-sample/src/extension.ts
+++ b/test-provider-sample/src/extension.ts
@@ -1,5 +1,5 @@
 import * as vscode from 'vscode';
-import { MarkdownTestData, TestCase, testData, TestFile } from './testTree';
+import { getContentFromFilesystem, MarkdownTestData, TestCase, testData, TestFile } from './testTree';
 
 export function activate(context: vscode.ExtensionContext) {
   const ctrl = vscode.test.createTestController('mathTestController');
@@ -14,6 +14,9 @@ export function activate(context: vscode.ExtensionContext) {
   ctrl.runHandler = (request: vscode.TestRunRequest, cancellation: vscode.CancellationToken) => {
     const queue: { test: vscode.TestItem; data: TestCase }[] = [];
     const run = ctrl.createTestRun(request);
+    // map of file uris to statments on each line:
+    const coveredLines = new Map</* file uri */ string, (vscode.StatementCoverage | undefined)[]>();
+
     const discoverTests = async (tests: Iterable<vscode.TestItem>) => {
       for (const test of tests) {
         if (request.exclude?.includes(test)) {
@@ -31,6 +34,20 @@ export function activate(context: vscode.ExtensionContext) {
 
           await discoverTests(test.children.values());
         }
+
+        if (test.uri && !coveredLines.has(test.uri.toString())) {
+          try {
+            const lines = (await getContentFromFilesystem(test.uri)).split('\n');
+            coveredLines.set(
+              test.uri.toString(),
+              lines.map((lineText, lineNo) =>
+                lineText.trim().length ? new vscode.StatementCoverage(0, new vscode.Position(lineNo, 0)) : undefined
+              )
+            );
+          } catch {
+            // ignored
+          }
+        }
       }
     };
 
@@ -43,12 +60,35 @@ export function activate(context: vscode.ExtensionContext) {
           run.setState(test, vscode.TestResultState.Running);
           await data.run(test, run);
         }
+
+        const lineNo = test.range!.start.line;
+        const fileCoverage = coveredLines.get(test.uri!.toString());
+        if (fileCoverage) {
+          fileCoverage[lineNo]!.executionCount++;
+        }
+
         run.appendOutput(`Completed ${test.id}\r\n`);
       }
 
       run.end();
     };
 
+    run.coverageProvider = {
+      provideFileCoverage() {
+        const coverage: vscode.FileCoverage[] = [];
+        for (const [uri, statements] of coveredLines) {
+          coverage.push(
+            vscode.FileCoverage.fromDetails(
+              vscode.Uri.parse(uri),
+              statements.filter((s): s is vscode.StatementCoverage => !!s)
+            )
+          );
+        }
+
+        return coverage;
+      },
+    };
+
     discoverTests(request.tests).then(runTestQueue);
   };
 
diff --git a/test-provider-sample/src/testTree.ts b/test-provider-sample/src/testTree.ts
index 6e7bc2d5..6ab5306e 100644
--- a/test-provider-sample/src/testTree.ts
+++ b/test-provider-sample/src/testTree.ts
@@ -10,7 +10,7 @@ export const testData = new WeakMap<vscode.TestItem, MarkdownTestData>();
 
 let generationCounter = 0;
 
-const getContentFromFilesystem = async (uri: vscode.Uri) => {
+export const getContentFromFilesystem = async (uri: vscode.Uri) => {
   try {
     const rawContent = await vscode.workspace.fs.readFile(uri);
     return textDecoder.decode(rawContent);