more test

open-gpdb · Mar 23, 2024 · 18917cc · 18917cc
1 parent c996dcd
commit 18917cc
Showing 1 changed file with 94 additions and 0 deletions.
diff --git a/pkg/proc/yrreader_test.go b/pkg/proc/yrreader_test.go
@@ -1,6 +1,7 @@
 package proc_test
 
 import (
+	"fmt"
 	"io"
 	"testing"
 
@@ -30,3 +31,96 @@ func TestYproxyRetryReaderEmpty(t *testing.T) {
 
 	assert.Nil(t, yr.Close())
 }
+
+func TestYproxyRetryReaderSimpleRead(t *testing.T) {
+
+	ctrl := gomock.NewController(t)
+
+	rr := mock.NewMockRestartReader(ctrl)
+
+	yr := proc.NewYRetryReader(rr)
+
+	buf := []byte{0, 0, 0}
+
+	rr.EXPECT().Restart(int64(0)).Return(nil)
+	rr.EXPECT().Read(gomock.Any()).Do(
+		func(rbuf []byte) {
+			rbuf[0] = 1
+			rbuf[1] = 27
+			rbuf[2] = 33
+		},
+	).Return(3, nil)
+	rr.EXPECT().Read(gomock.Any()).Return(0, io.EOF)
+	rr.EXPECT().Close().Times(1)
+
+	n, err := yr.Read(buf)
+	assert.NoError(t, err)
+	assert.Equal(t, 3, n)
+
+	assert.Equal(t, buf[0], uint8(1))
+	assert.Equal(t, buf[1], uint8(27))
+	assert.Equal(t, buf[2], uint8(33))
+
+	_, err = yr.Read(buf)
+
+	assert.Equal(t, io.EOF, err)
+
+	assert.Nil(t, yr.Close())
+}
+
+func TestYproxyRetryReaderSimpleReadRetry(t *testing.T) {
+
+	ctrl := gomock.NewController(t)
+
+	rr := mock.NewMockRestartReader(ctrl)
+
+	yr := proc.NewYRetryReader(rr)
+
+	buf := []byte{0, 0, 0}
+
+	rr.EXPECT().Restart(int64(0)).Return(nil).Times(1)
+	rr.EXPECT().Restart(int64(3)).Return(nil).Times(1)
+	rr.EXPECT().Read(buf).Do(
+		func(rbuf []byte) {
+			rbuf[0] = 1
+			rbuf[1] = 27
+			rbuf[2] = 33
+		},
+	).Return(3, nil)
+	rr.EXPECT().Read(gomock.Any()).Return(0, fmt.Errorf("no"))
+	rr.EXPECT().Read(gomock.Any()).Do(
+		func(rbuf []byte) {
+			rbuf[0] = 1
+			rbuf[1] = 27
+			rbuf[2] = 33
+		},
+	).Return(3, nil)
+	rr.EXPECT().Read(gomock.Any()).Return(0, io.EOF)
+
+	rr.EXPECT().Close().Times(2)
+
+	n, err := yr.Read(buf)
+	assert.NoError(t, err)
+	assert.Equal(t, 3, n)
+
+	assert.Equal(t, buf[0], uint8(1))
+	assert.Equal(t, buf[1], uint8(27))
+	assert.Equal(t, buf[2], uint8(33))
+
+	// yr got error, but retries
+	buf2 := []byte{0, 0, 0}
+
+	n, err = yr.Read(buf2)
+	assert.NoError(t, err)
+	assert.Equal(t, 3, n)
+
+	assert.Equal(t, buf2[0], uint8(1))
+	assert.Equal(t, buf2[1], uint8(27))
+	assert.Equal(t, buf2[2], uint8(33))
+
+	_, err = yr.Read([]byte{0, 0, 0})
+
+	assert.Equal(t, io.EOF, err)
+
+	assert.Nil(t, yr.Close())
+}