github.com/muhammadn/cortex@v1.9.1-0.20220510110439-46bb7000d03d/pkg/ingester/client/cortex_util_test.go (about)

     1  package client
     2  
     3  import (
     4  	"context"
     5  	"net"
     6  	"sync"
     7  	"testing"
     8  	"time"
     9  
    10  	"github.com/grafana/dskit/grpcutil"
    11  	"github.com/stretchr/testify/assert"
    12  	"github.com/stretchr/testify/mock"
    13  	"github.com/stretchr/testify/require"
    14  	"google.golang.org/grpc"
    15  	"google.golang.org/grpc/test/bufconn"
    16  
    17  	"github.com/cortexproject/cortex/pkg/util/test"
    18  )
    19  
    20  func TestSendQueryStream(t *testing.T) {
    21  	// Create a new gRPC server with in-memory communication.
    22  	listen := bufconn.Listen(1024 * 1024)
    23  	server := grpc.NewServer()
    24  	bufDialer := func(context.Context, string) (net.Conn, error) {
    25  		return listen.Dial()
    26  	}
    27  
    28  	conn, err := grpc.DialContext(context.Background(), "bufnet", grpc.WithContextDialer(bufDialer), grpc.WithInsecure())
    29  	require.NoError(t, err)
    30  	defer conn.Close()
    31  
    32  	// Create a cancellable context for the client.
    33  	clientCtx, clientCancel := context.WithCancel(context.Background())
    34  
    35  	// Create a WaitGroup used to wait until the mocked server assertions
    36  	// complete before returning.
    37  	wg := sync.WaitGroup{}
    38  	wg.Add(1)
    39  
    40  	serverMock := &IngesterServerMock{}
    41  	serverMock.On("QueryStream", mock.Anything, mock.Anything).Return(nil).Run(func(args mock.Arguments) {
    42  		defer wg.Done()
    43  
    44  		stream := args.Get(1).(grpc.ServerStream)
    45  
    46  		// Cancel the client request.
    47  		clientCancel()
    48  
    49  		// Wait until the cancelling has been propagated to the server.
    50  		test.Poll(t, time.Second, context.Canceled, func() interface{} {
    51  			return stream.Context().Err()
    52  		})
    53  
    54  		// Try to send the response and assert the error we get is the context.Canceled
    55  		// and not transport.ErrIllegalHeaderWrite. This is the assertion we care about
    56  		// in this test.
    57  		err := SendQueryStream(stream.(Ingester_QueryStreamServer), &QueryStreamResponse{})
    58  		assert.Equal(t, context.Canceled, err)
    59  	})
    60  
    61  	RegisterIngesterServer(server, serverMock)
    62  
    63  	go func() {
    64  		require.NoError(t, server.Serve(listen))
    65  	}()
    66  
    67  	client := NewIngesterClient(conn)
    68  	stream, err := client.QueryStream(clientCtx, &QueryRequest{})
    69  	require.NoError(t, err)
    70  
    71  	// Try to receive the response and assert the error we get is the context.Canceled
    72  	// wrapped within a gRPC error.
    73  	_, err = stream.Recv()
    74  	assert.Equal(t, true, grpcutil.IsGRPCContextCanceled(err))
    75  
    76  	// Wait until the assertions in the server mock have completed.
    77  	wg.Wait()
    78  }