ScalaTest support for Doobie

doobie_scalatest.scala

import org.scalatest._

import scala.reflect.runtime.universe.WeakTypeTag
import doobie.imports._
import doobie.contrib.postgresql.imports._
import doobie.contrib.postgresql.pgtypes._
import DoobiePgTypeImplicits._
import doobie.free.connection.{ConnectionIO, setAutoCommit, commit, rollback, close, delay}
import doobie.util.transactor.Transactor
import doobie.util.query.{Query, Query0}
import doobie.util.update.{Update, Update0}
import doobie.util.analysis._
import doobie.util.pretty._

import scalaz.concurrent.Task
import scalaz._
import Scalaz._

trait QueryChecker extends FunSuite {
  def transactor: Transactor[Task]

  def check[A, B](q: Query[A, B])(implicit A: WeakTypeTag[A], B: WeakTypeTag[B]) =
    checkAnalysis(s"Query[${typeName(A)}, ${typeName(B)}]", q.stackFrame, q.sql, q.analysis)

  def check[A](q: Query0[A])(implicit A: WeakTypeTag[A]) =
    checkAnalysis(s"Query0[${typeName(A)}]", q.stackFrame, q.sql, q.analysis)

  def checkOutput[A](q: Query0[A])(implicit A: WeakTypeTag[A]) =
    checkAnalysis(s"Query0[${typeName(A)}]", q.stackFrame, q.sql, q.outputAnalysis)

  def check[A](q: Update[A])(implicit A: WeakTypeTag[A]) =
    checkAnalysis(s"Update[${typeName(A)}]", q.stackFrame, q.sql, q.analysis)

  def check[A](q: Update0)(implicit A: WeakTypeTag[A]) =
    checkAnalysis(s"Update0", q.stackFrame, q.sql, q.analysis)

  /** Check if the analysis has an error */
  private def hasError(analysis: ConnectionIO[Analysis]): Boolean = {
    transactor.trans(analysis).attemptRun match {
      case -\/(e) => true
      case \/-(a) => !(a.paramDescriptions.map { case (s, es) => es.isEmpty } ++ a.columnDescriptions.map { case (s, es) => es.isEmpty }).forall(x => x)
    }
  }

  private def checkAnalysis(typeName: String, stackFrame: Option[StackTraceElement], sql: String, analysis: ConnectionIO[Analysis]) = {
    if (hasError(analysis)) {
      val analysisOutput = transactor.trans(analysis).attemptRun match {
        case -\/(e) =>
          failure("SQL Compiles and Typechecks", formatError(e.getMessage))
        case \/-(a) =>
          success("SQL Compiles and Typechecks", None) +
            a.paramDescriptions.map { case (s, es) => assertEmpty(s, es) }.mkString("\n") +
            a.columnDescriptions.map { case (s, es) => assertEmpty(s, es) }.mkString("\n")
      }

      println(formatSql(sql))
      println(analysisOutput)
      fail()
    }
  }

  private val packagePrefix = "\\b[a-z]+\\.".r

  private def assertEmpty(name: String, es: List[AlignmentError]): String =
    if (es.isEmpty) success(name, None)
    else failure(name, es.map(formatError).mkString("\n"))

  private def typeName[A](tag: WeakTypeTag[A]): String =
    packagePrefix.replaceAllIn(tag.tpe.toString, "")

  private def formatError(e: AlignmentError): String =
    formatError(e.msg)

  private def formatError(s: String): String =
    (wrap(80)(s) match {
      case s :: ss => (s"${Console.RED}  - $s${Console.RESET}") :: ss.map(s => s"${Console.RED}    $s${Console.RESET}")
      case Nil => Nil
    }).mkString("\n")

  def formatSql(sql: String): String = {
    val line = sql.lines.dropWhile(_.trim.isEmpty).map(s => s"  \033[37m$s${Console.RESET}").mkString("")
    s"\n$line\n"
  }

  def failure(name: String, desc: String): String =
    s"${Console.RED}  ✕ ${Console.RESET}$name\n" + desc.lines.map(s => s"  $s").mkString("\n")

  def success(name: String, desc: Option[String]): String =
    s"${Console.GREEN}  ✓ ${Console.RESET}$name\n" + desc.mkString("\n")
}

class QueryTest extends FunSuite with QueryChecker with TestHelpers {
  val transactor = createDoobieTransactor()

  test("check exampleQuery2") {
    case class ExampleQuery2Row(a: Int, b: Vector[Int])
    val exampleQuery2: Query0[ExampleQuery2Row] = dsql"SELECT 'hello'::VARCHAR AS a, array[1, 2, 3]::SMALLINT [] AS b".query[ExampleQuery2Row]
    check(exampleQuery2)
  }
}