From adffccbe0fa16e58622212265f469b249cde9dda Mon Sep 17 00:00:00 2001
From: Sheng Chen <sheche@microsoft.com>
Date: Thu, 6 May 2021 11:01:59 +0800
Subject: [PATCH 1/4] feat: Support generate tests from a source file

---
 .../META-INF/MANIFEST.MF                      |   1 +
 .../java/test/plugin/util/JUnitPlugin.java    |  16 +-
 .../test/plugin/util/TestGenerationUtils.java | 291 +++++++++++++++++-
 src/commands/generationCommands.ts            |  61 +++-
 src/extension.ts                              |   8 +-
 src/provider/codeActionProvider.ts            |   8 +-
 src/utils/commandUtils.ts                     |   4 +-
 7 files changed, 369 insertions(+), 20 deletions(-)

diff --git a/java-extension/com.microsoft.java.test.plugin/META-INF/MANIFEST.MF b/java-extension/com.microsoft.java.test.plugin/META-INF/MANIFEST.MF
index 185d4989..994fe38c 100644
--- a/java-extension/com.microsoft.java.test.plugin/META-INF/MANIFEST.MF
+++ b/java-extension/com.microsoft.java.test.plugin/META-INF/MANIFEST.MF
@@ -17,6 +17,7 @@ Require-Bundle: org.eclipse.jdt.core,
  org.eclipse.core.resources,
  org.eclipse.jdt.core.manipulation,
  org.eclipse.text,
+ org.eclipse.ltk.core.refactoring,
  org.eclipse.debug.core,
  org.eclipse.jdt.junit.core,
  org.eclipse.jdt.junit.runtime,
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 6966230f..2d753ab5 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
@@ -26,6 +26,15 @@
 
 public class JUnitPlugin implements BundleActivator {
 
+    /**
+     * Client registered command to handle the server-side request to ask user for input
+     */
+    private static final String JAVA_TEST_ASK_CLIENT_FOR_INPUT = "_java.test.askClientForInput";
+
+    /**
+     * Client registered command to handle the server-side request to ask user for a choice
+     */
+    private static final String JAVA_TEST_ASK_CLIENT_FOR_CHOICE = "_java.test.askClientForChoice";
     public static final String PLUGIN_ID = "java.test.runner";
     private static ClasspathUpdateHandler handler = new ClasspathUpdateHandler();
     private static BundleContext context;
@@ -86,6 +95,11 @@ public static Object askClientForChoice(String placeHolder, List<Option> choices
 
     public static Object askClientForChoice(String placeHolder, List<Option> choices, boolean canPickMany) {
         return JavaLanguageServerPlugin.getInstance().getClientConnection()
-                    .executeClientCommand("_java.test.askClientForChoice", placeHolder, choices, canPickMany);
+                    .executeClientCommand(JAVA_TEST_ASK_CLIENT_FOR_CHOICE, placeHolder, choices, canPickMany);
+    }
+
+    public static Object askClientForInput(String prompt, String defaultValue) {
+        return JavaLanguageServerPlugin.getInstance().getClientConnection()
+                    .executeClientCommand(JAVA_TEST_ASK_CLIENT_FOR_INPUT, prompt, defaultValue);
     }
 }
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 adada903..c44481f0 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
@@ -17,12 +17,15 @@
 
 import org.eclipse.core.runtime.CoreException;
 import org.eclipse.core.runtime.IProgressMonitor;
+import org.eclipse.core.runtime.NullProgressMonitor;
 import org.eclipse.jdt.core.IClasspathEntry;
 import org.eclipse.jdt.core.ICompilationUnit;
 import org.eclipse.jdt.core.IImportDeclaration;
 import org.eclipse.jdt.core.IJavaElement;
 import org.eclipse.jdt.core.IJavaProject;
 import org.eclipse.jdt.core.IMethod;
+import org.eclipse.jdt.core.IPackageFragment;
+import org.eclipse.jdt.core.IPackageFragmentRoot;
 import org.eclipse.jdt.core.IType;
 import org.eclipse.jdt.core.JavaModelException;
 import org.eclipse.jdt.core.dom.AST;
@@ -43,23 +46,31 @@
 import org.eclipse.jdt.core.dom.rewrite.ImportRewrite;
 import org.eclipse.jdt.core.dom.rewrite.ImportRewrite.ImportRewriteContext;
 import org.eclipse.jdt.core.dom.rewrite.ListRewrite;
+import org.eclipse.jdt.core.formatter.CodeFormatter;
+import org.eclipse.jdt.core.refactoring.CompilationUnitChange;
 import org.eclipse.jdt.internal.core.manipulation.StubUtility;
 import org.eclipse.jdt.internal.corext.codemanipulation.CodeGenerationSettings;
 import org.eclipse.jdt.internal.corext.codemanipulation.ContextSensitiveImportRewriteContext;
 import org.eclipse.jdt.internal.corext.codemanipulation.StubUtility2Core;
 import org.eclipse.jdt.internal.corext.dom.ASTNodeFactory;
+import org.eclipse.jdt.internal.corext.refactoring.changes.CreateCompilationUnitChange;
+import org.eclipse.jdt.internal.corext.util.CodeFormatterUtil;
+import org.eclipse.jdt.ls.core.internal.ChangeUtil;
 import org.eclipse.jdt.ls.core.internal.JDTUtils;
 import org.eclipse.jdt.ls.core.internal.JavaLanguageServerPlugin;
 import org.eclipse.jdt.ls.core.internal.handlers.CodeGenerationUtils;
 import org.eclipse.jdt.ls.core.internal.text.correction.SourceAssistProcessor;
 import org.eclipse.lsp4j.MessageType;
 import org.eclipse.lsp4j.WorkspaceEdit;
+import org.eclipse.text.edits.InsertEdit;
 import org.eclipse.text.edits.MalformedTreeException;
 import org.eclipse.text.edits.MultiTextEdit;
 import org.eclipse.text.edits.TextEdit;
 import java.util.ArrayList;
+import java.util.Arrays;
 import java.util.HashMap;
 import java.util.HashSet;
+import java.util.LinkedList;
 import java.util.List;
 import java.util.Map;
 import java.util.Objects;
@@ -152,7 +163,7 @@ public static WorkspaceEdit generateTests(List<Object> arguments, IProgressMonit
                 if (ProjectTestUtils.isTestEntry(entry)) {
                     return generateTestsFromTest(unit, root, (TypeDeclaration) coveringNode, binding, cursorOffset);
                 } else {
-                    generateTestsFromSource();
+                    return generateTestsFromSource(unit, binding, cursorOffset);
                 }
             }
         }
