plugins/testng/src/com/theoryinpractice/testng/TestNGFramework.java

/*
 * Copyright 2000-2013 JetBrains s.r.o.
 *
 * 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.theoryinpractice.testng;

import com.intellij.CommonBundle;
import com.intellij.codeInsight.AnnotationUtil;
import com.intellij.codeInsight.intention.AddAnnotationFix;
import com.intellij.ide.fileTemplates.FileTemplateDescriptor;
import com.intellij.openapi.application.ApplicationManager;
import com.intellij.openapi.ui.Messages;
import com.intellij.psi.*;
import com.intellij.psi.codeStyle.JavaCodeStyleManager;
import com.intellij.testIntegration.JavaTestFramework;
import com.intellij.util.IncorrectOperationException;
import com.intellij.util.PathUtil;
import com.theoryinpractice.testng.util.TestNGUtil;
import icons.TestngIcons;
import org.jetbrains.annotations.NonNls;
import org.jetbrains.annotations.NotNull;
import org.jetbrains.annotations.Nullable;
import org.testng.annotations.BeforeMethod;
import org.testng.annotations.Test;

import javax.swing.*;

public class TestNGFramework extends JavaTestFramework {
  @NotNull
  public String getName() {
    return "TestNG";
  }

  @NotNull
  @Override
  public Icon getIcon() {
    return TestngIcons.TestNG;
  }

  protected String getMarkerClassFQName() {
    return "org.testng.annotations.Test";
  }

  @NotNull
  public String getLibraryPath() {
    try {
      return PathUtil.getJarPathForClass(Class.forName("org.testng.annotations.Test"));
    }
    catch (ClassNotFoundException e) {
      throw new RuntimeException(e);
    }
  }

  @Nullable
  public String getDefaultSuperClass() {
    return null;
  }

  public boolean isTestClass(PsiClass clazz, boolean canBePotential) {
    if (canBePotential) return isUnderTestSources(clazz);
    return TestNGUtil.isTestNGClass(clazz);
  }

  @Nullable
  @Override
  protected PsiMethod findSetUpMethod(@NotNull PsiClass clazz) {
    for (PsiMethod each : clazz.getMethods()) {
      if (AnnotationUtil.isAnnotated(each, "org.testng.annotations.BeforeMethod", false)) return each;
    }
    return null;
  }

  @Nullable
  @Override
  protected PsiMethod findTearDownMethod(@NotNull PsiClass clazz) {
    for (PsiMethod each : clazz.getMethods()) {
      if (AnnotationUtil.isAnnotated(each, "org.testng.annotations.AfterMethod", false)) return each;
    }
    return null;
  }

  @Override
  protected PsiMethod findOrCreateSetUpMethod(PsiClass clazz) throws IncorrectOperationException {
    PsiMethod method = findSetUpMethod(clazz);
    if (method != null) return method;

    final PsiManager manager = clazz.getManager();
    final PsiElementFactory factory = JavaPsiFacade.getInstance(manager.getProject()).getElementFactory();
    String setUpName = "setUp";
    PsiMethod patternMethod = createSetUpPatternMethod(factory);
    PsiMethod inClass = clazz.findMethodBySignature(patternMethod, false);
    if (inClass != null) {
      int exit = ApplicationManager.getApplication().isUnitTestMode() ?
                 Messages.YES :
                 Messages.showYesNoDialog(manager.getProject(), "Method \'" + setUpName + "\' already exist but is not annotated as @BeforeMethod.",
                                          CommonBundle.getWarningTitle(),
                                          "Annotate",
                                          "Create new method",
                                          Messages.getWarningIcon());
      if (exit == Messages.YES) {
        new AddAnnotationFix(BeforeMethod.class.getName(), inClass).invoke(inClass.getProject(), null, inClass.getContainingFile());
        return inClass;
      } else if (exit == Messages.NO) {
        inClass = null;
        int i = 0;
        while (clazz.findMethodBySignature(patternMethod, false) != null) {
          patternMethod.setName(setUpName + (++i));
        }
        setUpName = patternMethod.getName();
      }
    }

    final PsiClass superClass = clazz.getSuperClass();
    if (superClass != null) {
      final PsiMethod[] methods = superClass.findMethodsBySignature(patternMethod, false);
      if (methods.length > 0) {
        final PsiModifierList modifierList = methods[0].getModifierList();
        if (!modifierList.hasModifierProperty(PsiModifier.PRIVATE)) { //do not override private method
          @NonNls String pattern = "@" + BeforeMethod.class.getName() + "\n";
          if (modifierList.hasModifierProperty(PsiModifier.PROTECTED)) {
            pattern += "protected ";
          }
          else if (modifierList.hasModifierProperty(PsiModifier.PUBLIC)) {
            pattern += "public ";
          }
          patternMethod = factory.createMethodFromText(pattern + "void " + setUpName + "() throws Exception {\nsuper." + setUpName + "();\n}", null);
        }
      }
    }

    final PsiMethod[] psiMethods = clazz.getMethods();

    PsiMethod testMethod = null;
    for (PsiMethod psiMethod : psiMethods) {
      if (inClass == null && AnnotationUtil.isAnnotated(psiMethod, BeforeMethod.class.getName(), false)) {
        inClass = psiMethod;
      }
      if (testMethod == null && AnnotationUtil.isAnnotated(psiMethod, Test.class.getName(), false) && !psiMethod.hasModifierProperty(PsiModifier.PRIVATE)) {
        testMethod = psiMethod;
      }
    }
    if (inClass == null) {
      final PsiMethod psiMethod;
      if (testMethod != null) {
        psiMethod = (PsiMethod)clazz.addBefore(patternMethod, testMethod);
      } else {
        psiMethod = (PsiMethod)clazz.add(patternMethod);
      }
      JavaCodeStyleManager.getInstance(clazz.getProject()).shortenClassReferences(clazz);
      return psiMethod;
    }
    else if (inClass.getBody() == null) {
      return (PsiMethod)inClass.replace(patternMethod);
    }
    return inClass;
  }

  @Override
  public char getMnemonic() {
    return 'N';
  }

  @Override
  public FileTemplateDescriptor getTestClassFileTemplateDescriptor() {
    return new FileTemplateDescriptor("TestNG Test Class.java");
  }

  public FileTemplateDescriptor getSetUpMethodFileTemplateDescriptor() {
    return new FileTemplateDescriptor("TestNG SetUp Method.java");
  }

  public FileTemplateDescriptor getTearDownMethodFileTemplateDescriptor() {
    return new FileTemplateDescriptor("TestNG TearDown Method.java");
  }

  public FileTemplateDescriptor getTestMethodFileTemplateDescriptor() {
    return new FileTemplateDescriptor("TestNG Test Method.java");
  }

  @Nullable
  @Override
  public FileTemplateDescriptor getParametersMethodFileTemplateDescriptor() {
    return new FileTemplateDescriptor("TestNG Parameters Method.java");
  }

  @Override
  public boolean isTestMethod(PsiElement element) {
    return element instanceof PsiMethod && TestNGUtil.hasTest((PsiModifierListOwner)element);
  }
}