cody auto generated test cases

Signed-off-by: Derek Wang <[email protected]>

pkg/sources/udsource/user_defined_source_test.go

@@ -0,0 +1,173 @@
+/*
+Copyright 2022 The Numaproj 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 udsource
+
+import (
+	"context"
+	"errors"
+	"io"
+	"testing"
+	"time"
+
+	"github.com/golang/mock/gomock"
+	sourcepb "github.com/numaproj/numaflow-go/pkg/apis/proto/source/v1"
+	"github.com/numaproj/numaflow-go/pkg/apis/proto/source/v1/sourcemock"
+	"github.com/stretchr/testify/assert"
+	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/status"
+	"google.golang.org/protobuf/types/known/emptypb"
+	"google.golang.org/protobuf/types/known/timestamppb"
+
+	"github.com/numaproj/numaflow/pkg/isb"
+)
+
+func TestUserDefinedSource_Read(t *testing.T) {
+	ctrl := gomock.NewController(t)
+	defer ctrl.Finish()
+
+	mockClient := sourcemock.NewMockSourceClient(ctrl)
+	mockReadClient := sourcemock.NewMockSource_ReadFnClient(ctrl)
+
+	req := &sourcepb.ReadRequest{
+		Request: &sourcepb.ReadRequest_Request{
+			NumRecords:  1,
+			TimeoutInMs: 1000,
+		},
+	}
+
+	offset := &sourcepb.Offset{Offset: []byte(`test_offset`), PartitionId: 0}
+	eventTime := time.Unix(1661169600, 0).UTC()
+	expectedResponse := &sourcepb.ReadResponse{
+		Result: &sourcepb.ReadResponse_Result{
+			Payload:   []byte(`test_payload`),
+			Offset:    offset,
+			EventTime: timestamppb.New(eventTime),
+			Keys:      []string{"test_key"},
+		},
+	}
+
+	mockReadClient.EXPECT().Recv().Return(expectedResponse, nil).Times(1)
+	mockReadClient.EXPECT().Recv().Return(nil, io.EOF).Times(1)
+	mockClient.EXPECT().ReadFn(gomock.Any(), &rpcMsg{msg: req}).Return(mockReadClient, nil)
+
+	ctx, cancel := context.WithTimeout(context.Background(), time.Second)
+	defer cancel()
+
+	u := NewMockUDSgRPCBasedUDSource(mockClient)
+	readMessages, err := u.ApplyReadFn(ctx, 1, time.Millisecond*1000)
+	assert.NoError(t, err)
+	assert.Equal(t, 1, len(readMessages))
+	assert.Equal(t, []byte(`test_payload`), readMessages[0].Body.Payload)
+	assert.Equal(t, []string{"test_key"}, readMessages[0].Keys)
+	assert.Equal(t, NewUserDefinedSourceOffset(offset), readMessages[0].ReadOffset)
+	assert.Equal(t, eventTime, readMessages[0].EventTime)
+}
+
+func TestUserDefinedSource_Read_Error(t *testing.T) {
+	ctrl := gomock.NewController(t)
+	defer ctrl.Finish()
+
+	mockClient := sourcemock.NewMockSourceClient(ctrl)
+	mockReadClient := sourcemock.NewMockSource_ReadFnClient(ctrl)
+
+	req := &sourcepb.ReadRequest{
+		Request: &sourcepb.ReadRequest_Request{
+			NumRecords:  1,
+			TimeoutInMs: 1000,
+		},
+	}
+
+	eventTime := time.Unix(1661169600, 0).UTC()
+	expectedResponse := &sourcepb.ReadResponse{
+		Result: &sourcepb.ReadResponse_Result{
+			Payload:   []byte(`test_payload`),
+			Offset:    &sourcepb.Offset{Offset: []byte(`test_offset`), PartitionId: 0},
+			EventTime: timestamppb.New(eventTime),
+			Keys:      []string{"test_key"},
+		},
+	}
+
+	mockReadClient.EXPECT().Recv().Return(expectedResponse, errors.New("mock error for read")).AnyTimes()
+	mockClient.EXPECT().ReadFn(gomock.Any(), &rpcMsg{msg: req}).Return(mockReadClient, nil)
+
+	ctx, cancel := context.WithTimeout(context.Background(), time.Second)
+	defer cancel()
+
+	u := NewMockUDSgRPCBasedUDSource(mockClient)
+	readMessages, err := u.ApplyReadFn(ctx, 1, time.Millisecond*1000)
+	assert.Error(t, err)
+	assert.Equal(t, 0, len(readMessages))
+}
+
+func TestUserDefinedSource_Ack(t *testing.T) {
+	ctrl := gomock.NewController(t)
+	defer ctrl.Finish()
+
+	offset1 := &sourcepb.Offset{Offset: []byte("test-offset-1"), PartitionId: 0}
+	offset2 := &sourcepb.Offset{Offset: []byte("test-offset-2"), PartitionId: 0}
+
+	mockClient := sourcemock.NewMockSourceClient(ctrl)
+	req := &sourcepb.AckRequest{
+		Request: &sourcepb.AckRequest_Request{
+			Offsets: []*sourcepb.Offset{
+				offset1,
+				offset2,
+			},
+		},
+	}
+
+	mockClient.EXPECT().AckFn(gomock.Any(), &rpcMsg{msg: req}).Return(&sourcepb.AckResponse{Result: &sourcepb.AckResponse_Result{Success: &emptypb.Empty{}}}, nil).AnyTimes()
+
+	ctx, cancel := context.WithTimeout(context.Background(), time.Second)
+	defer cancel()
+
+	u := NewMockUDSgRPCBasedUDSource(mockClient)
+	err := u.ApplyAckFn(ctx, []isb.Offset{
+		NewUserDefinedSourceOffset(offset1),
+		NewUserDefinedSourceOffset(offset2),
+	})
+	assert.NoError(t, err)
+}
+
+func TestUserDefinedSource_Ack_Error(t *testing.T) {
+	ctrl := gomock.NewController(t)
+	defer ctrl.Finish()
+
+	offset1 := &sourcepb.Offset{Offset: []byte("test-offset-1"), PartitionId: 0}
+	offset2 := &sourcepb.Offset{Offset: []byte("test-offset-2"), PartitionId: 0}
+
+	mockClient := sourcemock.NewMockSourceClient(ctrl)
+	req := &sourcepb.AckRequest{
+		Request: &sourcepb.AckRequest_Request{
+			Offsets: []*sourcepb.Offset{
+				offset1,
+				offset2,
+			},
+		},
+	}
+
+	mockClient.EXPECT().AckFn(gomock.Any(), &rpcMsg{msg: req}).Return(nil, status.New(codes.DeadlineExceeded, "mock test err").Err())
+	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
+	defer cancel()
+
+	u := NewMockUDSgRPCBasedUDSource(mockClient)
+	err := u.ApplyAckFn(ctx, []isb.Offset{
+		NewUserDefinedSourceOffset(offset1),
+		NewUserDefinedSourceOffset(offset2),
+	})
+	assert.ErrorIs(t, err, status.New(codes.DeadlineExceeded, "mock test err").Err())
+}