Merge "Implement Gradle awared \'Add JUnit dependency\' quick fix" into studio-1.4-dev automerge: 2e6ab42

automerge: 8bd7cc6

* commit '8bd7cc60b5e3eb22ed31b19bc183f376ce5af0f5':
  Implement Gradle awared 'Add JUnit dependency' quick fix

2 files changed, 150 insertions, 1 deletions

diff --git a/android/src/com/android/tools/idea/quickfix/AddGradleJUnitDependencyFix.java b/android/src/com/android/tools/idea/quickfix/AddGradleJUnitDependencyFix.java
new file mode 100644
index 00000000000..1c03473d97d
--- /dev/null
+++ b/android/src/com/android/tools/idea/quickfix/AddGradleJUnitDependencyFix.java
@@ -0,0 +1,110 @@
+/*
+ * Copyright (C) 2015 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.android.tools.idea.quickfix;
+
+import com.android.tools.idea.gradle.parser.Dependency;
+import com.android.tools.idea.gradle.project.GradleProjectImporter;
+import com.android.tools.idea.gradle.project.GradleSyncListener;
+import com.google.common.base.Function;
+import com.google.common.collect.ImmutableList;
+import com.google.common.collect.Lists;
+import com.intellij.codeInsight.daemon.QuickFixBundle;
+import com.intellij.codeInsight.daemon.impl.actions.AddImportAction;
+import com.intellij.openapi.application.Application;
+import com.intellij.openapi.application.ApplicationManager;
+import com.intellij.openapi.editor.Editor;
+import com.intellij.openapi.module.Module;
+import com.intellij.openapi.project.DumbService;
+import com.intellij.openapi.project.Project;
+import com.intellij.openapi.roots.ModuleRootManager;
+import com.intellij.openapi.vfs.VirtualFile;
+import com.intellij.psi.JavaPsiFacade;
+import com.intellij.psi.PsiClass;
+import com.intellij.psi.PsiFile;
+import com.intellij.psi.PsiReference;
+import com.intellij.psi.search.GlobalSearchScope;
+import com.intellij.psi.util.PsiUtilCore;
+import org.jetbrains.annotations.NotNull;
+import org.jetbrains.annotations.Nullable;
+
+import java.util.List;
+
+/**
+ * Quickfix to add JUnit dependency to gradle.build file and sync the project.
+ */
+public class AddGradleJUnitDependencyFix extends GradleDependencyFix {
+  private final Module myCurrentModule;
+  private final PsiReference myReference;
+  private final String myClassName;
+  private final boolean myIsJunit4;
+
+  public AddGradleJUnitDependencyFix(@NotNull Module currentModule, @NotNull PsiReference reference, @NotNull String className,
+                                     boolean isJunit4) {
+    myCurrentModule = currentModule;
+    myReference = reference;
+    myClassName = className;
+    myIsJunit4 = isJunit4;
+  }
+
+  @Override
+  @NotNull
+  public String getText() {
+    return QuickFixBundle.message("orderEntry.fix.add.junit.jar.to.classpath");
+  }
+
+  @Override
+  @NotNull
+  public String getFamilyName() {
+    return getText();
+  }
+
+  @Override
+  public boolean isAvailable(@NotNull Project project, @Nullable Editor editor, @Nullable PsiFile file) {
+    return !project.isDisposed() && !myCurrentModule.isDisposed();
+  }
+
+  @Override
+  public void invoke(@NotNull final Project project, @Nullable final Editor editor, @Nullable PsiFile file) {
+    String gradleDependencyEntry = myIsJunit4 ? "junit:junit:4.12" : "junit:junit:3.8.1";
+
+    VirtualFile location = PsiUtilCore.getVirtualFile(myReference.getElement());
+    boolean inTests = location != null && ModuleRootManager.getInstance(myCurrentModule).getFileIndex().isInTestSourceContent(location);
+
+    final Dependency dependency = new Dependency(getDependencyScope(myCurrentModule, inTests), Dependency.Type.EXTERNAL,
+                                                 gradleDependencyEntry);
+
+    final Application application = ApplicationManager.getApplication();
+    application.invokeAndWait(new Runnable() {
+      @Override
+      public void run() {
+        application.runWriteAction(new Runnable() {
+          @Override
+          public void run() {
+            addDependency(myCurrentModule, dependency);
+            gradleSyncAndImportClass(myCurrentModule, editor, myReference, new Function<Void, List<PsiClass>>() {
+               @Override
+               public List<PsiClass> apply(@Nullable Void input) {
+                 PsiClass aClass =
+                   JavaPsiFacade.getInstance(project).findClass(myClassName, GlobalSearchScope.moduleWithLibrariesScope(myCurrentModule));
+                 return aClass != null ? ImmutableList.of(aClass) : null;
+               }
+            });
+          }
+        });
+      }
+    }, application.getDefaultModalityState());
+  }
+}
diff --git a/android/src/com/android/tools/idea/quickfix/AndroidUnresolvedReferenceQuickFixProvider.java b/android/src/com/android/tools/idea/quickfix/AndroidUnresolvedReferenceQuickFixProvider.java
index 00dd73d9643..fdf50cea8ab 100644
--- a/android/src/com/android/tools/idea/quickfix/AndroidUnresolvedReferenceQuickFixProvider.java
+++ b/android/src/com/android/tools/idea/quickfix/AndroidUnresolvedReferenceQuickFixProvider.java
@@ -33,6 +33,9 @@ import com.intellij.packageDependencies.DependencyValidationManager;
 import com.intellij.psi.*;
 import com.intellij.psi.search.GlobalSearchScope;
 import com.intellij.psi.search.PsiShortNamesCache;
