From 3296b5e6032b1c54f496ffdda8fb13f9be46b8b5 Mon Sep 17 00:00:00 2001
From: lance6716 <lance6716@gmail.com>
Date: Tue, 8 Aug 2023 15:15:24 +0800
Subject: [PATCH] lightning/external: add UT and fix a bug (#45874)

ref pingcap/tidb#45719
---
 br/pkg/lightning/backend/external/BUILD.bazel |  3 +-
 .../backend/external/byte_reader_test.go      | 43 +++++++++++++++++
 br/pkg/lightning/backend/external/file.go     |  4 +-
 .../lightning/backend/external/file_test.go   | 48 ++++++++++++++++++-
 4 files changed, 94 insertions(+), 4 deletions(-)

diff --git a/br/pkg/lightning/backend/external/BUILD.bazel b/br/pkg/lightning/backend/external/BUILD.bazel
index d0304bba1a8fe..3bb3e1a6a985c 100644
--- a/br/pkg/lightning/backend/external/BUILD.bazel
+++ b/br/pkg/lightning/backend/external/BUILD.bazel
@@ -31,10 +31,11 @@ go_test(
     ],
     embed = [":external"],
     flaky = True,
-    shard_count = 5,
+    shard_count = 7,
     deps = [
         "//br/pkg/storage",
         "@com_github_pingcap_errors//:errors",
         "@com_github_stretchr_testify//require",
+        "@org_golang_x_exp//rand",
     ],
 )
diff --git a/br/pkg/lightning/backend/external/byte_reader_test.go b/br/pkg/lightning/backend/external/byte_reader_test.go
index 35b34d2a1a894..6a6fd8fb0d6cf 100644
--- a/br/pkg/lightning/backend/external/byte_reader_test.go
+++ b/br/pkg/lightning/backend/external/byte_reader_test.go
@@ -18,9 +18,11 @@ import (
 	"context"
 	"io"
 	"testing"
+	"time"
 
 	"github.com/pingcap/errors"
 	"github.com/stretchr/testify/require"
+	"golang.org/x/exp/rand"
 )
 
 // mockExtStore is only used for test.
@@ -178,3 +180,44 @@ func TestByteReaderAuxBuf(t *testing.T) {
 	require.Equal(t, []byte("0"), *y1)
 	require.Equal(t, []byte("12"), *y2)
 }
+
+func TestReset(t *testing.T) {
+	seed := time.Now().Unix()
+	rand.Seed(uint64(seed))
+	t.Logf("seed: %d", seed)
+	src := make([]byte, 256)
+	for i := range src {
+		src[i] = byte(i)
+	}
+	ms := &mockExtStore{src: src}
+	bufSize := rand.Intn(256)
+	br, err := newByteReader(context.Background(), ms, bufSize)
+	require.NoError(t, err)
+	end := 0
+	toCheck := make([]*[]byte, 0, 10)
+	for end < len(src) {
+		n := rand.Intn(len(src) - end)
+		if n == 0 {
+			n = 1
+		}
+		y, err := br.readNBytes(n)
+		require.NoError(t, err)
+		toCheck = append(toCheck, y)
+		end += n
+
+		l := end
+		r := end
+		for i := len(toCheck) - 1; i >= 0; i-- {
+			l -= len(*toCheck[i])
+			require.Equal(t, src[l:r], *toCheck[i])
+			r = l
+		}
+
+		if rand.Intn(2) == 0 {
+			br.reset()
+			toCheck = toCheck[:0]
+		}
+	}
+	_, err = br.readNBytes(1)
+	require.Equal(t, io.EOF, err)
+}
diff --git a/br/pkg/lightning/backend/external/file.go b/br/pkg/lightning/backend/external/file.go
index 91ffdf47cb499..a87f0a255aa5f 100644
--- a/br/pkg/lightning/backend/external/file.go
+++ b/br/pkg/lightning/backend/external/file.go
@@ -65,7 +65,7 @@ func (s *KeyValueStore) AddKeyValue(key, value []byte) error {
 	// data layout: keyLen + key + valueLen + value
 	_, err := s.dataWriter.Write(
 		s.ctx,
-		binary.BigEndian.AppendUint64(b[:], uint64(len(key))),
+		binary.BigEndian.AppendUint64(b[:0], uint64(len(key))),
 	)
 	if err != nil {
 		return err
@@ -76,7 +76,7 @@ func (s *KeyValueStore) AddKeyValue(key, value []byte) error {
 	}
 	_, err = s.dataWriter.Write(
 		s.ctx,
-		binary.BigEndian.AppendUint64(b[:], uint64(len(value))),
+		binary.BigEndian.AppendUint64(b[:0], uint64(len(value))),
 	)
 	if err != nil {
 		return err
diff --git a/br/pkg/lightning/backend/external/file_test.go b/br/pkg/lightning/backend/external/file_test.go
index d6ed0cb2524d3..fa2b5ccbd79c6 100644
--- a/br/pkg/lightning/backend/external/file_test.go
+++ b/br/pkg/lightning/backend/external/file_test.go
@@ -16,10 +16,13 @@ package external
 
 import (
 	"context"
+	"io"
 	"testing"
+	"time"
 
 	"github.com/pingcap/tidb/br/pkg/storage"
 	"github.com/stretchr/testify/require"
+	"golang.org/x/exp/rand"
 )
 
 func TestAddKeyValueMaintainRangeProperty(t *testing.T) {
@@ -104,4 +107,47 @@ func TestAddKeyValueMaintainRangeProperty(t *testing.T) {
 	require.NoError(t, err)
 }
 
-// TODO(lance6716): add more tests when the usage of other functions are merged into master.
+func TestKVReadWrite(t *testing.T) {
+	seed := time.Now().Unix()
+	rand.Seed(uint64(seed))
+	t.Logf("seed: %d", seed)
+	ctx := context.Background()
+	memStore := storage.NewMemStorage()
+	writer, err := memStore.Create(ctx, "/test", nil)
+	require.NoError(t, err)
+	rc := &rangePropertiesCollector{
+		propSizeIdxDistance: 100,
+		propKeysIdxDistance: 2,
+	}
+	rc.reset()
+	kvStore, err := NewKeyValueStore(ctx, writer, rc, 1, 1)
+	require.NoError(t, err)
+
+	kvCnt := rand.Intn(10) + 10
+	keys := make([][]byte, kvCnt)
+	values := make([][]byte, kvCnt)
+	for i := 0; i < kvCnt; i++ {
+		randLen := rand.Intn(10) + 1
+		keys[i] = make([]byte, randLen)
+		rand.Read(keys[i])
+		randLen = rand.Intn(10) + 1
+		values[i] = make([]byte, randLen)
+		rand.Read(values[i])
+		err = kvStore.AddKeyValue(keys[i], values[i])
+		require.NoError(t, err)
+	}
+	err = writer.Close(ctx)
+	require.NoError(t, err)
+
+	bufSize := rand.Intn(100) + 1
+	kvReader, err := newKVReader(ctx, "/test", memStore, 0, bufSize)
+	require.NoError(t, err)
+	for i := 0; i < kvCnt; i++ {
+		key, value, err := kvReader.nextKV()
+		require.NoError(t, err)
+		require.Equal(t, keys[i], key)
+		require.Equal(t, values[i], value)
+	}
+	_, _, err = kvReader.nextKV()
+	require.Equal(t, io.EOF, err)
+}