ThenReturn.scala

import scala.meta._
import scalafix.v1._

class ThenReturn extends SyntacticRule("ThenReturn") {
  val exclude = Seq("")
  override def fix(implicit doc: SyntacticDocument): Patch = {
    doc.input match {
      case f: Input.VirtualFile if exclude.exists(f.path.contains) =>
        Patch.empty
      case _ =>
        Option
          .when(doc.tree.collect { case Term.Name("andThen") =>
            ()
          }.isEmpty) {
            val result = doc.tree.collect {
              case Term.ApplyInfix.After_4_6_0(
                    a,
                    returns @ Term.Name("returns"),
                    Type.ArgClause(Nil),
                    Term.ArgClause(
                      List(
                        b
                      ),
                      None
                    )
                  ) if Option(a).collect {
                    case Term.Select(
                          Term.Placeholder(),
                          _: Term.Name
                        ) =>
                      ()
                  }.isEmpty && a.collect { case Term.Name("when") => () }.isEmpty && !b.is[Term.PartialFunction] =>
                Seq(
                  Patch.addAround(a, "Mockito.when(", ")."),
                  Patch.replaceTree(returns, "thenReturn("),
                  Patch.addRight(b, ")")
                ).asPatch
              case Term.Apply.After_4_6_0(
                    Term.Select(
                      a,
                      returns @ Term.Name("returns")
                    ),
                    Term.ArgClause(
                      List(
                        b
                      ),
                      None
                    )
                  ) if Option(a).collect {
                    case Term.Select(
                          Term.Placeholder(),
                          _: Term.Name
                        ) =>
                      ()
                  }.isEmpty && a.collect { case Term.Name("when") => () }.isEmpty && !b.is[Term.PartialFunction] =>
                Seq(
                  Patch.addAround(a, "Mockito.when(", ")"),
                  Patch.replaceTree(returns, "thenReturn")
                ).asPatch
            }

            Option
              .when(result.nonEmpty) {
                Seq(
                  result.asPatch,
                  Option
                    .when(doc.tree.collect { case importer"org.mockito.Mockito" =>
                      ()
                    }.isEmpty) {
                      Patch.addGlobalImport(importer"org.mockito.Mockito")
                    }
                    .asPatch
                ).asPatch
              }
              .asPatch
          }
          .asPatch
    }
  }
}