+import com.intellij.psi.util.PsiTreeUtil;
+import com.intellij.psi.util.PsiUtil;
+import org.jetbrains.annotations.NonNls;
 import org.jetbrains.annotations.NotNull;
 
 import java.util.ArrayList;
@@ -87,11 +90,21 @@ public class AndroidUnresolvedReferenceQuickFixProvider extends UnresolvedRefere
       return;
     }
 
-    // TODO implement a quickfix that could properly "add junit dependency" to the gradle file.
     PsiElement psiElement = reference.getElement();
     String referenceName = reference.getRangeInElement().substring(psiElement.getText());
     Project project = psiElement.getProject();
 
+    // Check if it is a JUnit class reference.
+    if ("TestCase".equals(referenceName) || isAnnotation(psiElement) && isJunitAnnotationName(referenceName, psiElement)) {
+      final boolean isJunit4 = !referenceName.equals("TestCase");
+      String className = isJunit4 ? "org.junit." + referenceName : "junit.framework.TestCase";
+      PsiClass found =
+        JavaPsiFacade.getInstance(project).findClass(className, contextModule.getModuleWithDependenciesAndLibrariesScope(true));
+      if (found == null) {
+        registrar.register(new AddGradleJUnitDependencyFix(contextModule, reference, className, isJunit4));
+      }
+    }
+
     // Check if we could fix it by introduce gradle dependency.
     PsiClass[] classes = PsiShortNamesCache.getInstance(project).getClassesByName(referenceName, GlobalSearchScope.allScope(project));
     List<PsiClass> allowedDependencies = filterAllowedDependencies(psiElement, classes);
@@ -165,6 +178,32 @@ public class AndroidUnresolvedReferenceQuickFixProvider extends UnresolvedRefere
     return result;
   }
 
+  // Duplicated from com.intellij.codeInsight.daemon.impl.quickfix.OrderEntryFix.isAnnotation
+  private static boolean isAnnotation(@NotNull final PsiElement psiElement) {
+    return PsiTreeUtil.getParentOfType(psiElement, PsiAnnotation.class) != null && PsiUtil.isLanguageLevel5OrHigher(psiElement);
+  }
+
+  // Duplicated from com.intellij.codeInsight.daemon.impl.quickfix.OrderEntryFix.isJunitAnnotationName
+  private static boolean isJunitAnnotationName(@NonNls final String referenceName, @NotNull final PsiElement psiElement) {
+    if ("Test".equals(referenceName) || "Ignore".equals(referenceName) || "RunWith".equals(referenceName) ||
+        "Before".equals(referenceName) || "BeforeClass".equals(referenceName) ||
+        "After".equals(referenceName) || "AfterClass".equals(referenceName)) {
+      return true;
+    }
+    final PsiElement parent = psiElement.getParent();
+    if (parent != null && !(parent instanceof PsiAnnotation)) {
+      final PsiReference reference = parent.getReference();
+      if (reference != null) {
+        final String referenceText = parent.getText();
+        if (isJunitAnnotationName(reference.getRangeInElement().substring(referenceText), parent)) {
+          final int lastDot = referenceText.lastIndexOf('.');
+          return lastDot > -1 && referenceText.substring(0, lastDot).equals("org.junit");
+        }
+      }
+    }
+    return false;
+  }
+
   @NotNull
   @Override
   public Class<PsiJavaCodeReferenceElement> getReferenceClass() {