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/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) ⇒
+        case OnNext(n) if f.isDefinedAt(n) ⇒ f(n)
-          assert(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.

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)