Test stub matching partial function mapping request to response

2 files changed, 46 insertions, 8 deletions

diff --git a/core/src/main/scala/com/softwaremill/sttp/testing/SttpBackendStub.scala b/core/src/main/scala/com/softwaremill/sttp/testing/SttpBackendStub.scala
index aa04ef3..72d3490 100644
--- a/core/src/main/scala/com/softwaremill/sttp/testing/SttpBackendStub.scala
+++ b/core/src/main/scala/com/softwaremill/sttp/testing/SttpBackendStub.scala
@@ -31,10 +31,17 @@ class SttpBackendStub[R[_], S] private (rm: MonadError[R],
   def whenRequestMatches(p: Request[_, _] => Boolean): WhenRequest =
     new WhenRequest(p)
 
+  def whenRequestMatchesPartial(
+      partial: PartialFunction[Request[_, _], Response[_]]) = {
+    val m = Matcher(partial)
+    val vector: Vector[Matcher[_]] = matchers :+ m
+    new SttpBackendStub(rm, vector, fallback)
+  }
+
   override def send[T](request: Request[T, S]): R[Response[T]] = {
     matchers
       .collectFirst {
-        case matcher if matcher(request) => matcher.response
+        case matcher: Matcher[T] if matcher(request) => matcher.response(request).get
       } match {
       case Some(response) => wrapResponse(response)
       case None =>
@@ -64,8 +71,11 @@ class SttpBackendStub[R[_], S] private (rm: MonadError[R],
       thenRespond(Response[Nothing](Left(msg), code, Nil, Nil))
     def thenRespond[T](body: T): SttpBackendStub[R, S] =
       thenRespond(Response[T](Right(body), 200, Nil, Nil))
-    def thenRespond[T](resp: Response[T]): SttpBackendStub[R, S] =
-      new SttpBackendStub(rm, matchers :+ Matcher(p, resp), fallback)
+    def thenRespond[T](resp: Response[T]): SttpBackendStub[R, S] = {
+      val m = Matcher[T](p, resp)
+      new SttpBackendStub(rm, matchers :+ m, fallback)
+
+    }
   }
 }
 
@@ -101,9 +111,22 @@ object SttpBackendStub {
   private val DefaultResponse =
     Response[Nothing](Left("Not Found"), 404, Nil, Nil)
 
-  private case class Matcher[T](p: Request[T, _] => Boolean,
-                                response: Response[T]) {
-    def apply(request: Request[_, _]): Boolean =
-      p(request.asInstanceOf[Request[T, _]])
+  private case class Matcher[T](p: PartialFunction[Request[T, _], Response[_]]) {
+
+    def apply(request: Request[T, _]): Boolean =
+      p.isDefinedAt(request)
+
+    def response[S](request: Request[T, S]): Option[Response[T]] = {
+      p.lift(request).asInstanceOf[Option[Response[T]]]
+    }
+  }
+
+  private object Matcher {
+
+    def apply[T](p: Request[T, _] => Boolean, response: Response[T]) = {
+      new Matcher[T]({
+        case r if p(r) => response
+      })
+    }
   }
 }
diff --git a/core/src/test/scala/com/softwaremill/sttp/testing/SttpBackendStubTests.scala b/core/src/test/scala/com/softwaremill/sttp/testing/SttpBackendStubTests.scala
index b80ae76..d10d1b6 100644
--- a/core/src/test/scala/com/softwaremill/sttp/testing/SttpBackendStubTests.scala
+++ b/core/src/test/scala/com/softwaremill/sttp/testing/SttpBackendStubTests.scala
@@ -12,6 +12,10 @@ class SttpBackendStubTests extends FlatSpec with Matchers with ScalaFutures {
     .thenRespond(10)
     .whenRequestMatches(_.method == Method.GET)
     .thenRespondServerError()
+    .whenRequestMatchesPartial({
+      case r if r.method == Method.POST && r.uri.path.endsWith(List("partial10")) => Response(Right(10), 200, Nil, Nil)
+      case r if r.method == Method.POST && r.uri.path.endsWith(List("partialAda")) => Response(Right("Ada"), 200, Nil, Nil)
+    })
 
   "backend stub" should "use the first rule if it matches" in {
     implicit val b = testingStub
@@ -38,7 +42,7 @@ class SttpBackendStubTests extends FlatSpec with Matchers with ScalaFutures {
 
   it should "use the default response if no rule matches" in {
     implicit val b = testingStub
-    val r = sttp.post(uri"http://example.org/d").send()
+    val r = sttp.put(uri"http://example.org/d").send()
     r.code should be(404)
   }
 
@@ -49,6 +53,17 @@ class SttpBackendStubTests extends FlatSpec with Matchers with ScalaFutures {
     r.futureValue.code should be(404)
   }
 
+  it should "use rules in partial function" in {
+    implicit val s = testingStub
+    val r = sttp.post(uri"http://example.org/partial10").send()
+    r.is200 should be(true)
+    r.body should be(Right(10))
+
+    val ada = sttp.post(uri"http://example.org/partialAda").send()
+    ada.is200 should be(true)
+    ada.body should be(Right("Ada"))
+  }
+
   val testingStubWithFallback = SttpBackendStub
     .withFallback(testingStub)
     .whenRequestMatches(_.uri.path.startsWith(List("c")))