1 files changed, 33 insertions, 0 deletions
diff --git a/src/test/scala/xyz/driver/core/rest/RestTest.scala b/src/test/scala/xyz/driver/core/rest/RestTest.scala
index 68fe419..5403765 100644
--- a/src/test/scala/xyz/driver/core/rest/RestTest.scala
+++ b/src/test/scala/xyz/driver/core/rest/RestTest.scala
@@ -1,5 +1,6 @@
 package xyz.driver.core.rest
 
+import akka.http.scaladsl.marshalling.ToResponseMarshallable
 import akka.http.scaladsl.model.StatusCodes
 import akka.http.scaladsl.server.{Directives, Route, ValidationRejection}
 import akka.http.scaladsl.testkit.ScalatestRouteTest
@@ -7,6 +8,9 @@ import akka.util.ByteString
 import org.scalatest.{Matchers, WordSpec}
 import xyz.driver.core.rest
 
+import scala.concurrent.Future
+import scala.util.Random
+
 class RestTest extends WordSpec with Matchers with ScalatestRouteTest with Directives {
   "`escapeScriptTags` function" should {
     "escape script tags properly" in {
@@ -70,4 +74,33 @@ class RestTest extends WordSpec with Matchers with ScalatestRouteTest with Direc
       }
     }
   }
+
+  "completeWithPagination directive" should {
+    import akka.http.scaladsl.marshallers.sprayjson.SprayJsonSupport._
+    import spray.json.DefaultJsonProtocol._
+
+    val data = Seq.fill(103)(Random.alphanumeric.take(10).mkString)
+    val route: Route = rest.completeWithPagination[String] {
+      case Some(pagination) =>
+        val filtered = data.slice(pagination.offset, pagination.offset + pagination.pageSize)
+        Future.successful(ListResponse(filtered, data.size, Some(pagination)))
+      case None => Future.successful(ListResponse(data, data.size, None))
+    }
+
+    "return a response with pagination headers when pagination has been passed" in {
+      Get("/?pageNumber=2&pageSize=10") ~> route ~> check {
+        responseAs[Seq[String]] shouldBe data.slice(10, 20)
+        header(ContextHeaders.ResourceCount).map(_.value) should contain("103")
+        header(ContextHeaders.PageCount).map(_.value) should contain("11")
+      }
+    }
+
+    "return a response with pagination headers when no pagination has been passed" in {
+      Get("/") ~> route ~> check {
+        responseAs[Seq[String]] shouldBe data
+        header(ContextHeaders.ResourceCount).map(_.value) should contain("103")
+        header(ContextHeaders.PageCount).map(_.value) should contain("1")
+      }
+    }
+  }
 }