From fcd5e310d54ede82ee48eeba89fbe12cd85c2c9e Mon Sep 17 00:00:00 2001
From: Sheng Chen <sheche@microsoft.com>
Date: Mon, 24 May 2021 15:33:41 +0800
Subject: [PATCH 1/2] fix: Can select peer project if current project does not
 contain test framework

---
 .../test/plugin/util/TestGenerationUtils.java | 44 +++++++++++++++++--
 1 file changed, 40 insertions(+), 4 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 250cb823..5195dbd1 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
@@ -61,6 +61,7 @@
 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.ProjectUtils;
 import org.eclipse.jdt.ls.core.internal.handlers.CodeGenerationUtils;
 import org.eclipse.jdt.ls.core.internal.text.correction.SourceAssistProcessor;
 import org.eclipse.lsp4j.MessageType;
@@ -180,14 +181,17 @@ public static WorkspaceEdit generateTests(List<Object> arguments, IProgressMonit
      */
     private static WorkspaceEdit generateTestsFromSource(ICompilationUnit unit, ITypeBinding typeBinding,
             int cursorOffset) throws CoreException {
-        final IJavaProject javaProject = unit.getJavaProject();
+        final IJavaProject javaProject = determineTestProject(unit);
+        if (javaProject == null) {
+            return null;
+        }
         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 IClasspathEntry testEntry = getTestClasspathEntry(javaProject, unit);
         if (testEntry == null) {
             JavaLanguageServerPlugin.getInstance().getClientConnection().showNotificationMessage(MessageType.Error,
                     "Cannot find a valid classpath entry to generate tests.");
@@ -352,8 +356,8 @@ private static String constructMethodStub(ICompilationUnit testUnit, TestKind te
                 lineDelimiter, testUnit.getJavaProject().getOptions(true));
     }
 
-    private static IClasspathEntry getTestClasspathEntry(ICompilationUnit unit) throws JavaModelException {
-        final IJavaProject javaProject = unit.getJavaProject();
+    private static IClasspathEntry getTestClasspathEntry(IJavaProject javaProject, ICompilationUnit unit)
+            throws JavaModelException {
         // 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.
         IClasspathEntry testEntry = javaProject.getClasspathEntryFor(
@@ -563,6 +567,38 @@ private static List<String> determineMethodsToGenerate(List<String> lifecycleAnn
                 methodList, true /*pickMany*/);
     }
 
+    private static IJavaProject determineTestProject(ICompilationUnit unit) {
+        final IJavaProject javaProject = unit.getJavaProject();
+        if (TestKindProvider.getTestKindsFromCache(javaProject).size() > 0) {
+            return javaProject;
+        }
+
+        final IJavaProject[] javaProjects = ProjectUtils.getJavaProjects();
+        final List<Option> javaTestProjects = new LinkedList<>();
+        for (final IJavaProject project : javaProjects) {
+            if (project.equals(JavaLanguageServerPlugin.getProjectsManager().getDefaultProject())) {
+                continue;
+            }
+            if (TestKindProvider.getTestKindsFromCache(project).size() > 0) {
+                javaTestProjects.add(new Option(project.getHandleIdentifier(), project.getElementName(),
+                    project.getProject().getLocation().toOSString()));
+            }
+        }
+
+        if (javaTestProjects.size() == 0) {
+            JavaLanguageServerPlugin.getInstance().getClientConnection().showNotificationMessage(MessageType.Error,
+                    "The projects in the workspace do not have a test framework," +
+                    " please make sure they are configured correctly.");
+            return null;
+        }
+        final Object result = JUnitPlugin.askClientForChoice("Select a project where tests are generated in",
+                javaTestProjects);
+        if (result == null) {
+            return null;
+        }
+        return (IJavaProject) JavaCore.create((String) result);
+    }
+
     private static TestKind determineTestFramework(Set<TestKind> availableFrameworks) throws CoreException {
         if (availableFrameworks.size() == 0) {
             JavaLanguageServerPlugin.getInstance().getClientConnection().showNotificationMessage(MessageType.Error,

From 72fe5af63b0abb8447e38920c902609d05502b2c Mon Sep 17 00:00:00 2001
From: Sheng Chen <sheche@microsoft.com>
Date: Tue, 25 May 2021 10:06:58 +0800
Subject: [PATCH 2/2] Address comments

---
 .../test/plugin/util/TestGenerationUtils.java     | 15 +++++++++------
 1 file changed, 9 insertions(+), 6 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 5195dbd1..1a8dbfb0 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
@@ -286,7 +286,8 @@ private static WorkspaceEdit addTestMethodsToExistingTestClass(CompilationUnit t
             return new MethodMetaData(methodName, testAnnotation);
         }).collect(Collectors.toList());
 
-        final TextEdit edit = getTextEdit(kind, metadata, testRoot, typeNode, typeBinding, insertPosition);
+        final TextEdit edit = getTextEdit(kind, metadata, testRoot, typeNode, typeBinding, insertPosition,
+                false /*fromTest*/);
         return SourceAssistProcessor.convertToWorkspaceEdit((ICompilationUnit) testRoot.getJavaElement(), edit);
     }
 
@@ -587,8 +588,7 @@ private static IJavaProject determineTestProject(ICompilationUnit unit) {
 
         if (javaTestProjects.size() == 0) {
             JavaLanguageServerPlugin.getInstance().getClientConnection().showNotificationMessage(MessageType.Error,
-                    "The projects in the workspace do not have a test framework," +
-                    " please make sure they are configured correctly.");
+                "No test library found in your workspace, please add a test library to your project classpath first.");
             return null;
         }
         final Object result = JUnitPlugin.askClientForChoice("Select a project where tests are generated in",
@@ -655,11 +655,11 @@ private static TextEdit createTextEditFromTestFile(TestKind kind, List<String> m
             return new MethodMetaData(methodName, prefix + annotationName);
         }).collect(Collectors.toList());
 
-        return getTextEdit(kind, metadata, root, typeNode, typeBinding, insertPosition);
+        return getTextEdit(kind, metadata, root, typeNode, typeBinding, insertPosition, true /*fromTest*/);
     }
 
     private static TextEdit getTextEdit(TestKind kind, List<MethodMetaData> methodMetadata, CompilationUnit root,
-            TypeDeclaration typeNode, ITypeBinding typeBinding, IJavaElement insertPosition)
+            TypeDeclaration typeNode, ITypeBinding typeBinding, IJavaElement insertPosition, boolean fromTest)
             throws CoreException {
         final ASTRewrite astRewrite = ASTRewrite.create(root.getAST());
         final ImportRewrite importRewrite = StubUtility.createImportRewrite(root, true);
@@ -703,7 +703,10 @@ private static TextEdit getTextEdit(TestKind kind, List<MethodMetaData> methodMe
             }
 
             final ASTNode insertion = StubUtility2Core.getNodeToInsertBefore(listRewrite, insertPosition);
-            if (insertion != null) {
+            // Only try to insert according to the cursor position when it's triggered from test file,
+            // If it's triggered from source file, since code snippets will be generated into a new file,
+            // they will always insert to the last.
+            if (insertion != null && fromTest) {
                 listRewrite.insertBefore(decl, insertion, null);
             } else {
                 listRewrite.insertLast(decl, null);