client_stream_test.go

// Copyright 2021-2025 The Connect Authors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package connect

import (
	"errors"
	"fmt"
	"net/http"
	"testing"

	"connectrpc.com/connect/internal/assert"
	pingv1 "connectrpc.com/connect/internal/gen/connect/ping/v1"
)

const expectedStreamErrorMessage = "no stream initialized"

func TestClientStreamForClient_InitErrNoPanics(t *testing.T) {
	t.Parallel()
	initErr := errors.New("client init failure")
	clientStream := &ClientStreamForClient[pingv1.PingRequest, pingv1.PingResponse]{err: initErr}
	assert.ErrorIs(t, clientStream.Send(&pingv1.PingRequest{}), initErr)
	verifyHeaders(t, clientStream.RequestHeader())
	res, err := clientStream.CloseAndReceive()
	assert.Nil(t, res)
	assert.ErrorIs(t, err, initErr)
	conn, err := clientStream.Conn()
	assert.NotNil(t, err)
	assert.Nil(t, conn)
}

func TestClientStreamForClientSimple_InitErrNoPanics(t *testing.T) {
	t.Parallel()
	initErr := errors.New("client init failure")
	clientStream := &ClientStreamForClientSimple[pingv1.PingRequest, pingv1.PingResponse]{
		stream: &ClientStreamForClient[pingv1.PingRequest, pingv1.PingResponse]{err: initErr},
	}
	assert.ErrorIs(t, clientStream.Send(&pingv1.PingRequest{}), initErr)
	res, err := clientStream.CloseAndReceive()
	assert.Nil(t, res)
	assert.ErrorIs(t, err, initErr)
	assert.NotNil(t, err)
}

func TestClientStreamForClientSimple_NilStreamNoPanics(t *testing.T) {
	t.Parallel()
	clientStream := &ClientStreamForClientSimple[pingv1.PingRequest, pingv1.PingResponse]{}
	// Should not panic
	clientStream.Peer()
	clientStream.Spec()
	err := clientStream.Send(&pingv1.PingRequest{})
	assert.NotNil(t, err)
	assert.Equal(t, err.Error(), expectedStreamErrorMessage)
	res, err := clientStream.CloseAndReceive()
	assert.Nil(t, res)
	assert.NotNil(t, err)
	assert.Equal(t, err.Error(), expectedStreamErrorMessage)
}

func TestServerStreamForClient_InitErrNoPanics(t *testing.T) {
	t.Parallel()
	initErr := errors.New("client init failure")
	serverStream := &ServerStreamForClient[pingv1.PingResponse]{constructErr: initErr}
	assert.ErrorIs(t, serverStream.Err(), initErr)
	assert.ErrorIs(t, serverStream.Close(), initErr)
	assert.NotNil(t, serverStream.Msg())
	assert.False(t, serverStream.Receive())
	verifyHeaders(t, serverStream.ResponseHeader())
	verifyHeaders(t, serverStream.ResponseTrailer())
	conn, err := serverStream.Conn()
	assert.NotNil(t, err)
	assert.Nil(t, conn)
}

func TestServerStreamForClient(t *testing.T) {
	t.Parallel()
	stream := &ServerStreamForClient[pingv1.PingResponse]{
		conn: &nopStreamingClientConn{},
	}
	// Ensure that each call to Receive allocates a new message. This helps
	// vtprotobuf, which doesn't automatically zero messages before unmarshaling
	// (see https://connectrpc.com/connect/issues/345), and it's also
	// less error-prone for users.
	assert.True(t, stream.Receive())
	first := fmt.Sprintf("%p", stream.Msg())
	assert.True(t, stream.Receive())
	second := fmt.Sprintf("%p", stream.Msg())
	assert.NotEqual(t, first, second)
	conn, err := stream.Conn()
	assert.Nil(t, err)
	assert.NotNil(t, conn)
}

func TestBidiStreamForClient_InitErrNoPanics(t *testing.T) {
	t.Parallel()
	initErr := errors.New("client init failure")
	bidiStream := &BidiStreamForClient[pingv1.CumSumRequest, pingv1.CumSumResponse]{err: initErr}
	res, err := bidiStream.Receive()
	assert.Nil(t, res)
	assert.ErrorIs(t, err, initErr)
	verifyHeaders(t, bidiStream.RequestHeader())
	verifyHeaders(t, bidiStream.ResponseHeader())
	verifyHeaders(t, bidiStream.ResponseTrailer())
	assert.ErrorIs(t, bidiStream.Send(&pingv1.CumSumRequest{}), initErr)
	assert.ErrorIs(t, bidiStream.CloseRequest(), initErr)
	assert.ErrorIs(t, bidiStream.CloseResponse(), initErr)
	conn, err := bidiStream.Conn()
	assert.NotNil(t, err)
	assert.Nil(t, conn)
}

func TestBidiStreamForClientSimple_InitErrNoPanics(t *testing.T) {
	t.Parallel()
	initErr := errors.New("client init failure")
	bidiStream := &BidiStreamForClientSimple[pingv1.CumSumRequest, pingv1.CumSumResponse]{
		stream: &BidiStreamForClient[pingv1.CumSumRequest, pingv1.CumSumResponse]{err: initErr},
	}
	res, err := bidiStream.Receive()
	assert.Nil(t, res)
	assert.ErrorIs(t, err, initErr)
	verifyHeaders(t, bidiStream.ResponseHeader())
	verifyHeaders(t, bidiStream.ResponseTrailer())
	assert.ErrorIs(t, bidiStream.Send(&pingv1.CumSumRequest{}), initErr)
	assert.ErrorIs(t, bidiStream.CloseRequest(), initErr)
	assert.ErrorIs(t, bidiStream.CloseResponse(), initErr)
	assert.NotNil(t, err)
}

func TestBidiStreamForClientSimple_NilStreamNoPanics(t *testing.T) {
	t.Parallel()
	bidiStream := &BidiStreamForClientSimple[pingv1.PingRequest, pingv1.PingResponse]{}
	// Should not panic
	bidiStream.Peer()
	bidiStream.Spec()
	bidiStream.ResponseHeader()
	bidiStream.ResponseTrailer()
	err := bidiStream.Send(&pingv1.PingRequest{})
	assert.NotNil(t, err)
	assert.Equal(t, err.Error(), expectedStreamErrorMessage)
	res, err := bidiStream.Receive()
	assert.Nil(t, res)
	assert.NotNil(t, err)
	assert.Equal(t, err.Error(), expectedStreamErrorMessage)
	err = bidiStream.CloseRequest()
	assert.Nil(t, res)
	assert.NotNil(t, err)
	assert.Equal(t, err.Error(), expectedStreamErrorMessage)
	err = bidiStream.CloseResponse()
	assert.Nil(t, res)
	assert.NotNil(t, err)
	assert.Equal(t, err.Error(), expectedStreamErrorMessage)
}

func verifyHeaders(t *testing.T, headers http.Header) {
	t.Helper()
	assert.Equal(t, headers, http.Header{})

	// Verify set/del don't panic
	headers.Set("A", "b")
	headers.Del("A")
}

type nopStreamingClientConn struct {
	StreamingClientConn
}

func (c *nopStreamingClientConn) Receive(msg any) error {
	return nil
}

func (c *nopStreamingClientConn) Spec() Spec {
	return Spec{}
}