From 7eb660b4974acd8500f60c3a92da16f1c8465cea Mon Sep 17 00:00:00 2001
From: Kirill Yankov <myworkpostbox@gmail.com>
Date: Mon, 17 Oct 2016 18:08:09 +0300
Subject: [PATCH] Add expectNextChainingPF method  to streams testkit #21614
 (#21617)

Add expectNextChainingPF method to streams testkit #21614

Improved error messages on expectNextPf, expectNextPfChaining, expectEventPF
---
 .../akka/stream/testkit/StreamTestKit.scala   | 15 +++++--
 .../stream/testkit/StreamTestKitSpec.scala    | 45 +++++++++++++++++++
 2 files changed, 56 insertions(+), 4 deletions(-)

diff --git a/akka-stream-testkit/src/main/scala/akka/stream/testkit/StreamTestKit.scala b/akka-stream-testkit/src/main/scala/akka/stream/testkit/StreamTestKit.scala
index 1c84c1d2af..3f088f5100 100644
--- a/akka-stream-testkit/src/main/scala/akka/stream/testkit/StreamTestKit.scala
+++ b/akka-stream-testkit/src/main/scala/akka/stream/testkit/StreamTestKit.scala
@@ -530,14 +530,21 @@ object TestSubscriber {
 
     def expectNextPF[T](f: PartialFunction[Any, T]): T = {
       expectEventPF {
-        case OnNext(n) ⇒
-          assert(f.isDefinedAt(n))
-          f(n)
+        case OnNext(n) if f.isDefinedAt(n) ⇒ f(n)
       }
     }
 
+    /**
+     * Expect next element and test it with partial function.
+     *
+     * Allows chaining probe methods.
+     */
+    def expectNextChainingPF(f: PartialFunction[Any, Any]): Self = {
+      expectNextPF(f.andThen(_ ⇒ self))
+    }
+
     def expectEventPF[T](f: PartialFunction[SubscriberEvent, T]): T =
-      probe.expectMsgPF[T]()(f.asInstanceOf[PartialFunction[Any, T]])
+      probe.expectMsgPF[T](hint = "message matching partial function")(f.asInstanceOf[PartialFunction[Any, T]])
 
     /**
      * Receive messages for a given duration or until one does not match a given partial function.
diff --git a/akka-stream-testkit/src/test/scala/akka/stream/testkit/StreamTestKitSpec.scala b/akka-stream-testkit/src/test/scala/akka/stream/testkit/StreamTestKitSpec.scala
index 11320644a4..2d22b2feb4 100644
--- a/akka-stream-testkit/src/test/scala/akka/stream/testkit/StreamTestKitSpec.scala
+++ b/akka-stream-testkit/src/test/scala/akka/stream/testkit/StreamTestKitSpec.scala
@@ -110,6 +110,51 @@ class StreamTestKitSpec extends AkkaSpec {
         .expectNextOrComplete(1337)
     }
 
+    "#expectNextPF should pass with right element" in {
+      val result = Source.single(1).runWith(TestSink.probe)
+        .request(1)
+        .expectNextPF {
+          case 1 ⇒ "success"
+        }
+      result should be("success")
+    }
+
+    "#expectNextPF should fail with wrong element" in {
+      intercept[AssertionError] {
+        Source.single(1).runWith(TestSink.probe)
+          .request(1)
+          .expectNextPF {
+            case 2 ⇒
+          }
+      }.getMessage should include("message matching partial function")
+    }
+
+    "#expectNextChainingPF should pass with right element" in {
+      Source.single(1).runWith(TestSink.probe)
+        .request(1)
+        .expectNextChainingPF {
+          case 1 ⇒
+        }
+    }
+
+    "#expectNextChainingPF should allow to chain test methods" in {
+      Source(1 to 2).runWith(TestSink.probe)
+        .request(2)
+        .expectNextChainingPF {
+          case 1 ⇒
+        }.expectNext(2)
+    }
+
+    "#expectNextChainingPF should fail with wrong element" in {
+      intercept[AssertionError] {
+        Source.single(1).runWith(TestSink.probe)
+          .request(1)
+          .expectNextChainingPF {
+            case 2 ⇒
+          }
+      }.getMessage should include("message matching partial function")
+    }
+
     "#expectNextN given a number of elements" in {
       Source(1 to 4).runWith(TestSink.probe)
         .request(4)