fix bug add test

exchanges/stream/websocket_connection.go

@@ -344,14 +344,15 @@ func (w *WebsocketConnection) waitForResponses(ctx context.Context, signature an
 	defer timeout.Stop()
 
 	resps := make([][]byte, 0, expected)
+inspection:
 	for range expected {
 		select {
 		case resp := <-ch:
 			resps = append(resps, resp)
 			// Checks recently received message to determine if this is in fact the final message in a sequence of messages.
 			if messageInspector != nil && messageInspector.IsFinal(resp) {
 				w.Match.RemoveSignature(signature)
-				return resps, nil
+				break inspection
 			}
 		case <-timeout.C:
 			w.Match.RemoveSignature(signature)

exchanges/stream/websocket_test.go

@@ -769,9 +769,36 @@ func TestSendMessageReturnResponse(t *testing.T) {
 	assert.ErrorIs(t, err, ErrSignatureTimeout, "SendMessageReturnResponse should error when request ID not found")
 }
 
-type inspection struct{}
+func TestWaitForResponses(t *testing.T) {
+	t.Parallel()
+	dummy := &WebsocketConnection{
+		ResponseMaxLimit: time.Nanosecond,
+		Match:            NewMatch(),
+	}
+	_, err := dummy.waitForResponses(context.Background(), "silly", nil, 1, inspection{})
+	require.ErrorIs(t, err, ErrSignatureTimeout)
+
+	ctx, cancel := context.WithCancel(context.Background())
+	cancel()
+	_, err = dummy.waitForResponses(ctx, "silly", nil, 1, inspection{})
+	require.ErrorIs(t, err, context.Canceled)
+
+	// test break early and hit verbose path
+	ch := make(chan []byte, 1)
+	ch <- []byte("hello")
+	ctx = request.WithVerbose(context.Background())
+	dummy.ResponseMaxLimit = time.Second
+	got, err := dummy.waitForResponses(ctx, "silly", ch, 2, inspection{breakEarly: true})
+	require.NoError(t, err, context.Canceled)
+	require.Len(t, got, 1)
+	assert.Equal(t, "hello", string(got[0]))
+}
+
+type inspection struct {
+	breakEarly bool
+}
 
-func (inspection) IsFinal([]byte) bool { return false }
+func (i inspection) IsFinal([]byte) bool { return i.breakEarly }
 
 type reporter struct {
 	name string