From 8121a2c5a78eab7d87e36a03859f2a99e5ac5326 Mon Sep 17 00:00:00 2001
From: Sheng Chen <sheche@microsoft.com>
Date: Mon, 26 Apr 2021 17:29:28 +0800
Subject: [PATCH] fix: Display more friendly options when generating tests

---
 .../java/test/plugin/model/Option.java        | 28 +++++++++++++++++++
 .../java/test/plugin/util/JUnitPlugin.java    |  6 ++--
 .../test/plugin/util/TestGenerationUtils.java | 19 ++++++++-----
 src/commands/generationCommands.ts            | 27 +++++++++++++++---
 4 files changed, 66 insertions(+), 14 deletions(-)
 create mode 100644 java-extension/com.microsoft.java.test.plugin/src/main/java/com/microsoft/java/test/plugin/model/Option.java

diff --git a/java-extension/com.microsoft.java.test.plugin/src/main/java/com/microsoft/java/test/plugin/model/Option.java b/java-extension/com.microsoft.java.test.plugin/src/main/java/com/microsoft/java/test/plugin/model/Option.java
new file mode 100644
index 00000000..87c23716
--- /dev/null
+++ b/java-extension/com.microsoft.java.test.plugin/src/main/java/com/microsoft/java/test/plugin/model/Option.java
@@ -0,0 +1,28 @@
+/*******************************************************************************
+* Copyright (c) 2021 Microsoft Corporation and others.
+* All rights reserved. This program and the accompanying materials
+* are made available under the terms of the Eclipse Public License v1.0
+* which accompanies this distribution, and is available at
+* http://www.eclipse.org/legal/epl-v10.html
+*
+* Contributors:
+*     Microsoft Corporation - initial API and implementation
+*******************************************************************************/
+
+package com.microsoft.java.test.plugin.model;
+
+public class Option {
+    public String value;
+    public String label;
+    public String description;
+
+    public Option(String label) {
+        this.label = label;
+    }
+
+    public Option(String value, String label, String description) {
+        this.value = value;
+        this.label = label;
+        this.description = description;
+    }
+}
diff --git a/java-extension/com.microsoft.java.test.plugin/src/main/java/com/microsoft/java/test/plugin/util/JUnitPlugin.java b/java-extension/com.microsoft.java.test.plugin/src/main/java/com/microsoft/java/test/plugin/util/JUnitPlugin.java
index 8049f3c0..6966230f 100644
--- a/java-extension/com.microsoft.java.test.plugin/src/main/java/com/microsoft/java/test/plugin/util/JUnitPlugin.java
+++ b/java-extension/com.microsoft.java.test.plugin/src/main/java/com/microsoft/java/test/plugin/util/JUnitPlugin.java
@@ -12,6 +12,7 @@
 package com.microsoft.java.test.plugin.util;
 
 import com.microsoft.java.test.plugin.handler.ClasspathUpdateHandler;
+import com.microsoft.java.test.plugin.model.Option;
 
 import org.eclipse.core.runtime.CoreException;
 import org.eclipse.core.runtime.IStatus;
@@ -79,13 +80,12 @@ public static void logException(String message, Throwable ex) {
         }
     }
 
