From e87f10e6b8a78cde5d2124007ec21f88cb9debee Mon Sep 17 00:00:00 2001
From: Link <linkdoyle@163.com>
Date: Fri, 6 Nov 2020 22:35:07 +0800
Subject: [PATCH] feat. Support custom test case in source code.

---
 src/codelens/CustomCodeLensProvider.ts | 24 +++++++++++++++++++++++-
 src/commands/test.ts                   | 15 ++++++++++++++-
 src/extension.ts                       |  2 +-
 3 files changed, 38 insertions(+), 3 deletions(-)

diff --git a/src/codelens/CustomCodeLensProvider.ts b/src/codelens/CustomCodeLensProvider.ts
index 4b9b6491..d045cc32 100644
--- a/src/codelens/CustomCodeLensProvider.ts
+++ b/src/codelens/CustomCodeLensProvider.ts
@@ -56,10 +56,32 @@ export class CustomCodeLensProvider implements vscode.CodeLensProvider {
         }
 
         if (shortcuts.indexOf("test") >= 0) {
+            let testLineStart: number = -1;
+            for (let i: number = document.lineCount - 1; i >= 0; i--) {
+                const lineContent: string = document.lineAt(i).text;
+                if (lineContent.indexOf("@lc test=start") >= 0) {
+                    testLineStart = i;
+                    break;
+                }
+            }
+            let testLineEnd: number = -1;
+            if (testLineStart >= 0) {
+                for (let i: number = document.lineCount - 1; i >= 0; i--) {
+                    const lineContent: string = document.lineAt(i).text;
+                    if (lineContent.indexOf("@lc test=end") >= 0) {
+                        testLineEnd = i;
+                        break;
+                    }
+                }
+            }
+            const testStringInSource: string | null = testLineEnd >= 0
+                ? document.getText(new vscode.Range(testLineStart + 1, 0, testLineEnd, 0))
+                : null;
+
             codeLens.push(new vscode.CodeLens(range, {
                 title: "Test",
                 command: "leetcode.testSolution",
-                arguments: [document.uri],
+                arguments: [document.uri, testStringInSource],
             }));
         }
 
diff --git a/src/commands/test.ts b/src/commands/test.ts
index d070c9aa..b061e05c 100644
--- a/src/commands/test.ts
+++ b/src/commands/test.ts
@@ -12,7 +12,7 @@ import { getActiveFilePath } from "../utils/workspaceUtils";
 import * as wsl from "../utils/wslUtils";
 import { leetCodeSubmissionProvider } from "../webview/leetCodeSubmissionProvider";
 
-export async function testSolution(uri?: vscode.Uri): Promise<void> {
+export async function testSolution(uri?: vscode.Uri, testContentInSource?: string): Promise<void> {
     try {
         if (leetCodeManager.getStatus() === UserStatus.SignedOut) {
             return;
@@ -30,6 +30,12 @@ export async function testSolution(uri?: vscode.Uri): Promise<void> {
                 detail: "Test with the default cases",
                 value: ":default",
             },
+            {
+                label: "$(code) From source",
+                description: "",
+                detail: "Test with the cases in source code",
+                value: ":source",
+            },
             {
                 label: "$(pencil) Write directly...",
                 description: "",
@@ -75,6 +81,13 @@ export async function testSolution(uri?: vscode.Uri): Promise<void> {
                     }
                 }
                 break;
+            case ":source":
+                if (testContentInSource) {
+                    result = await leetCodeExecutor.testSolution(filePath, parseTestString(testContentInSource.replace(/\r?\n/g, "\\n")));
+                } else {
+                    vscode.window.showErrorMessage("Failed to load test cases from source code.");
+                }
+                break;
             default:
                 break;
         }
diff --git a/src/extension.ts b/src/extension.ts
index 5bd026f1..c39ec7f9 100644
--- a/src/extension.ts
+++ b/src/extension.ts
@@ -59,7 +59,7 @@ export async function activate(context: vscode.ExtensionContext): Promise<void>
             vscode.commands.registerCommand("leetcode.searchProblem", () => show.searchProblem()),
             vscode.commands.registerCommand("leetcode.showSolution", (input: LeetCodeNode | vscode.Uri) => show.showSolution(input)),
             vscode.commands.registerCommand("leetcode.refreshExplorer", () => leetCodeTreeDataProvider.refresh()),
-            vscode.commands.registerCommand("leetcode.testSolution", (uri?: vscode.Uri) => test.testSolution(uri)),
+            vscode.commands.registerCommand("leetcode.testSolution", (uri?: vscode.Uri, testContentInSource?: string) => test.testSolution(uri, testContentInSource)),
             vscode.commands.registerCommand("leetcode.submitSolution", (uri?: vscode.Uri) => submit.submitSolution(uri)),
             vscode.commands.registerCommand("leetcode.switchDefaultLanguage", () => switchDefaultLanguage()),
             vscode.commands.registerCommand("leetcode.addFavorite", (node: LeetCodeNode) => star.addFavorite(node)),