@@ -160,8 +171,272 @@ public static WorkspaceEdit generateTests(List<Object> arguments, IProgressMonit
         return null;
     }
 
-    private static void generateTestsFromSource() {
-        // TODO: unimplement
+    /**
+     * Generate test methods from a focal source file
+     */
+    private static WorkspaceEdit generateTestsFromSource(ICompilationUnit unit, ITypeBinding typeBinding,
+            int cursorOffset) throws CoreException {
+        final IJavaProject javaProject = unit.getJavaProject();
+        final List<TestKind> testFrameworksInProject = TestKindProvider.getTestKindsFromCache(javaProject);
+        final TestKind testKind = determineTestFramework(new HashSet<>(testFrameworksInProject));
+        if (testKind == null) {
+            return null;
+        }
+
+        final IClasspathEntry testEntry = getTestClasspathEntry(unit);
+        final String testFullyQualifiedName = getTestFullyQualifiedName(typeBinding, javaProject, testEntry);
+        if (testFullyQualifiedName == null) {
+            return null;
+        }
+        final List<String> methodsToTest = getMethodsToTest(typeBinding);
+        if (methodsToTest == null || methodsToTest.size() == 0) {
+            return null;
+        }
+
+        final ICompilationUnit testUnit = getTestCompilationUnit(javaProject, testEntry, testFullyQualifiedName);
+        return createTextEditFromSourceFile(testKind, testUnit, methodsToTest, cursorOffset);
+    }
+
+    private static WorkspaceEdit createTextEditFromSourceFile(TestKind kind, ICompilationUnit testUnit,
+            List<String> methodsToTest, int cursorPosition) throws CoreException {
+        if (testUnit.exists()) {
+            final IType[] types = testUnit.getAllTypes();
+            if (types.length == 0) {
+                // an empty java file
+                return getTestCompilationUnit(kind, testUnit, methodsToTest);
+            }
+
+            final CompilationUnit root = (CompilationUnit) TestSearchUtils.parseToAst(testUnit,
+                false /* fromCache */, new NullProgressMonitor());
+
+            if (root == null) {
+                return null;
+            }
+
+            IType type = testUnit.findPrimaryType();
+            if (type == null) {
+                type = types[0];
+            }
+
+            final ASTNode typeNode = root.findDeclaringNode(type.getKey());
+            if (!(typeNode instanceof TypeDeclaration)) {
+                return null;
+            }
+
+            final ITypeBinding binding = ((TypeDeclaration) typeNode).resolveBinding();
+            if (binding == null) {
+                return null;
+            }
+
+            IJavaElement insertPosition = null; // Insert to the last by default.
+            try {
+                insertPosition = CodeGenerationUtils.findInsertElement(type, cursorPosition);
+            } catch (Throwable e) {
+                // ignore if the upstream does not support insert position preference.
+            }
+
+            return createEditToExistingTestFile(root, kind, methodsToTest, (TypeDeclaration) typeNode,
+                    binding, insertPosition);
+        } else {
+            return createAndGetNewTestCompilationUnit(kind, testUnit, methodsToTest);
+        }
+    }
+
+    /**
+     * Get a compilation unit change without creating.
+     */
+    private static WorkspaceEdit getTestCompilationUnit(TestKind kind, ICompilationUnit testUnit,
+            List<String> methodsToTest) throws CoreException {
+        final CompilationUnitChange cuChange = new CompilationUnitChange("", testUnit);
+        final String cuContent = constructNewCU(testUnit, methodsToTest, kind);
+        cuChange.setEdit(new InsertEdit(0, cuContent));
+        return ChangeUtil.convertToWorkspaceEdit(cuChange);
+    }
+
+    /**
+     * Create and get a compilation unit change.
+     */
+    private static WorkspaceEdit createAndGetNewTestCompilationUnit(TestKind kind, ICompilationUnit testUnit,
+            List<String> methodsToTest) throws CoreException {
+        final String cuContent = constructNewCU(testUnit, methodsToTest, kind);
+        final CreateCompilationUnitChange change =
+                new CreateCompilationUnitChange(testUnit, cuContent, "");
+        return ChangeUtil.convertToWorkspaceEdit(change);
+    }
+
+    private static WorkspaceEdit createEditToExistingTestFile(CompilationUnit testRoot, TestKind kind,
+            List<String> methodsToTest, TypeDeclaration typeNode, ITypeBinding typeBinding,
+            IJavaElement insertPosition) throws JavaModelException, CoreException {
+        final String testAnnotation = getTestAnnotation(kind);
+        final List<MethodMetaData> metadata = methodsToTest.stream().map(method -> {
+            final String methodName = getTestMethodName(method);
+            return new MethodMetaData(methodName, testAnnotation);
+        }).collect(Collectors.toList());
+
+        final TextEdit edit = getTextEdit(kind, metadata, testRoot, typeNode, typeBinding, insertPosition);
+        return SourceAssistProcessor.convertToWorkspaceEdit((ICompilationUnit) testRoot.getJavaElement(), edit);
+    }
+
+    private static String constructNewCU(ICompilationUnit testUnit,
+            List<String> methods, TestKind testKind) throws CoreException {
+        final String delimiter = StubUtility.getLineDelimiterUsed(testUnit);
+        final String typeStub = constructTypeStub(testUnit, methods, testKind, delimiter);
+        final String cuContent = constructCUContent(testUnit, testKind, typeStub, delimiter);
+        final String formattedCuStub = CodeFormatterUtil.format(
+                CodeFormatter.K_COMPILATION_UNIT, cuContent, 0, delimiter, testUnit.getJavaProject().getOptions(true));
+        return formattedCuStub;
+    }
+
+    private static String constructCUContent(ICompilationUnit testUnit, TestKind testKind, 
+            String typeContent, String delimiter) throws CoreException {
+        final IPackageFragment packageFragment = (IPackageFragment) testUnit.getParent();
+        final StringBuilder buf = new StringBuilder();
+        if (!packageFragment.isDefaultPackage()) {
+            buf.append("package ")
+                .append(packageFragment.getElementName())
+                .append(";")
+                .append(delimiter)
+                .append(delimiter)
+                .append("import ")
+                .append(getTestAnnotation(testKind))
+                .append(";")
+                .append(delimiter)
+                .append(delimiter);
+        }
+        buf.append(typeContent);
+        return buf.toString();
+    }
+
+    private static String constructTypeStub(ICompilationUnit testUnit, List<String> methods,
+            TestKind testKind, String delimiter) throws CoreException {
+        final String typeName = testUnit.getElementName().replace(".java", "");
+        final StringBuilder buf = new StringBuilder();
+        buf.append("public class ").append(typeName).append(" {").append(delimiter);
+        for (final String method : methods) {
+            buf.append(constructMethodStub(testUnit, testKind, method, delimiter)).append(delimiter);
+        }
+        buf.append("}").append(delimiter);
+        return buf.toString();
+    }
+
+    private static String constructMethodStub(ICompilationUnit testUnit, TestKind testKind,
+            String method, String delimiter) {
+        final StringBuilder buf = new StringBuilder();
+        buf.append("@Test").append(delimiter);
+        if (testKind == TestKind.JUnit) {
+            buf.append("public ");
+        }
+        final String methodName = getTestMethodName(method);
+        buf.append("void ").append(methodName).append("() {").append(delimiter).append(delimiter).append("}");
+        final String methodContent = buf.toString();
+        // TODO: get test unit options directly
+        return CodeFormatterUtil.format(CodeFormatter.K_STATEMENTS, methodContent, 1,
+                delimiter, testUnit.getJavaProject().getOptions(true));
+    }
+
+    private static IClasspathEntry getTestClasspathEntry(ICompilationUnit unit) throws JavaModelException {
+        final IJavaProject javaProject = unit.getJavaProject();
+        // In most cases, this is the classpath entry used for testing, we first find the target entry by hard-code
+        // to avoid go into the generated entries.
+        final IClasspathEntry testEntry = javaProject.getClasspathEntryFor(
+            javaProject.getPath().append("src/test/java"));
+        if (ProjectTestUtils.isTestEntry(testEntry)) {
+            return testEntry;
+        }
+
+        final IClasspathEntry[] entries = javaProject.readRawClasspath();
+        for (final IClasspathEntry entry : entries) {
+            if (ProjectTestUtils.isTestEntry(entry)) {
+                return entry;
+            }
+        }
+
+        return entries[0];
+    }
+
+    private static String getTestFullyQualifiedName(ITypeBinding typeBinding, IJavaProject project,
+            IClasspathEntry testEntry) throws JavaModelException {
+        String promptName = typeBinding.getBinaryName() + "Tests";
+        final ICompilationUnit testCompilationUnit = getTestCompilationUnit(project, testEntry, promptName);
+        if (!testCompilationUnit.exists()) {
+            promptName = typeBinding.getBinaryName() + "Test";
+        }
+
+        final String fullyQualifiedName = (String) JUnitPlugin.askClientForInput(
+            "Please type the target test class name", promptName);
+        if (fullyQualifiedName == null) {
+            return null;
+        }
+
+        if (fullyQualifiedName.charAt(fullyQualifiedName.length() - 1) == '.') {
+            JavaLanguageServerPlugin.getInstance().getClientConnection().showNotificationMessage(MessageType.Error,
+                    "Invalid Java class name: " + fullyQualifiedName);
+            return null;
+        }
+
+        final String[] identifiers = fullyQualifiedName.split("\\.");
+        for (final String identifier : identifiers) {
+            final char[] chars = identifier.toCharArray();
+            if (!Character.isJavaIdentifierStart(chars[0])) {
+                JavaLanguageServerPlugin.getInstance().getClientConnection().showNotificationMessage(MessageType.Error,
+                    "Invalid Java identifier: " + identifier);
+                return null;
+            }
+            for (int i = 1; i < chars.length; i++) {
+                if (!Character.isJavaIdentifierPart(chars[i])) {
+                    JavaLanguageServerPlugin.getInstance().getClientConnection().showNotificationMessage(
+                        MessageType.Error, "Invalid Java identifier: " + identifier);
+                    return null;
+                }
+            }
+        }
+
+        return fullyQualifiedName;
+    }
+
+    private static ICompilationUnit getTestCompilationUnit(IJavaProject javaProject, IClasspathEntry testEntry,
+            String testFullyQualifiedName) throws JavaModelException {
+        final IPackageFragmentRoot packageRoot = javaProject.findPackageFragmentRoot(testEntry.getPath());
+        final String packageQualifiedName = testFullyQualifiedName.substring(0,
+                testFullyQualifiedName.lastIndexOf("."));
+        final IPackageFragment packageFragment = packageRoot.getPackageFragment(packageQualifiedName);
+        final String compilationUnitName = testFullyQualifiedName.substring(
+                testFullyQualifiedName.lastIndexOf(".") + 1) + ".java";
+        final ICompilationUnit testUnit = packageFragment.getCompilationUnit(compilationUnitName);
+        return testUnit;
+    }
+
+    private static List<String> getMethodsToTest(ITypeBinding typeBinding) {
+        // TODO: list all the accessible methods via advanced option button
+        final List<IMethodBinding> allMethods = new LinkedList<>();
+        final IMethodBinding[] typeMethods = typeBinding.getDeclaredMethods();
+        for (final IMethodBinding method : typeMethods) {
+            final int modifiers = method.getModifiers();
+            if (!method.isConstructor() && !Modifier.isPrivate(modifiers) && !method.isSynthetic()) {
+                allMethods.add(method);
+            }
+        }
+
+        final List<Option> methodNames = allMethods.stream()
+            .map(method -> {
+                final String returnValue = method.getReturnType().getName();
+                final ITypeBinding[] paramTypes = method.getParameterTypes();
+                final String params = String.join(", ",
+                        Arrays.stream(paramTypes).map(t -> t.getName()).toArray(String[]::new));
+
+                return new Option(method.getName(), method.getName() + "(" + params + ")", ": " + returnValue);
+            })
+            .sorted((methodA, methodB) -> {
+                return methodA.label.compareTo(methodB.label);
+            })
+            .collect(Collectors.toList());
+
+        return (List<String>) JUnitPlugin.askClientForChoice("Select the methods to test",
+                methodNames, true /*pickMany*/);
+    }
+
+    private static String getTestMethodName(String methodName) {
+        return "test" + Character.toUpperCase(methodName.charAt(0)) + methodName.substring(1);
     }
 
     private static WorkspaceEdit generateTestsFromTest(ICompilationUnit unit, CompilationUnit root,
@@ -495,6 +770,16 @@ private static boolean needsOverrideAnnotation(boolean isStatic, IMethodBinding
         return true;
     }
 
+    private static String getTestAnnotation(TestKind testKind) {
+        if (testKind == TestKind.JUnit) {
+            return JUNIT4_LIFECYCLE_ANNOTATION_PREFIX + "Test";
+        } else if (testKind == TestKind.TestNG) {
+            return TESTNG_LIFECYCLE_ANNOTATION_PREFIX + "Test";
+        }
+
+        return JUNIT5_LIFECYCLE_ANNOTATION_PREFIX + "Test";
+    }
+
     private static String getUniqueMethodName(IJavaElement type, Map<String, IMethodBinding> methodsMap,
             String suggestedName, boolean isStatic) throws JavaModelException {
         if (type instanceof IType) {
diff --git a/src/commands/generationCommands.ts b/src/commands/generationCommands.ts
index f12588de..bcc883a7 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, QuickPickItem, Uri, window, workspace, WorkspaceEdit } from 'vscode';
+import { commands, ExtensionContext, QuickPickItem, TextEdit, Uri, window, workspace, WorkspaceEdit } from 'vscode';
 import * as protocolConverter from 'vscode-languageclient/lib/protocolConverter';
 import * as commandUtils from '../utils/commandUtils';
 
@@ -9,11 +9,17 @@ const converter: protocolConverter.Converter = protocolConverter.createConverter
 export async function generateTests(uri: Uri, cursorOffset: number): Promise<void> {
     const edit: WorkspaceEdit = converter.asWorkspaceEdit(await commandUtils.generateTests(uri, cursorOffset));
     if (edit) {
-        workspace.applyEdit(edit);
+        await workspace.applyEdit(edit);
+        const entries: Array<[Uri, TextEdit[]]> = edit.entries();
+        if (entries?.[0]?.[0]) {
+            await window.showTextDocument(entries[0][0], {
+                preserveFocus: true,
+            });
+        }
     }
 }
 
-export async function registerSelectTestFrameworkCommand(context: ExtensionContext): Promise<void> {
+export async function registerAskForChoiceCommand(context: ExtensionContext): Promise<void> {
     context.subscriptions.push(commands.registerCommand('_java.test.askClientForChoice', async (placeHolder: string, choices: IOption[], canPickMany: boolean) => {
         const options: IOption[] = [];
         for (const c of choices) {
@@ -38,6 +44,55 @@ export async function registerSelectTestFrameworkCommand(context: ExtensionConte
     }));
 }
 
+/**
+ * A command that server side can call to get a input value.
+ * Currently it's only used to get a Java qualified name, so the check is on the client side.
+ * @param context
+ */
+export async function registerAskForInputCommand(context: ExtensionContext): Promise<void> {
+    context.subscriptions.push(commands.registerCommand('_java.test.askClientForInput', async (prompt: string, value: string) => {
+        const ans: string | undefined = await window.showInputBox({
+            value,
+            prompt,
+            validateInput: checkJavaQualifiedName,
+        });
+        return ans;
+    }));
+}
+
+function checkJavaQualifiedName(value: string): string {
+    if (!value || !value.trim()) {
+        return 'Input cannot be empty.';
+    }
+
+    for (const part of value.split('.')) {
+        if (isKeyword(part)) {
+            return `Keyword '${part}' cannot be used.`;
+        }
+
+        if (!isJavaIdentifier(part)) {
+            return `Invalid Java qualified name.`;
+        }
+    }
+
+    return '';
+}
+
+const keywords: Set<string> = new Set([
+    'abstract', 'default', 'if', 'private', 'this', 'boolean', 'do', 'implements', 'protected', 'throw', 'break', 'double', 'import',
+    'public', 'throws', 'byte', 'else', 'instanceof', 'return', 'transient', 'case', 'extends', 'int', 'short', 'try', 'catch', 'final',
+    'interface', 'static', 'void', 'char', 'finally', 'long', 'strictfp', 'volatile', 'class', 'float', 'native', 'super', 'while',
+    'const', 'for', 'new', 'switch', 'continue', 'goto', 'package', 'synchronized', 'true', 'false', 'null', 'assert', 'enum',
+]);
+export function isKeyword(identifier: string): boolean {
+    return keywords.has(identifier);
+}
+
+const identifierRegExp: RegExp = /^([a-zA-Z_$][a-zA-Z\d_$]*)$/;
+export function isJavaIdentifier(identifier: string): boolean {
+    return identifierRegExp.test(identifier);
+}
+
 interface IOption extends QuickPickItem {
     value: string;
 }
diff --git a/src/extension.ts b/src/extension.ts
index eb238e7e..c7fa493f 100644
--- a/src/extension.ts
+++ b/src/extension.ts
@@ -10,7 +10,7 @@ import { sendInfo } from 'vscode-extension-telemetry-wrapper';
 import { testSourceProvider } from '../extension.bundle';
 import { testCodeLensController } from './codelens/TestCodeLensController';
 import { debugTestsFromExplorer, openTextDocument, runTestsFromExplorer, runTestsFromJavaProjectExplorer } from './commands/explorerCommands';
-import { generateTests, registerSelectTestFrameworkCommand } from './commands/generationCommands';
+import { generateTests, registerAskForChoiceCommand, registerAskForInputCommand } from './commands/generationCommands';
 import { openLogFile, showOutputChannel } from './commands/logCommands';
 import { runFromCodeLens } from './commands/runFromCodeLens';
 import { executeTestsFromUri } from './commands/runFromUri';
@@ -66,7 +66,6 @@ async function doActivate(_operationId: string, context: ExtensionContext): Prom
                 await testSourceProvider.initialize();
                 await testFileWatcher.registerListeners(true /*enableDebounce*/);
                 await testCodeLensController.registerCodeLensProvider();
-                await registerTestCodeActionProvider();
             }));
         }
 
@@ -82,7 +81,6 @@ async function doActivate(_operationId: string, context: ExtensionContext): Prom
                     await testSourceProvider.initialize();
                     await testFileWatcher.registerListeners();
                     await testCodeLensController.registerCodeLensProvider();
-                    await registerTestCodeActionProvider();
                 }
                 serverMode = mode;
             }));
