+testKit Add fishForMessages allowing selective expecting of msgs (#399)

* +testKit Add fishForMessages allowing selective expecting of msgs

* Update Sources/DistributedActorsTestKit/TestProbes.swift

Co-Authored-By: Yim Lee <[email protected]>

Co-authored-by: Yim Lee <[email protected]>

Author: ktoso

Sources/DistributedActorsTestKit/TestProbes.swift

@@ -194,6 +194,82 @@ extension ActorTestProbe {
 
         throw ExpectationError.noMessagesInQueue
     }
+
+    /// Allows for "fishing out" certain messages from the stream of incoming messages to this probe.
+    /// Messages can be caught or ignored using the passed in function.
+    ///
+    /// Allows transforming the caught messages to `CaughtMessage` (e.g. extracting a specific payload from all incoming messages).
+    /// If you need to aggregate the exact `Message` types, you prefer using `fishForMessages`.
+    ///
+    /// Once `MessageFishingDirective.catchComplete` or `MessageFishingDirective.complete` is returned,
+    /// the function returns all so-far accumulated messages.
+    public func fishFor<CaughtMessage>(
+        _ type: CaughtMessage.Type, within timeout: TimeAmount,
+        file: StaticString = #file, line: UInt = #line, column: UInt = #column,
+        _ fisher: (Message) throws -> FishingDirective<CaughtMessage>
+    ) throws -> [CaughtMessage] {
+        let deadline = Deadline.fromNow(timeout)
+
+        var caughtMessages: [CaughtMessage] = []
+        while deadline.hasTimeLeft() {
+            if let message = try self.maybeExpectMessage(within: deadline.timeLeft) {
+                switch try fisher(message) {
+                case .catchContinue(let caught):
+                    caughtMessages.append(caught)
+                case .catchComplete(let caught):
+                    caughtMessages.append(caught)
+                    return caughtMessages
+                case .ignore:
+                    ()
+                case .complete:
+                    return caughtMessages
+                }
+            }
+        }
+
+        return caughtMessages
+    }
+
+    public enum FishingDirective<CaughtMessage> {
+        case catchContinue(CaughtMessage)
+        case catchComplete(CaughtMessage)
+        case ignore
+        case complete
+    }
+
+    /// Allows for "fishing out" certain messages from the stream of incoming messages to this probe.
+    /// Messages can be caught or ignored using the passed in function.
+    ///
+    /// The accumulated messages are assumed to be transforming the caught messages to `CaughtMessage` (e.g. extracting a specific payload from all incoming messages).
+    /// If you need to aggregate the exact `Message` types, you prefer using `fishForMessages`.
+    ///
+    /// Once `MessageFishingDirective.catchComplete` or `MessageFishingDirective.complete` is returned,
+    /// the function returns all so-far accumulated messages.
+    public func fishForMessages(
+        within timeout: TimeAmount,
+        file: StaticString = #file, line: UInt = #line, column: UInt = #column,
+        _ fisher: (Message) throws -> MessageFishingDirective
+    ) throws -> [Message] {
+        try self.fishFor(Message.self, within: timeout, file: file, line: line, column: column) { message in
+            switch try fisher(message) {
+            case .catchContinue:
+                return .catchContinue(message)
+            case .catchComplete:
+                return .catchComplete(message)
+            case .ignore:
+                return .ignore
+            case .complete:
+                return .complete
+            }
+        }
+    }
+
+    public enum MessageFishingDirective {
+        case catchContinue
+        case catchComplete
+        case ignore
+        case complete
+    }
 }
 
 extension ActorTestProbe {

Tests/DistributedActorsTestKitTests/ActorTestKitTests+XCTest.swift

@@ -28,6 +28,9 @@ extension ActorTestKitTests {
             ("test_fail_shouldNotImmediatelyFailWithinEventuallyBlock", test_fail_shouldNotImmediatelyFailWithinEventuallyBlock),
             ("test_nestedEventually_shouldProperlyHandleFailures", test_nestedEventually_shouldProperlyHandleFailures),
             ("test_ensureRegistered_countAndRefs", test_ensureRegistered_countAndRefs),
+            ("test_fishForMessages", test_fishForMessages),
+            ("test_fishForTransformed", test_fishForTransformed),
+            ("test_fishFor_canThrow", test_fishFor_canThrow),
             ("test_ActorableTestProbe_shouldWork", test_ActorableTestProbe_shouldWork),
         ]
     }

Tests/DistributedActorsTestKitTests/ActorTestKitTests.swift

@@ -86,6 +86,77 @@ final class ActorTestKitTests: XCTestCase {
         try self.testKit.ensureRegistered(key: key, expectedCount: 2, expectedRefs: [greeterProbe2.ref, greeterProbe1.ref])
     }
 
+    func test_fishForMessages() throws {
+        let p = self.testKit.spawnTestProbe(expecting: String.self)
+
+        p.tell("yes-1")
+        p.tell("yes-2")
+        p.tell("no-1")
+        p.tell("no-2")
+        p.tell("yes-3")
+        p.tell("yes-end")
+
+        let messages = try p.fishForMessages(within: .seconds(30)) { message in
+            if message.contains("yes-end") {
+                return .catchComplete
+            } else if message.contains("yes") {
+                return .catchContinue
+            } else {
+                return .ignore
+            }
+        }
+
+        messages.shouldEqual([
+            "yes-1",
+            "yes-2",
+            "yes-3",
+            "yes-end",
+        ])
+    }
+
+    func test_fishForTransformed() throws {
+        let p = self.testKit.spawnTestProbe(expecting: String.self)
+
+        p.tell("yes-1")
+        p.tell("yes-2")
+        p.tell("no-1")
+        p.tell("no-2")
+        p.tell("yes-3")
+        p.tell("yes-end")
+
+        let messages = try p.fishFor(String.self, within: .seconds(30)) { message in
+            if message.contains("yes-end") {
+                return .catchComplete("\(message)!!!")
+            } else if message.contains("yes") {
+                return .catchContinue("\(message)!!!")
+            } else {
+                return .ignore
+            }
+        }
+
+        messages.shouldEqual([
+            "yes-1!!!",
+            "yes-2!!!",
+            "yes-3!!!",
+            "yes-end!!!",
+        ])
+    }
+
+    func test_fishFor_canThrow() throws {
+        let p = self.testKit.spawnTestProbe(expecting: String.self)
+
+        p.tell("yes-1")
+
+        do {
+            _ = try p.fishForMessages(within: .seconds(30)) { message in
+                throw TestError("Boom: \(message)")
+            }
+            throw self.testKit.fail("Should have thrown")
+        } catch {
+            "\(error)".shouldContain("Boom: yes-1")
+        }
+    }
+
     // ==== ------------------------------------------------------------------------------------------------------------
     // MARK: Actorable