Generator for test data

Generator.scala

import scala.util.Random

trait Generator[A] extends Function0[A] {
  val random = new Random

  def apply: A
}

object Generator {

  def of[A](as: A*): Generator[A] = {
    new Generator[A] {
      def apply: A = {
        val next = random.nextInt(as.length)
        as(next)
      }
    }
  }


  def withProbability[A](as: (A, Double)*): Generator[A] = {
    new Generator[A] {
      val normals = {
        val total = as.foldLeft(0.0) { case (acc, (tuple)) => acc + tuple._2 }
        var runningTotal = 0.0
        as.map { tuple =>
          runningTotal += tuple._2
          (runningTotal / total, tuple._1)
        }
      }

      def apply: A = {
        val value = random.nextDouble
        normals.dropWhile(value > _._1).head._2
      }
    }
  }
}