WasCalledTimes.scala

package fix

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

object WasCalledTimes {
  val Times: PartialFunction[Tree, Tree] = {
    case Term.Select(
          n,
          Term.Name("times")
        ) =>
      n
    case Term.Name("Once" | "once") =>
      Lit.Int(1)
    case Term.Name("twice") =>
      Lit.Int(2)
    case Term.Name("thrice" | "threeTimes") =>
      Lit.Int(3)
    case Term.Name("fourTimes") =>
      Lit.Int(4)
  }
}

class WasCalledTimes extends SyntacticRule("WasCalledTimes") {
  override def fix(implicit doc: SyntacticDocument): Patch = {
    val result = doc.tree.collect {
      case t @ Term.ApplyInfix.After_4_6_0(
            Term.Select(
              obj,
              method
            ),
            Term.Name("wasCalled"),
            Type.ArgClause(Nil),
            Term.ArgClause(
              List(
                WasCalledTimes.Times(n)
              ),
              None
            )
          ) =>
        Patch.replaceTree(t, s"Mockito.verify($obj, Mockito.times($n)).${method}")
      case t @ Term.ApplyInfix.After_4_6_0(
            Term.Apply.After_4_6_0(
              Term.Select(
                obj,
                method
              ),
              Term.ArgClause(
                args,
                None
              )
            ),
            Term.Name("wasCalled"),
            Type.ArgClause(Nil),
            Term.ArgClause(
              List(
                WasCalledTimes.Times(n)
              ),
              None
            )
          ) =>
        Patch.replaceTree(t, s"Mockito.verify($obj, Mockito.times($n)).${method}${args.mkString("(", ", ", ")")}")
      case t @ Term.ApplyInfix.After_4_6_0(
            Term.Apply.After_4_6_0(
              Term.Apply.After_4_6_0(
                Term.Select(
                  obj,
                  method
                ),
                Term.ArgClause(
                  args1,
                  None
                )
              ),
              Term.ArgClause(
                args2,
                None
              )
            ),
            Term.Name("wasCalled"),
            Type.ArgClause(Nil),
            Term.ArgClause(
              List(
                WasCalledTimes.Times(n)
              ),
              None
            )
          ) =>
        Patch.replaceTree(
          t,
          s"Mockito.verify($obj, Mockito.times($n)).${method}${args1.mkString("(", ", ", ")")}${args2.mkString("(", ", ", ")")}"
        )
    }

    Option
      .when(result.nonEmpty) {
        Seq(
          result.asPatch,
          Option
            .when(doc.tree.collect {
              case importer"org.mockito.Mockito" =>
                ()
              case Import(
                    List(
                      Importer(
                        Term.Select(
                          Term.Name("org"),
                          Term.Name("mockito")
                        ),
                        xs
                      )
                    )
                  ) if xs.collect { case Importee.Name(Name("Mockito")) => () }.nonEmpty =>
                ()
            }.isEmpty) {
              Patch.addGlobalImport(importer"org.mockito.Mockito")
            }
            .asPatch
        ).asPatch
      }
      .asPatch
  }
}

WasNever.scala

package fix

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

class WasNever extends SyntacticRule("WasNever") {
  override def fix(implicit doc: SyntacticDocument): Patch = {
    val result = doc.tree.collect {
      case t @ Term.ApplyInfix.After_4_6_0(
            Term.Apply.After_4_6_0(
              Term.Apply.After_4_6_0(
                Term.Select(
                  obj,
                  method
                ),
                args1
              ),
              args2
            ),
            Term.Name("wasNever"),
            Type.ArgClause(Nil),
            Term.ArgClause(
              List(
                Term.Name("called")
              ),
              None
            )
          ) =>
        Seq(
          Patch.replaceTree(t, s"Mockito.verify($obj, Mockito.never()).${method}${args1}${args2}")
        ).asPatch
      case t @ Term.ApplyInfix.After_4_6_0(
            Term.Select(
              obj,
              method
            ),
            Term.Name("wasNever"),
            Type.ArgClause(Nil),
            Term.ArgClause(
              List(
                Term.Name("called")
              ),
              None
            )
          ) =>
        Patch.replaceTree(t, s"Mockito.verify($obj, Mockito.never()).${method}")
      case t @ Term.ApplyInfix.After_4_6_0(
            Term.Apply.After_4_6_0(
              Term.Select(
                obj,
                method
              ),
              args
            ),
            Term.Name("wasNever"),
            Type.ArgClause(Nil),
            Term.ArgClause(
              List(
                Term.Name("called")
              ),
              None
            )
          ) =>
        Seq(
          Patch.replaceTree(t, s"Mockito.verify($obj, Mockito.never()).${method}${args}"),
          Patch.removeTokens(args.tokens)
        ).asPatch
    }
    Option
      .when(result.nonEmpty) {
        Seq(
          result.asPatch,
          Option
            .when(doc.tree.collect {
              case importer"org.mockito.Mockito" =>
                ()
              case Import(
                    List(
                      Importer(
                        Term.Select(
                          Term.Name("org"),
                          Term.Name("mockito")
                        ),
                        xs
                      )
                    )
                  ) if xs.collect { case Importee.Name(Name("Mockito")) => () }.nonEmpty =>
                ()
            }.isEmpty) {
              Patch.addGlobalImport(importer"org.mockito.Mockito")
            }
            .asPatch
        ).asPatch
      }
      .asPatch
  }
}