@Override
  protected boolean setupConfigurationFromContext(
      TestNGConfiguration configuration,
      ConfigurationContext context,
      Ref<PsiElement> sourceElement) {
    // TODO: check TestNG Pattern running first, before method/class (see
    // TestNGInClassConfigurationProducer for logic)
    // TODO: and PsiClassOwner not handled, which is in TestNGInClassConfigurationProducer

    Location location = context.getLocation();
    if (location == null) {
      return false;
    }

    Project project = context.getProject();
    PsiElement leaf = location.getPsiElement();

    if (!ProjectRootsUtil.isInProjectOrLibSource(leaf)) {
      return false;
    }

    if (!(leaf.getContainingFile() instanceof KtFile)) {
      return false;
    }

    KtFile jetFile = (KtFile) leaf.getContainingFile();

    if (ProjectStructureUtil.isJsKotlinModule(jetFile)) {
      return false;
    }

    KtNamedDeclaration declarationToRun = getDeclarationToRun(leaf);

    if (declarationToRun instanceof KtNamedFunction) {
      KtNamedFunction function = (KtNamedFunction) declarationToRun;

      @SuppressWarnings("unchecked")
      KtElement owner = PsiTreeUtil.getParentOfType(function, KtFunction.class, KtClass.class);

      if (owner instanceof KtClass) {
        PsiClass delegate = LightClassUtil.INSTANCE.getPsiClass((KtClass) owner);
        if (delegate != null) {
          for (PsiMethod method : delegate.getMethods()) {
            if (method.getNavigationElement() == function) {
              if (TestNGUtil.hasTest(method)) {
                return configure(configuration, location, context, project, delegate, method);
              }
              break;
            }
          }
        }
      }
    }

    if (declarationToRun instanceof KtClass) {
      PsiClass delegate = LightClassUtil.INSTANCE.getPsiClass((KtClassOrObject) declarationToRun);
      if (!isTestNGClass(delegate)) {
        return false;
      }

      return configure(configuration, location, context, project, delegate, null);
    }

    return false;
  }
  @Override
  public void onFirstRun(
      ConfigurationFromContext configuration,
      ConfigurationContext context,
      Runnable startRunnable) {
    KtNamedDeclaration declarationToRun = getDeclarationToRun(configuration.getSourceElement());
    final PsiNamedElement lightElement =
        CollectionsKt.firstOrNull(LightClassUtilsKt.toLightElements(declarationToRun));

    // Copied from TestNGInClassConfigurationProducer.onFirstRun()
    if (lightElement instanceof PsiMethod || lightElement instanceof PsiClass) {
      PsiMethod psiMethod;
      PsiClass containingClass;

      if (lightElement instanceof PsiMethod) {
        psiMethod = (PsiMethod) lightElement;
        containingClass = psiMethod.getContainingClass();
      } else {
        psiMethod = null;
        containingClass = (PsiClass) lightElement;
      }

      InheritorChooser inheritorChooser =
          new InheritorChooser() {
            @Override
            protected void runForClasses(
                List<PsiClass> classes,
                PsiMethod method,
                ConfigurationContext context,
                Runnable performRunnable) {
              ((TestNGConfiguration) context.getConfiguration().getConfiguration())
                  .bePatternConfiguration(classes, method);
              super.runForClasses(classes, method, context, performRunnable);
            }

            @Override
            protected void runForClass(
                PsiClass aClass,
                PsiMethod psiMethod,
                ConfigurationContext context,
                Runnable performRunnable) {
              if (lightElement instanceof PsiMethod) {
                Project project = psiMethod.getProject();
                MethodLocation methodLocation =
                    new MethodLocation(project, psiMethod, PsiLocation.fromPsiElement(aClass));
                ((TestNGConfiguration) context.getConfiguration().getConfiguration())
                    .setMethodConfiguration(methodLocation);
              } else {
                ((TestNGConfiguration) context.getConfiguration().getConfiguration())
                    .setClassConfiguration(aClass);
              }
              super.runForClass(aClass, psiMethod, context, performRunnable);
            }
          };
      if (inheritorChooser.runMethodInAbstractClass(
          context,
          startRunnable,
          psiMethod,
          containingClass,
          new Condition<PsiClass>() {
            @Override
            public boolean value(PsiClass aClass) {
              return aClass.hasModifierProperty(PsiModifier.ABSTRACT) && TestNGUtil.hasTest(aClass);
            }
          })) return;
    }

    super.onFirstRun(configuration, context, startRunnable);
  }