expectNextPF with timeout #21314 (#21825)

* expectNextPF with timeout #21314

improved scaladocs

* renamed functions with timeout, returned previous comments for the sake of unification

akka-stream-testkit/src/main/scala/akka/stream/testkit/StreamTestKit.scala

@@ -551,23 +551,46 @@ object TestSubscriber {
       self
     }
 
-    def expectNextPF[T](f: PartialFunction[Any, T]): T = {
-      expectEventPF {
-        case OnNext(n) if f.isDefinedAt(n) ⇒ f(n)
-      }
-    }
+    /**
+     * Expect a stream element and test it with partial function.
+     *
+     */
+    def expectNextPF[T](f: PartialFunction[Any, T]): T =
+      expectNextWithTimeoutPF(Duration.Undefined, f)
 
     /**
-     * Expect next element and test it with partial function.
+     * Expect a stream element and test it with partial function.
+     *
+     * @param max wait no more than max time, otherwise throw AssertionError
+     */
+    def expectNextWithTimeoutPF[T](max: Duration, f: PartialFunction[Any, T]): T =
+      expectEventWithTimeoutPF(max, {
+        case OnNext(n) if f.isDefinedAt(n) ⇒ f(n)
+      })
+
+    /**
+     * Expect a stream element during specified time or timeout and test it with partial function.
+     *
+     * Allows chaining probe methods.
+     *
+     * @param max wait no more than max time, otherwise throw AssertionError
+     */
+    def expectNextChainingPF(max: Duration, f: PartialFunction[Any, Any]): Self =
+      expectNextWithTimeoutPF(max, f.andThen(_ ⇒ self))
+
+    /**
+     * Expect a stream element during specified time or timeout and test it with partial function.
      *
      * Allows chaining probe methods.
      */
-    def expectNextChainingPF(f: PartialFunction[Any, Any]): Self = {
-      expectNextPF(f.andThen(_ ⇒ self))
-    }
+    def expectNextChainingPF(f: PartialFunction[Any, Any]): Self =
+      expectNextChainingPF(Duration.Undefined, f)
+
+    def expectEventWithTimeoutPF[T](max: Duration, f: PartialFunction[SubscriberEvent, T]): T =
+      probe.expectMsgPF[T](max, hint = "message matching partial function")(f.asInstanceOf[PartialFunction[Any, T]])
 
     def expectEventPF[T](f: PartialFunction[SubscriberEvent, T]): T =
-      probe.expectMsgPF[T](hint = "message matching partial function")(f.asInstanceOf[PartialFunction[Any, T]])
+      expectEventWithTimeoutPF(Duration.Undefined, f)
 
     /**
      * Receive messages for a given duration or until one does not match a given partial function.