github.com/ydb-platform/ydb-go-sdk/v3@v3.57.0/internal/grpcwrapper/rawtopic/rawtopicwriter/streamwriter_test.go (about)

     1  package rawtopicwriter
     2  
     3  import (
     4  	"testing"
     5  
     6  	"github.com/stretchr/testify/require"
     7  	"github.com/ydb-platform/ydb-go-genproto/protos/Ydb_Topic"
     8  	"google.golang.org/grpc/codes"
     9  	grpcStatus "google.golang.org/grpc/status"
    10  )
    11  
    12  func TestSendWriteRequest(t *testing.T) {
    13  	t.Run("OK", func(t *testing.T) {
    14  		expected := &Ydb_Topic.StreamWriteMessage_FromClient_WriteRequest{
    15  			WriteRequest: &Ydb_Topic.StreamWriteMessage_WriteRequest{
    16  				Codec: 123,
    17  				Messages: []*Ydb_Topic.StreamWriteMessage_WriteRequest_MessageData{
    18  					{
    19  						SeqNo: 1,
    20  					},
    21  				},
    22  			},
    23  		}
    24  
    25  		sendCounter := 0
    26  		var send sendFunc = func(req *Ydb_Topic.StreamWriteMessage_FromClient) error {
    27  			sendCounter++
    28  			require.Equal(t, expected, req.GetClientMessage())
    29  
    30  			return nil
    31  		}
    32  		err := sendWriteRequest(send, expected)
    33  		require.NoError(t, err)
    34  		require.Equal(t, 1, sendCounter)
    35  	})
    36  
    37  	t.Run("Split", func(t *testing.T) {
    38  		originalMessage := &Ydb_Topic.StreamWriteMessage_WriteRequest{
    39  			Codec: 123,
    40  			Messages: []*Ydb_Topic.StreamWriteMessage_WriteRequest_MessageData{
    41  				{
    42  					SeqNo: 1,
    43  				},
    44  				{
    45  					SeqNo: 2,
    46  				},
    47  				{
    48  					SeqNo: 3,
    49  				},
    50  			},
    51  		}
    52  
    53  		split1 := &Ydb_Topic.StreamWriteMessage_WriteRequest{
    54  			Codec: 123,
    55  			Messages: []*Ydb_Topic.StreamWriteMessage_WriteRequest_MessageData{
    56  				{
    57  					SeqNo: 1,
    58  				},
    59  			},
    60  		}
    61  
    62  		split2 := &Ydb_Topic.StreamWriteMessage_WriteRequest{
    63  			Codec: 123,
    64  			Messages: []*Ydb_Topic.StreamWriteMessage_WriteRequest_MessageData{
    65  				{
    66  					SeqNo: 2,
    67  				},
    68  				{
    69  					SeqNo: 3,
    70  				},
    71  			},
    72  		}
    73  
    74  		getWriteRequest := func(req *Ydb_Topic.StreamWriteMessage_FromClient) *Ydb_Topic.StreamWriteMessage_WriteRequest {
    75  			return req.GetClientMessage().(*Ydb_Topic.StreamWriteMessage_FromClient_WriteRequest).WriteRequest
    76  		}
    77  
    78  		sendCounter := 0
    79  		var send sendFunc = func(commonReq *Ydb_Topic.StreamWriteMessage_FromClient) error {
    80  			sendCounter++
    81  			req := getWriteRequest(commonReq)
    82  			switch sendCounter {
    83  			case 1:
    84  				require.Equal(t, originalMessage, req)
    85  
    86  				return grpcStatus.Error(codes.ResourceExhausted, "test resource exhausted")
    87  			case 2:
    88  				require.Equal(t, split1, req)
    89  			case 3:
    90  				require.Equal(t, split2, req)
    91  			default:
    92  				t.Fatal()
    93  			}
    94  
    95  			return nil
    96  		}
    97  
    98  		err := sendWriteRequest(send, &Ydb_Topic.StreamWriteMessage_FromClient_WriteRequest{WriteRequest: originalMessage})
    99  		require.NoError(t, err)
   100  		require.Equal(t, 3, sendCounter)
   101  	})
   102  }