fixup

vici/session.go

@@ -274,13 +274,18 @@ func (s *Session) CallStreaming(ctx context.Context, cmd string, event string, i
 // to stream while the command request is active. The complete stream of messages received from
 // the server is returned once the request is complete.
 func (s *Session) StreamedCommandRequest(cmd string, event string, msg *Message) ([]*Message, error) {
-	s.mu.Lock()
-	defer s.mu.Unlock()
-	if s.cc == nil {
-		return nil, errors.New("session closed")
+	resp, err := s.CallStreaming(context.Background(), cmd, event, msg)
+	if err != nil {
+		return nil, err
 	}
 
-	return s.streamedRequest(context.Background(), cmd, event, msg)
+	messages := make([]*Message, 0)
+
+	for m := range resp {
+		messages = append(messages, m)
+	}
+
+	return messages, nil
 }
 
 // Subscribe registers the session to listen for all events given. To receive
@@ -345,37 +350,6 @@ func (s *Session) request(ctx context.Context, cmd string, in *Message) (*Messag
 	return p.msg, nil
 }
 
-func (s *Session) streamedRequest(ctx context.Context, cmd string, event string, in *Message) ([]*Message, error) {
-	if err := s.eventRegister(ctx, event); err != nil {
-		return nil, err
-	}
-	// nolint
-	defer s.eventUnregister(ctx, event)
-
-	if err := s.cc.packetWrite(ctx, newPacket(pktCmdRequest, cmd, in)); err != nil {
-		return nil, err
-	}
-
-	messages := make([]*Message, 0)
-	for {
-		p, err := s.cc.packetRead(ctx)
-		if err != nil {
-			return nil, err
-		}
-
-		switch p.ptype {
-		case pktEvent:
-			messages = append(messages, p.msg)
-		case pktCmdResponse:
-			// End of event stream
-			messages = append(messages, p.msg)
-			return messages, nil
-		default:
-			return nil, fmt.Errorf("%v: %v", errUnexpectedResponse, p.ptype)
-		}
-	}
-}
-
 func (s *Session) eventRequest(ctx context.Context, ptype uint8, event string) error {
 	p := newPacket(ptype, event, nil)