github.com/ydb-platform/ydb-go-sdk/v3@v3.57.0/internal/topic/topicreaderinternal/reader_test.go (about)

     1  package topicreaderinternal
     2  
     3  import (
     4  	"context"
     5  	"errors"
     6  	"runtime"
     7  	"testing"
     8  
     9  	"github.com/stretchr/testify/require"
    10  	"go.uber.org/mock/gomock"
    11  
    12  	"github.com/ydb-platform/ydb-go-sdk/v3/internal/empty"
    13  	"github.com/ydb-platform/ydb-go-sdk/v3/internal/xcontext"
    14  	"github.com/ydb-platform/ydb-go-sdk/v3/internal/xtest"
    15  )
    16  
    17  func TestReader_Close(t *testing.T) {
    18  	xtest.TestManyTimes(t, func(t testing.TB) {
    19  		mc := gomock.NewController(t)
    20  		defer mc.Finish()
    21  
    22  		testErr := errors.New("test error")
    23  		readerContext, readerCancel := xcontext.WithCancel(context.Background())
    24  		baseReader := NewMockbatchedStreamReader(mc)
    25  		baseReader.EXPECT().ReadMessageBatch(gomock.Any(), ReadMessageBatchOptions{}).Do(func(_, _ interface{}) {
    26  			<-readerContext.Done()
    27  		}).Return(nil, testErr)
    28  		baseReader.EXPECT().ReadMessageBatch(
    29  			gomock.Any(),
    30  			ReadMessageBatchOptions{batcherGetOptions: batcherGetOptions{MaxCount: 1, MinCount: 1}},
    31  		).Do(func(_, _ interface{}) {
    32  			<-readerContext.Done()
    33  		}).Return(nil, testErr)
    34  		baseReader.EXPECT().Commit(gomock.Any(), gomock.Any()).Do(func(_, _ interface{}) {
    35  			<-readerContext.Done()
    36  		}).Return(testErr)
    37  		baseReader.EXPECT().CloseWithError(gomock.Any(), gomock.Any()).Do(func(_, _ interface{}) {
    38  			readerCancel()
    39  		})
    40  
    41  		reader := &Reader{
    42  			reader: baseReader,
    43  		}
    44  
    45  		type callState struct {
    46  			callCompleted empty.Chan
    47  			err           error
    48  		}
    49  
    50  		isCallCompleted := func(state *callState) bool {
    51  			select {
    52  			case <-state.callCompleted:
    53  				return true
    54  			default:
    55  				return false
    56  			}
    57  		}
    58  
    59  		var allStates []*callState
    60  		newCallState := func() *callState {
    61  			state := &callState{
    62  				callCompleted: make(empty.Chan),
    63  			}
    64  			allStates = append(allStates, state)
    65  
    66  			return state
    67  		}
    68  
    69  		readerCommitState := newCallState()
    70  		readerReadMessageState := newCallState()
    71  		readerReadMessageBatchState := newCallState()
    72  
    73  		go func() {
    74  			readerCommitState.err = reader.Commit(context.Background(), &PublicMessage{
    75  				commitRange: commitRange{
    76  					partitionSession: &partitionSession{},
    77  				},
    78  			})
    79  			close(readerCommitState.callCompleted)
    80  		}()
    81  
    82  		go func() {
    83  			_, readerReadMessageState.err = reader.ReadMessage(context.Background())
    84  			close(readerReadMessageState.callCompleted)
    85  		}()
    86  
    87  		go func() {
    88  			_, readerReadMessageBatchState.err = reader.ReadMessageBatch(context.Background())
    89  			close(readerReadMessageBatchState.callCompleted)
    90  		}()
    91  
    92  		runtime.Gosched()
    93  
    94  		// check about no methods finished before close
    95  		for i := range allStates {
    96  			require.False(t, isCallCompleted(allStates[i]))
    97  		}
    98  		require.NoError(t, reader.Close(context.Background()))
    99  
   100  		// check about all methods stop work after close
   101  		for i := range allStates {
   102  			<-allStates[i].callCompleted
   103  			require.Error(t, allStates[i].err, i)
   104  		}
   105  	})
   106  }
   107  
   108  func TestReader_Commit(t *testing.T) {
   109  	t.Run("OK", func(t *testing.T) {
   110  		mc := gomock.NewController(t)
   111  		defer mc.Finish()
   112  
   113  		readerID := nextReaderID()
   114  		baseReader := NewMockbatchedStreamReader(mc)
   115  		reader := &Reader{
   116  			reader:   baseReader,
   117  			readerID: readerID,
   118  		}
   119  
   120  		expectedRangeOk := commitRange{
   121  			commitOffsetStart: 1,
   122  			commitOffsetEnd:   10,
   123  			partitionSession: &partitionSession{
   124  				readerID:           readerID,
   125  				partitionSessionID: 10,
   126  			},
   127  		}
   128  		baseReader.EXPECT().Commit(gomock.Any(), expectedRangeOk).Return(nil)
   129  		require.NoError(t, reader.Commit(context.Background(), &PublicMessage{commitRange: expectedRangeOk}))
   130  
   131  		expectedRangeErr := commitRange{
   132  			commitOffsetStart: 15,
   133  			commitOffsetEnd:   20,
   134  			partitionSession: &partitionSession{
   135  				readerID:           readerID,
   136  				partitionSessionID: 30,
   137  			},
   138  		}
   139  
   140  		testErr := errors.New("test err")
   141  		baseReader.EXPECT().Commit(gomock.Any(), expectedRangeErr).Return(testErr)
   142  		require.ErrorIs(t, reader.Commit(context.Background(), &PublicMessage{commitRange: expectedRangeErr}), testErr)
   143  	})
   144  
   145  	t.Run("CommitFromOtherReader", func(t *testing.T) {
   146  		ctx := xtest.Context(t)
   147  		reader := &Reader{readerID: 1}
   148  		forCommit := commitRange{
   149  			commitOffsetStart: 1,
   150  			commitOffsetEnd:   2,
   151  			partitionSession:  &partitionSession{readerID: 2},
   152  		}
   153  		err := reader.Commit(ctx, forCommit)
   154  		require.ErrorIs(t, err, errCommitSessionFromOtherReader)
   155  	})
   156  }
   157  
   158  func TestReader_WaitInit(t *testing.T) {
   159  	mc := gomock.NewController(t)
   160  	defer mc.Finish()
   161  
   162  	readerID := nextReaderID()
   163  	baseReader := NewMockbatchedStreamReader(mc)
   164  	reader := &Reader{
   165  		reader:   baseReader,
   166  		readerID: readerID,
   167  	}
   168  
   169  	baseReader.EXPECT().WaitInit(gomock.Any())
   170  	err := reader.WaitInit(context.Background())
   171  	require.NoError(t, err)
   172  }