-    public static Object askClientForChoice(String placeHolder, List<String> choices) {
+    public static Object askClientForChoice(String placeHolder, List<Option> choices) {
         return askClientForChoice(placeHolder, choices, false);
     }
 
-    public static Object askClientForChoice(String placeHolder, List<String> choices, boolean canPickMany) {
+    public static Object askClientForChoice(String placeHolder, List<Option> choices, boolean canPickMany) {
         return JavaLanguageServerPlugin.getInstance().getClientConnection()
                     .executeClientCommand("_java.test.askClientForChoice", placeHolder, choices, canPickMany);
     }
-
 }
diff --git a/java-extension/com.microsoft.java.test.plugin/src/main/java/com/microsoft/java/test/plugin/util/TestGenerationUtils.java b/java-extension/com.microsoft.java.test.plugin/src/main/java/com/microsoft/java/test/plugin/util/TestGenerationUtils.java
index d8fcad0c..6f872107 100644
--- a/java-extension/com.microsoft.java.test.plugin/src/main/java/com/microsoft/java/test/plugin/util/TestGenerationUtils.java
+++ b/java-extension/com.microsoft.java.test.plugin/src/main/java/com/microsoft/java/test/plugin/util/TestGenerationUtils.java
@@ -11,6 +11,7 @@
 
 package com.microsoft.java.test.plugin.util;
 
+import com.microsoft.java.test.plugin.model.Option;
 import com.microsoft.java.test.plugin.model.TestKind;
 import com.microsoft.java.test.plugin.provider.TestKindProvider;
 
@@ -185,10 +186,11 @@ private static WorkspaceEdit generateTestsFromTest(ICompilationUnit unit, Compil
     }
 
     private static List<String> determineMethodsToGenerate(List<String> lifecycleAnnotations) {
-        final List<String> methodList = lifecycleAnnotations.stream()
-                .map(annotation -> "@" + annotation + " Method")
+        final List<Option> methodList = lifecycleAnnotations.stream()
+                .map(annotation -> new Option(annotation, "@" + annotation, 
+                        capitalize(getSuggestedMethodNameByAnnotation(annotation)) + " Method"))
                 .collect(Collectors.toList());
-        methodList.add(0, "@Test Method");
+        methodList.add(0, new Option("Test", "@Test", "Test Method"));
         return (List<String>) JUnitPlugin.askClientForChoice("Select methods to generate",
                 methodList, true /*pickMany*/);
     }
@@ -202,9 +204,9 @@ private static TestKind determineTestFramework(Set<TestKind> availableFrameworks
         if (availableFrameworks.size() == 1) {
             return availableFrameworks.iterator().next();
         } else {
-            final List<String> frameworkList = availableFrameworks.stream()
+            final List<Option> frameworkList = availableFrameworks.stream()
                 .sorted((kind1, kind2) -> kind1.getValue() - kind2.getValue())
-                .map(framework -> framework.toString())
+                .map(framework -> new Option(framework.toString()))
                 .collect(Collectors.toList());
 
             final Object result = JUnitPlugin.askClientForChoice("Select a test framework to use", frameworkList);
@@ -244,8 +246,7 @@ private static TextEdit createTextEditFromTestFile(TestKind kind, List<String> m
         }
 
         final String prefix = annotationPrefix;
-        final List<MethodMetaData> metadata = methodsToGenerate.stream().map(method -> {
-            final String annotationName = method.substring(method.indexOf("@") + 1, method.lastIndexOf(" "));
+        final List<MethodMetaData> metadata = methodsToGenerate.stream().map(annotationName -> {
             final String methodName = getSuggestedMethodNameByAnnotation(annotationName);
             return new MethodMetaData(methodName, prefix + annotationName);
         }).collect(Collectors.toList());
@@ -417,6 +418,10 @@ private static boolean hasMethod(IMethod[] methods, String name) {
         return false;
     }
 
+    private static String capitalize(String str) {
+        return Character.toUpperCase(str.charAt(0)) + str.substring(1);
+    }
+
     static class MethodMetaData {
         public String methodName;
         public String annotation;
diff --git a/src/commands/generationCommands.ts b/src/commands/generationCommands.ts
index dfc678d1..f12588de 100644
--- a/src/commands/generationCommands.ts
+++ b/src/commands/generationCommands.ts
@@ -1,7 +1,7 @@
 // Copyright (c) Microsoft Corporation. All rights reserved.
 // Licensed under the MIT license.
 
-import { commands, ExtensionContext, Uri, window, workspace, WorkspaceEdit } from 'vscode';
+import { commands, ExtensionContext, QuickPickItem, Uri, window, workspace, WorkspaceEdit } from 'vscode';
 import * as protocolConverter from 'vscode-languageclient/lib/protocolConverter';
 import * as commandUtils from '../utils/commandUtils';
 
@@ -14,11 +14,30 @@ export async function generateTests(uri: Uri, cursorOffset: number): Promise<voi
 }
 
 export async function registerSelectTestFrameworkCommand(context: ExtensionContext): Promise<void> {
-    context.subscriptions.push(commands.registerCommand('_java.test.askClientForChoice', async (placeHolder: string, choices: string[], canPickMany: boolean) => {
-        const choice: string | undefined = await window.showQuickPick(choices, {
+    context.subscriptions.push(commands.registerCommand('_java.test.askClientForChoice', async (placeHolder: string, choices: IOption[], canPickMany: boolean) => {
+        const options: IOption[] = [];
+        for (const c of choices) {
+            options.push({
+                label: c.label,
+                description: c.description,
+                value: c.value,
+            });
+        }
+        const ans: any = await window.showQuickPick(options, {
             placeHolder,
             canPickMany,
         });
-        return choice;
+
+        if (!ans) {
+            return undefined;
+        } else if (Array.isArray(ans)) {
+            return ans.map((a: IOption) => a.value || a.label);
+        }
+
+        return ans.value || ans.label;
     }));
 }
+
+interface IOption extends QuickPickItem {
+    value: string;
+}