Allow receiving all messages with a non-null 'status'

src/main/kotlin/org/phoenixframework/Push.kt

@@ -47,6 +47,9 @@ class Push(
   /** Hooks into a Push. Where .receive("ok", callback(Payload)) are stored */
   var receiveHooks: MutableMap<String, List<((message: Message) -> Unit)>> = HashMap()
 
+  /** Hooks into a Push. Where .receiveAll(callback(status, message)) are stored */
+  private var receiveAllHooks: MutableList<(status: String, message: Message) -> Unit> = mutableListOf()
+
   /** True if the Push has been sent */
   var sent: Boolean = false
 
@@ -117,6 +120,28 @@ class Push(
     return this
   }
 
+  /**
+   * Receives any event that was a response to an outbound message.
+   *
+   * Example:
+   *    channel
+   *        .send("event", mPayload)
+   *        .receive { status, message ->
+   *            print(status) // "ok"
+   *        }
+   */
+  fun receive(callback: (status: String, message: Message) -> Unit): Push {
+    // If the message has already been received, pass it to the callback.
+    receivedMessage?.let {
+      val status = it.status
+      if (status != null) {
+        callback(status, it)
+      }
+    }
+    receiveAllHooks.add(callback)
+    return this
+  }
+
   //------------------------------------------------------------------------------
   // Internal
   //------------------------------------------------------------------------------
@@ -182,6 +207,7 @@ class Push(
    */
   private fun matchReceive(status: String, message: Message) {
     receiveHooks[status]?.forEach { it(message) }
+    receiveAllHooks.forEach { it(status, message) }
   }
 
   /** Removes receive hook from Channel regarding this Push */

src/test/kotlin/org/phoenixframework/ChannelTest.kt

@@ -28,6 +28,7 @@ class ChannelTest {
 
   @Mock lateinit var socket: Socket
   @Mock lateinit var mockCallback: ((Message) -> Unit)
+  @Mock lateinit var mockStatusCallback: ((String, Message) -> Unit)
 
   private val kDefaultRef = "1"
   private val kDefaultTimeout = 10_000L
@@ -410,6 +411,11 @@ class ChannelTest {
       joinPush.trigger("error", mapOf("a" to "b"))
     }
 
+    private fun receivesApproved() {
+      fakeClock.tick(joinPush.timeout / 2)
+      joinPush.trigger("approved", mapOf("a" to "b"))
+    }
+
     @Nested
     @DisplayName("receives 'ok'")
     inner class ReceivesOk {
@@ -652,6 +658,52 @@ class ChannelTest {
       /* End ReceivesError */
     }
 
+
+    @Nested
+    @DisplayName("receives 'all status'")
+    inner class ReceivesAllStatus {
+      @Test
+      internal fun `triggers receive('error') callback after error response`() {
+        assertThat(channel.state).isEqualTo(Channel.State.JOINING)
+        joinPush.receive(mockStatusCallback)
+
+        receivesError()
+        joinPush.trigger("error", kEmptyPayload)
+        verify(mockStatusCallback, times(1)).invoke(any(), any())
+      }
+
+      @Test
+      internal fun `triggers receive('error') callback if error response already received`() {
+        receivesError()
+
+        joinPush.receive(mockStatusCallback)
+
+        verify(mockStatusCallback).invoke(any(), any())
+      }
+
+      @Test
+      internal fun `triggers receive('approved') callback after approved response`() {
+        assertThat(channel.state).isEqualTo(Channel.State.JOINING)
+        joinPush.receive(mockStatusCallback)
+
+        receivesApproved()
+        joinPush.trigger("approved", kEmptyPayload)
+        verify(mockStatusCallback, times(1)).invoke(any(), any())
+
+      }
+
+      @Test
+      internal fun `triggers receive('approved') callback if approved response already received`() {
+        receivesApproved()
+
+        joinPush.receive(mockStatusCallback)
+
+        verify(mockStatusCallback).invoke(any(), any())
+      }
+
+      /* End ReceivesAllStatus */
+    }
+
     /* End JoinPush */
   }