Add func MatchReturnResponses

Author: Ryan O'Hara-Reid

exchanges/stream/stream_types.go

@@ -35,6 +35,10 @@ type Connection interface {
 	SetProxy(string)
 	GetURL() string
 	Shutdown() error
+	// MatchReturnResponses sets up a channel to listen for an expected number of responses. This is used for when a
+	// request is sent and a response is expected in a different connection. Please see implementation in
+	// websocket_connection.go
+	MatchReturnResponses(ctx context.Context, signature any, expected int) (<-chan MatchedResponse, error)
 }
 
 // Inspector is a hook that allows for custom message inspection

exchanges/stream/websocket_connection.go

@@ -366,6 +366,32 @@ func (w *WebsocketConnection) waitForResponses(ctx context.Context, signature an
 	return resps, nil
 }
 
+// MatchedResponse encapsulates the matched responses along with any errors encountered.
+type MatchedResponse struct {
+	Responses [][]byte
+	Err       error
+}
+
+// MatchReturnResponses sets up a channel to listen for an expected number of responses. These responses may not
+// originate from the same connection as the request, but can come from an alternative connection. It returns a channel
+// that will receive a MatchedResponse containing the collected responses or an error.
+func (w *WebsocketConnection) MatchReturnResponses(ctx context.Context, signature any, expected int) (<-chan MatchedResponse, error) {
+	connectionListen, err := w.Match.Set(signature, expected)
+	if err != nil {
+		return nil, err
+	}
+
+	out := make(chan MatchedResponse, 1) // buffered so routine below doesn't leak
+
+	go func() {
+		resps, err := w.waitForResponses(ctx, signature, connectionListen, expected)
+		out <- MatchedResponse{Responses: resps, Err: err}
+		close(out)
+	}()
+
+	return out, nil
+}
+
 func removeURLQueryString(url string) string {
 	if index := strings.Index(url, "?"); index != -1 {
 		return url[:index]

exchanges/stream/websocket_test.go

@@ -1525,3 +1525,24 @@ func TestGetConnection(t *testing.T) {
 	require.NoError(t, err)
 	assert.Same(t, expected, conn)
 }
+
+func TestMatchReturnResponses(t *testing.T) {
+	t.Parallel()
+
+	conn := WebsocketConnection{Match: NewMatch()}
+	_, err := conn.MatchReturnResponses(context.Background(), nil, 0)
+	require.ErrorIs(t, err, errInvalidBufferSize)
+
+	ch, err := conn.MatchReturnResponses(context.Background(), nil, 1)
+	require.NoError(t, err)
+
+	require.ErrorIs(t, (<-ch).Err, ErrSignatureTimeout)
+	conn.ResponseMaxLimit = time.Millisecond
+
+	ch, err = conn.MatchReturnResponses(context.Background(), nil, 1)
+	require.NoError(t, err)
+
+	exp := []byte("test")
+	require.True(t, conn.Match.IncomingWithData(nil, exp))
+	require.Equal(t, (<-ch).Responses[0], exp)
+}