stream test

antonzherdev · antonzherdev · commit 05cd33b4d18e · 2015-08-06T20:38:23.000+12:00
diff --git a/postgresql-async/src/test/scala/com/github/mauricio/async/db/postgresql/DatabaseTestHelper.scala b/postgresql-async/src/test/scala/com/github/mauricio/async/db/postgresql/DatabaseTestHelper.scala
@@ -19,10 +19,11 @@ package com.github.mauricio.async.db.postgresql
 import java.util.concurrent.{TimeUnit, TimeoutException}
 
 import com.github.mauricio.async.db.util.Log
-import com.github.mauricio.async.db.{Configuration, Connection}
+import com.github.mauricio.async.db.{RowData, Configuration, Connection}
+import org.reactivestreams.{Subscription, Subscriber}
 
 import scala.concurrent.duration._
-import scala.concurrent.{Await, Future}
+import scala.concurrent.{Promise, Await, Future}
 
 object DatabaseTestHelper {
   val log = Log.get[DatabaseTestHelper]
@@ -105,9 +106,44 @@ trait DatabaseTestHelper {
     } )
   }
 
+  def executeStream(handler: PostgreSQLConnection, statement: String, windowSize : Int = 1000, values: Array[Any] = Array.empty[Any]) : IndexedSeq[RowData] = {
+    handleTimeout(handler, {
+      val subscriber: TestSubscriber = new TestSubscriber
+      handler.streamQuery(statement, windowSize).subscribe(subscriber)
+      Await.result(subscriber.promise.future, Duration(5, SECONDS))
+    })
+  }
+
   def await[T](future: Future[T]): T = {
     Await.result(future, Duration(10, TimeUnit.SECONDS))
   }
 
+  class TestSubscriber extends Subscriber[RowData] {
+    val promise = Promise[IndexedSeq[RowData]]()
+    override def onError(t: Throwable): Unit = {
+      promise.failure(t)
+    }
+    
+    var subscription: Option[Subscription] = None
+    var requested : Long = 0
+    override def onSubscribe(subscription: Subscription): Unit = {
+      this.subscription = Some(subscription)
+      requested = 10
+      subscription.request(10)
+    }
+
+    override def onComplete(): Unit = {
+      promise.success(rows)
+      requested -= 1
+      if (requested <= 2) {
+        subscription.get.request(8)
+        requested += 8
+      }
+    }
 
+    var rows = IndexedSeq[RowData]()
+    override def onNext(t: RowData): Unit = {
+      rows = rows :+ t
+    }
+  }
 }
diff --git a/postgresql-async/src/test/scala/com/github/mauricio/async/db/postgresql/PostgreSQLConnectionSpec.scala b/postgresql-async/src/test/scala/com/github/mauricio/async/db/postgresql/PostgreSQLConnectionSpec.scala
@@ -196,6 +196,26 @@ class PostgreSQLConnectionSpec extends Specification with DatabaseTestHelper {
 
     }
 
+    "stream a statement" in {
+      withHandler {
+        handler =>
+          executeDdl(handler, this.preparedStatementCreate)
+          executeDdl(handler, this.preparedStatementInsert, 1)
+          executeDdl(handler, this.preparedStatementInsert2, 1)
+          val result = executeStream(handler, this.preparedStatementSelect)
+
+          result.size must_== 2
+
+          val row = result.head
+          row(0) === 1
+          row(1) === "John Doe"
+
+          val row2 = result(1)
+          row2(0) === 2
+          row2(1) === "Mary Jane"
+      }
+    }
+
     "execute a prepared statement with parameters" in {
 
       withHandler {
