github.com/pingcap/tiflow@v0.0.0-20240520035814-5bf52d54e205/pkg/p2p/client_batch_sender_test.go (about)

     1  // Copyright 2021 PingCAP, Inc.
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  //     http://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // See the License for the specific language governing permissions and
    12  // limitations under the License.
    13  
    14  package p2p
    15  
    16  import (
    17  	"context"
    18  	"fmt"
    19  	"math"
    20  	"testing"
    21  
    22  	proto "github.com/pingcap/tiflow/proto/p2p"
    23  	"github.com/stretchr/testify/mock"
    24  	"github.com/stretchr/testify/require"
    25  )
    26  
    27  func TestClientBatchSenderMaxCount(t *testing.T) {
    28  	ctx, cancel := context.WithCancel(context.Background())
    29  	defer cancel()
    30  
    31  	grpcStream := newMockSendMessageClient(ctx)
    32  	sender := newClientBatchSender(grpcStream, 100, math.MaxInt64)
    33  
    34  	grpcStream.AssertNotCalled(t, "Send", mock.Anything)
    35  	// 99 messages
    36  	for i := 1; i < 100; i++ {
    37  		err := sender.Append(&proto.MessageEntry{
    38  			Topic:    "test-topic",
    39  			Content:  []byte(fmt.Sprintf("test-%d", i)),
    40  			Sequence: int64(i),
    41  		})
    42  		require.NoError(t, err)
    43  	}
    44  
    45  	grpcStream.On("Send", mock.Anything).Return(nil).Run(func(args mock.Arguments) {
    46  		msg := args.Get(0).(*proto.MessagePacket)
    47  		require.Len(t, msg.Entries, 100)
    48  		for i, entry := range msg.Entries {
    49  			require.Equal(t, "test-topic", entry.Topic)
    50  			require.Equal(t, int64(i+1), entry.Sequence)
    51  			require.Equal(t, fmt.Sprintf("test-%d", i+1), string(entry.Content))
    52  		}
    53  	})
    54  	// one more message
    55  	err := sender.Append(&proto.MessageEntry{
    56  		Topic:    "test-topic",
    57  		Content:  []byte("test-100"),
    58  		Sequence: 100,
    59  	})
    60  	require.NoError(t, err)
    61  }
    62  
    63  func TestClientBatchSenderMaxSize(t *testing.T) {
    64  	ctx, cancel := context.WithCancel(context.Background())
    65  	defer cancel()
    66  
    67  	grpcStream := newMockSendMessageClient(ctx)
    68  	sender := newClientBatchSender(grpcStream, math.MaxInt64, 1000)
    69  
    70  	grpcStream.AssertNotCalled(t, "Send", mock.Anything)
    71  	i := 1
    72  	for size := 0; size < 220; {
    73  		msg := &proto.MessageEntry{
    74  			Topic:    "test-topic",
    75  			Content:  []byte(fmt.Sprintf("test-%d", i)),
    76  			Sequence: int64(i),
    77  		}
    78  		i++
    79  		size += msg.Size()
    80  		err := sender.Append(msg)
    81  		require.NoError(t, err)
    82  	}
    83  
    84  	grpcStream.On("Send", mock.Anything).Return(nil).Run(func(args mock.Arguments) {
    85  		msg := args.Get(0).(*proto.MessagePacket)
    86  		require.Len(t, msg.Entries, i)
    87  		for i, entry := range msg.Entries {
    88  			require.Equal(t, "test-topic", entry.Topic)
    89  			require.Equal(t, int64(i+1), entry.Sequence)
    90  			require.Equal(t, "11byteshere", string(entry.Content))
    91  		}
    92  	})
    93  	// one more message
    94  	err := sender.Append(&proto.MessageEntry{
    95  		Topic:    "test-topic",
    96  		Content:  []byte(fmt.Sprintf("test-%d", i)),
    97  		Sequence: int64(i),
    98  	})
    99  	require.NoError(t, err)
   100  }
   101  
   102  func TestClientBatchSenderFlush(t *testing.T) {
   103  	ctx, cancel := context.WithCancel(context.Background())
   104  	defer cancel()
   105  
   106  	grpcStream := newMockSendMessageClient(ctx)
   107  	sender := newClientBatchSender(grpcStream, math.MaxInt64, 1000)
   108  
   109  	grpcStream.AssertNotCalled(t, "Send", mock.Anything)
   110  	err := sender.Flush()
   111  	require.NoError(t, err)
   112  }