@@ -94,7 +92,6 @@ async function doActivate(_operationId: string, context: ExtensionContext): Prom
                 await testSourceProvider.initialize();
                 await testFileWatcher.registerListeners(true /*enableDebounce*/);
                 await testCodeLensController.registerCodeLensProvider();
-                await registerTestCodeActionProvider();
             }));
         }
 
@@ -113,7 +110,8 @@ async function doActivate(_operationId: string, context: ExtensionContext): Prom
     const testTreeView: TreeView<ITestItem> = window.createTreeView(testExplorer.testExplorerViewId, { treeDataProvider: testExplorer, showCollapseAll: true });
     runnerScheduler.initialize(context);
     testReportProvider.initialize(context, javaLanguageSupportVersion);
-    registerSelectTestFrameworkCommand(context);
+    registerAskForChoiceCommand(context);
+    registerAskForInputCommand(context);
 
     context.subscriptions.push(
         testExplorer,
diff --git a/src/provider/codeActionProvider.ts b/src/provider/codeActionProvider.ts
index eb43546f..ce817f66 100644
--- a/src/provider/codeActionProvider.ts
+++ b/src/provider/codeActionProvider.ts
@@ -1,9 +1,7 @@
 // Copyright (c) Microsoft Corporation. All rights reserved.
 // Licensed under the MIT license.
 
-import { CancellationToken, CodeAction, CodeActionContext, CodeActionKind, CodeActionProvider, Disposable, DocumentSelector, languages, Range, RelativePattern, Selection, TextDocument } from 'vscode';
-import { testSourceProvider } from '../../extension.bundle';
-import { parseDocumentSelector } from '../utils/uiUtils';
+import { CancellationToken, CodeAction, CodeActionContext, CodeActionKind, CodeActionProvider, Disposable, languages, Range, Selection, TextDocument } from 'vscode';
 
 let provider: Disposable;
 export class TestCodeActionProvider implements CodeActionProvider {
@@ -34,10 +32,8 @@ export async function registerTestCodeActionProvider(): Promise<Disposable> {
     if (provider) {
         provider.dispose();
     }
-    const patterns: RelativePattern[] = await testSourceProvider.getTestSourcePattern(false /*containsGeneralProject*/);
-    const documentSelector: DocumentSelector = parseDocumentSelector(patterns);
     provider = languages.registerCodeActionsProvider(
-        documentSelector,
+        {language: 'java', scheme: 'file', pattern: '**/*.java'},
         new TestCodeActionProvider(),
         { providedCodeActionKinds: [CodeActionKind.Source.append('generate.tests')] },
     );
diff --git a/src/utils/commandUtils.ts b/src/utils/commandUtils.ts
index 16c211a9..6d3561c3 100644
--- a/src/utils/commandUtils.ts
+++ b/src/utils/commandUtils.ts
@@ -42,8 +42,8 @@ export async function resolveStackTraceLocation(trace: string, projectNames: str
         JavaLanguageServerCommands.RESOLVE_STACKTRACE_LOCATION, trace, projectNames) || '';
 }
 
-export async function generateTests(uri: Uri, startPosition: number): Promise<any> {
-    return await executeJavaLanguageServerCommand<any>(JavaTestRunnerDelegateCommands.GENERATE_TESTS, uri.toString(), startPosition);
+export async function generateTests(uri: Uri, cursorOffset: number): Promise<any> {
+    return await executeJavaLanguageServerCommand<any>(JavaTestRunnerDelegateCommands.GENERATE_TESTS, uri.toString(), cursorOffset);
 }
 
 export async function resolveJUnitLaunchArguments(uri: string, fullName: string, testName: string, project: string,

From e3d2f07ce30fed893433fa8c167b31bb362f4fb1 Mon Sep 17 00:00:00 2001
From: Sheng Chen <sheche@microsoft.com>
Date: Fri, 7 May 2021 13:58:11 +0800
Subject: [PATCH 2/4] Address comments

---
 .../test/plugin/util/TestGenerationUtils.java | 26 +++++++++----------
 src/commands/generationCommands.ts            | 15 ++++++++---
 2 files changed, 24 insertions(+), 17 deletions(-)

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 c44481f0..eb1316d4 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
@@ -288,50 +288,50 @@ private static String constructNewCU(ICompilationUnit testUnit,
     }
 
     private static String constructCUContent(ICompilationUnit testUnit, TestKind testKind, 
-            String typeContent, String delimiter) throws CoreException {
+            String typeContent, String lineDelimiter) throws CoreException {
         final IPackageFragment packageFragment = (IPackageFragment) testUnit.getParent();
         final StringBuilder buf = new StringBuilder();
         if (!packageFragment.isDefaultPackage()) {
             buf.append("package ")
                 .append(packageFragment.getElementName())
                 .append(";")
-                .append(delimiter)
-                .append(delimiter)
+                .append(lineDelimiter)
+                .append(lineDelimiter)
                 .append("import ")
                 .append(getTestAnnotation(testKind))
                 .append(";")
-                .append(delimiter)
-                .append(delimiter);
+                .append(lineDelimiter)
+                .append(lineDelimiter);
         }
         buf.append(typeContent);
         return buf.toString();
     }
 
     private static String constructTypeStub(ICompilationUnit testUnit, List<String> methods,
-            TestKind testKind, String delimiter) throws CoreException {
+            TestKind testKind, String lineDelimiter) throws CoreException {
         final String typeName = testUnit.getElementName().replace(".java", "");
         final StringBuilder buf = new StringBuilder();
-        buf.append("public class ").append(typeName).append(" {").append(delimiter);
+        buf.append("public class ").append(typeName).append(" {").append(lineDelimiter);
         for (final String method : methods) {
-            buf.append(constructMethodStub(testUnit, testKind, method, delimiter)).append(delimiter);
+            buf.append(constructMethodStub(testUnit, testKind, method, lineDelimiter)).append(lineDelimiter);
         }
-        buf.append("}").append(delimiter);
+        buf.append("}").append(lineDelimiter);
         return buf.toString();
     }
 
     private static String constructMethodStub(ICompilationUnit testUnit, TestKind testKind,
-            String method, String delimiter) {
+            String method, String lineDelimiter) {
         final StringBuilder buf = new StringBuilder();
-        buf.append("@Test").append(delimiter);
+        buf.append("@Test").append(lineDelimiter);
         if (testKind == TestKind.JUnit) {
             buf.append("public ");
         }
         final String methodName = getTestMethodName(method);
-        buf.append("void ").append(methodName).append("() {").append(delimiter).append(delimiter).append("}");
+        buf.append("void ").append(methodName).append("() {").append(lineDelimiter).append(lineDelimiter).append("}");
         final String methodContent = buf.toString();
         // TODO: get test unit options directly
         return CodeFormatterUtil.format(CodeFormatter.K_STATEMENTS, methodContent, 1,
-                delimiter, testUnit.getJavaProject().getOptions(true));
+                lineDelimiter, testUnit.getJavaProject().getOptions(true));
     }
 
     private static IClasspathEntry getTestClasspathEntry(ICompilationUnit unit) throws JavaModelException {
diff --git a/src/commands/generationCommands.ts b/src/commands/generationCommands.ts
index bcc883a7..7d712a86 100644
--- a/src/commands/generationCommands.ts
+++ b/src/commands/generationCommands.ts
@@ -78,11 +78,18 @@ function checkJavaQualifiedName(value: string): string {
     return '';
 }
 
+// Copied from https://docs.oracle.com/javase/specs/jls/se8/html/jls-3.html#jls-Keyword
 const keywords: Set<string> = new Set([
-    'abstract', 'default', 'if', 'private', 'this', 'boolean', 'do', 'implements', 'protected', 'throw', 'break', 'double', 'import',
-    'public', 'throws', 'byte', 'else', 'instanceof', 'return', 'transient', 'case', 'extends', 'int', 'short', 'try', 'catch', 'final',
-    'interface', 'static', 'void', 'char', 'finally', 'long', 'strictfp', 'volatile', 'class', 'float', 'native', 'super', 'while',
-    'const', 'for', 'new', 'switch', 'continue', 'goto', 'package', 'synchronized', 'true', 'false', 'null', 'assert', 'enum',
+    'abstract', 'continue', 'for',        'new',       'switch',
+    'assert',   'default',  'if',         'package',   'synchronized',
+    'boolean',  'do',       'goto',       'private',   'this',
+    'break',    'double',   'implements', 'protected', 'throw',
+    'byte',     'else',     'import',     'public',    'throws',
+    'case',     'enum',     'instanceof', 'return',    'transient',
+    'catch',    'extends',  'int',        'short',     'try',
+    'char',     'final',    'interface',  'static',    'void',
+    'class',    'finally',  'long',       'strictfp',  'volatile',
+    'const',    'float',    'native',     'super',     'while',
 ]);
 export function isKeyword(identifier: string): boolean {
     return keywords.has(identifier);

From b03fcb39ca47ee687b85d62d93be9d84dca218df Mon Sep 17 00:00:00 2001
From: Sheng Chen <sheche@microsoft.com>
Date: Sat, 8 May 2021 10:17:10 +0800
Subject: [PATCH 3/4] Address the comments

---
 .../test/plugin/util/TestGenerationUtils.java | 83 +++++++++++--------
 1 file changed, 50 insertions(+), 33 deletions(-)

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 eb1316d4..d0ab9cb6 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
@@ -17,6 +17,7 @@
 
 import org.eclipse.core.runtime.CoreException;
 import org.eclipse.core.runtime.IProgressMonitor;
+import org.eclipse.core.runtime.IStatus;
 import org.eclipse.core.runtime.NullProgressMonitor;
 import org.eclipse.jdt.core.IClasspathEntry;
 import org.eclipse.jdt.core.ICompilationUnit;
@@ -27,6 +28,8 @@
 import org.eclipse.jdt.core.IPackageFragment;
 import org.eclipse.jdt.core.IPackageFragmentRoot;
 import org.eclipse.jdt.core.IType;
+import org.eclipse.jdt.core.JavaConventions;
+import org.eclipse.jdt.core.JavaCore;
 import org.eclipse.jdt.core.JavaModelException;
 import org.eclipse.jdt.core.dom.AST;
 import org.eclipse.jdt.core.dom.ASTNode;
@@ -194,16 +197,16 @@ private static WorkspaceEdit generateTestsFromSource(ICompilationUnit unit, ITyp
         }
 
         final ICompilationUnit testUnit = getTestCompilationUnit(javaProject, testEntry, testFullyQualifiedName);
-        return createTextEditFromSourceFile(testKind, testUnit, methodsToTest, cursorOffset);
+        return scaffoldTestFile(testKind, testUnit, methodsToTest, cursorOffset);
     }
 
-    private static WorkspaceEdit createTextEditFromSourceFile(TestKind kind, ICompilationUnit testUnit,
+    private static WorkspaceEdit scaffoldTestFile(TestKind kind, ICompilationUnit testUnit,
             List<String> methodsToTest, int cursorPosition) throws CoreException {
         if (testUnit.exists()) {
             final IType[] types = testUnit.getAllTypes();
             if (types.length == 0) {
                 // an empty java file
-                return getTestCompilationUnit(kind, testUnit, methodsToTest);
+                return addTestClassToExistingFile(kind, testUnit, methodsToTest);
             }
 
             final CompilationUnit root = (CompilationUnit) TestSearchUtils.parseToAst(testUnit,
@@ -235,17 +238,17 @@ private static WorkspaceEdit createTextEditFromSourceFile(TestKind kind, ICompil
                 // ignore if the upstream does not support insert position preference.
             }
 
-            return createEditToExistingTestFile(root, kind, methodsToTest, (TypeDeclaration) typeNode,
+            return addTestMethodsToExistingTestClass(root, kind, methodsToTest, (TypeDeclaration) typeNode,
                     binding, insertPosition);
         } else {
-            return createAndGetNewTestCompilationUnit(kind, testUnit, methodsToTest);
+            return createNewTestClass(kind, testUnit, methodsToTest);
         }
     }
 
     /**
      * Get a compilation unit change without creating.
      */
-    private static WorkspaceEdit getTestCompilationUnit(TestKind kind, ICompilationUnit testUnit,
+    private static WorkspaceEdit addTestClassToExistingFile(TestKind kind, ICompilationUnit testUnit,
             List<String> methodsToTest) throws CoreException {
         final CompilationUnitChange cuChange = new CompilationUnitChange("", testUnit);
         final String cuContent = constructNewCU(testUnit, methodsToTest, kind);
@@ -256,7 +259,7 @@ private static WorkspaceEdit getTestCompilationUnit(TestKind kind, ICompilationU
     /**
      * Create and get a compilation unit change.
      */
-    private static WorkspaceEdit createAndGetNewTestCompilationUnit(TestKind kind, ICompilationUnit testUnit,
+    private static WorkspaceEdit createNewTestClass(TestKind kind, ICompilationUnit testUnit,
             List<String> methodsToTest) throws CoreException {
         final String cuContent = constructNewCU(testUnit, methodsToTest, kind);
         final CreateCompilationUnitChange change =
@@ -264,7 +267,7 @@ private static WorkspaceEdit createAndGetNewTestCompilationUnit(TestKind kind, I
         return ChangeUtil.convertToWorkspaceEdit(change);
     }
 
-    private static WorkspaceEdit createEditToExistingTestFile(CompilationUnit testRoot, TestKind kind,
+    private static WorkspaceEdit addTestMethodsToExistingTestClass(CompilationUnit testRoot, TestKind kind,
             List<String> methodsToTest, TypeDeclaration typeNode, ITypeBinding typeBinding,
             IJavaElement insertPosition) throws JavaModelException, CoreException {
         final String testAnnotation = getTestAnnotation(kind);
@@ -356,42 +359,56 @@ private static IClasspathEntry getTestClasspathEntry(ICompilationUnit unit) thro
 
     private static String getTestFullyQualifiedName(ITypeBinding typeBinding, IJavaProject project,
             IClasspathEntry testEntry) throws JavaModelException {
-        String promptName = typeBinding.getBinaryName() + "Tests";
-        final ICompilationUnit testCompilationUnit = getTestCompilationUnit(project, testEntry, promptName);
-        if (!testCompilationUnit.exists()) {
-            promptName = typeBinding.getBinaryName() + "Test";
-        }
-
+        final String promptName = getDefaultTestFullyQualifiedName(typeBinding, project, testEntry);
         final String fullyQualifiedName = (String) JUnitPlugin.askClientForInput(
             "Please type the target test class name", promptName);
         if (fullyQualifiedName == null) {
             return null;
         }
 
-        if (fullyQualifiedName.charAt(fullyQualifiedName.length() - 1) == '.') {
-            JavaLanguageServerPlugin.getInstance().getClientConnection().showNotificationMessage(MessageType.Error,
-                    "Invalid Java class name: " + fullyQualifiedName);
-            return null;
+        final IStatus status = JavaConventions.validateJavaTypeName(fullyQualifiedName,
+                project.getOption(JavaCore.COMPILER_SOURCE, true),
+                project.getOption(JavaCore.COMPILER_COMPLIANCE, true),
+                project.getOption(JavaCore.COMPILER_PB_ENABLE_PREVIEW_FEATURES, true)
+        );
+
+        final int severity = status.getSeverity();
+        if (severity == IStatus.OK || severity == IStatus.WARNING) {
+            return fullyQualifiedName;
         }
 
-        final String[] identifiers = fullyQualifiedName.split("\\.");
-        for (final String identifier : identifiers) {
-            final char[] chars = identifier.toCharArray();
-            if (!Character.isJavaIdentifierStart(chars[0])) {
-                JavaLanguageServerPlugin.getInstance().getClientConnection().showNotificationMessage(MessageType.Error,
-                    "Invalid Java identifier: " + identifier);
-                return null;
-            }
-            for (int i = 1; i < chars.length; i++) {
-                if (!Character.isJavaIdentifierPart(chars[i])) {
-                    JavaLanguageServerPlugin.getInstance().getClientConnection().showNotificationMessage(
-                        MessageType.Error, "Invalid Java identifier: " + identifier);
-                    return null;
-                }
+        JavaLanguageServerPlugin.getInstance().getClientConnection().showNotificationMessage(MessageType.Error,
+                status.getMessage());
+        return null;
+    }
+
+    private static String getDefaultTestFullyQualifiedName(ITypeBinding typeBinding, IJavaProject project,
+            IClasspathEntry testEntry) throws JavaModelException {
+        final String defaultName = typeBinding.getBinaryName() + "Test";
+        final String attemptName = typeBinding.getBinaryName() + "Tests";
+        final ICompilationUnit testCompilationUnit = getTestCompilationUnit(project, testEntry, attemptName);
+        if (testCompilationUnit.exists()) {
+            return attemptName;
+        }
+
+        // check the majority naming under the package and use it as the default type name
+        final IPackageFragment packageFragment = (IPackageFragment) testCompilationUnit.getParent();
+        int counter = 0;
+        final ICompilationUnit[] compilationUnits = packageFragment.getCompilationUnits();
+        for (final ICompilationUnit unit : compilationUnits) {
+            final String name = unit.getElementName();
+            if (name.endsWith("Tests.java")) {
+                counter++;
+            } else if (name.endsWith("Test.java")) {
+                counter--;
             }
         }
 
-        return fullyQualifiedName;
+        if (counter > 0) {
+            return attemptName;
+        }
+
+        return defaultName;
     }
 
     private static ICompilationUnit getTestCompilationUnit(IJavaProject javaProject, IClasspathEntry testEntry,

From c76975ae13ad8cfe6fd0114f89320e253d236547 Mon Sep 17 00:00:00 2001
From: Sheng Chen <sheche@microsoft.com>
Date: Sat, 8 May 2021 13:40:39 +0800
Subject: [PATCH 4/4] Address comments

---
 .../test/plugin/util/TestGenerationUtils.java | 34 +++++++++++++------
 1 file changed, 23 insertions(+), 11 deletions(-)

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 d0ab9cb6..d9d344ce 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
@@ -187,6 +187,11 @@ private static WorkspaceEdit generateTestsFromSource(ICompilationUnit unit, ITyp
         }
 
         final IClasspathEntry testEntry = getTestClasspathEntry(unit);
+        if (testEntry == null) {
+            JavaLanguageServerPlugin.getInstance().getClientConnection().showNotificationMessage(MessageType.Error,
+                    "Cannot find a valid classpath entry to generate tests.");
+            return null;
+        }
         final String testFullyQualifiedName = getTestFullyQualifiedName(typeBinding, javaProject, testEntry);
         if (testFullyQualifiedName == null) {
             return null;
@@ -299,14 +304,16 @@ private static String constructCUContent(ICompilationUnit testUnit, TestKind tes
                 .append(packageFragment.getElementName())
                 .append(";")
                 .append(lineDelimiter)
-                .append(lineDelimiter)
-                .append("import ")
-                .append(getTestAnnotation(testKind))
-                .append(";")
-                .append(lineDelimiter)
                 .append(lineDelimiter);
         }
-        buf.append(typeContent);
+
+        buf.append("import ")
+            .append(getTestAnnotation(testKind))
+            .append(";")
+            .append(lineDelimiter)
+            .append(lineDelimiter)
+            .append(typeContent);
+
         return buf.toString();
     }
 
@@ -341,9 +348,9 @@ private static IClasspathEntry getTestClasspathEntry(ICompilationUnit unit) thro
         final IJavaProject javaProject = unit.getJavaProject();
         // In most cases, this is the classpath entry used for testing, we first find the target entry by hard-code
         // to avoid go into the generated entries.
-        final IClasspathEntry testEntry = javaProject.getClasspathEntryFor(
+        IClasspathEntry testEntry = javaProject.getClasspathEntryFor(
             javaProject.getPath().append("src/test/java"));
-        if (ProjectTestUtils.isTestEntry(testEntry)) {
+        if (testEntry != null && ProjectTestUtils.isTestEntry(testEntry)) {
             return testEntry;
         }
 
@@ -352,9 +359,13 @@ private static IClasspathEntry getTestClasspathEntry(ICompilationUnit unit) thro
             if (ProjectTestUtils.isTestEntry(entry)) {
                 return entry;
             }
-        }
 
-        return entries[0];
+            if (entry.getPath().isPrefixOf(unit.getPath())) {
+                testEntry = entry;
+            }
+        }
+        
+        return testEntry;
     }
 
     private static String getTestFullyQualifiedName(ITypeBinding typeBinding, IJavaProject project,
@@ -414,7 +425,8 @@ private static String getDefaultTestFullyQualifiedName(ITypeBinding typeBinding,
     private static ICompilationUnit getTestCompilationUnit(IJavaProject javaProject, IClasspathEntry testEntry,
             String testFullyQualifiedName) throws JavaModelException {
         final IPackageFragmentRoot packageRoot = javaProject.findPackageFragmentRoot(testEntry.getPath());
-        final String packageQualifiedName = testFullyQualifiedName.substring(0,
+        final int lastDelimiterIndex = testFullyQualifiedName.lastIndexOf(".");
+        final String packageQualifiedName = lastDelimiterIndex < 0 ? "" : testFullyQualifiedName.substring(0,
                 testFullyQualifiedName.lastIndexOf("."));
         final IPackageFragment packageFragment = packageRoot.getPackageFragment(packageQualifiedName);
         final String compilationUnitName = testFullyQualifiedName.substring(