109 lines
2.9 KiB
Go
109 lines
2.9 KiB
Go
// Copyright 2023 PingCAP, Inc.
|
|
//
|
|
// 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 external
|
|
|
|
import (
|
|
"bytes"
|
|
"context"
|
|
"fmt"
|
|
"slices"
|
|
"testing"
|
|
"time"
|
|
|
|
"github.com/pingcap/tidb/br/pkg/lightning/backend/kv"
|
|
"github.com/pingcap/tidb/br/pkg/lightning/common"
|
|
"github.com/pingcap/tidb/br/pkg/storage"
|
|
"github.com/pingcap/tidb/pkg/util/size"
|
|
"github.com/stretchr/testify/require"
|
|
"golang.org/x/exp/rand"
|
|
)
|
|
|
|
func TestReadAllDataBasic(t *testing.T) {
|
|
seed := time.Now().Unix()
|
|
rand.Seed(uint64(seed))
|
|
t.Logf("seed: %d", seed)
|
|
ctx := context.Background()
|
|
memStore := storage.NewMemStorage()
|
|
memSizeLimit := (rand.Intn(10) + 1) * 400
|
|
|
|
w := NewWriterBuilder().
|
|
SetPropSizeDistance(100).
|
|
SetPropKeysDistance(2).
|
|
SetMemorySizeLimit(uint64(memSizeLimit)).
|
|
SetBlockSize(memSizeLimit).
|
|
Build(memStore, "/test", "0")
|
|
|
|
writer := NewEngineWriter(w)
|
|
kvCnt := rand.Intn(10) + 10000
|
|
kvs := make([]common.KvPair, kvCnt)
|
|
for i := 0; i < kvCnt; i++ {
|
|
kvs[i] = common.KvPair{
|
|
Key: []byte(fmt.Sprintf("key%05d", i)),
|
|
Val: []byte("56789"),
|
|
}
|
|
}
|
|
|
|
require.NoError(t, writer.AppendRows(ctx, nil, kv.MakeRowsFromKvPairs(kvs)))
|
|
_, err := writer.Close(ctx)
|
|
require.NoError(t, err)
|
|
|
|
slices.SortFunc(kvs, func(i, j common.KvPair) int {
|
|
return bytes.Compare(i.Key, j.Key)
|
|
})
|
|
|
|
datas, stats, err := GetAllFileNames(ctx, memStore, "")
|
|
require.NoError(t, err)
|
|
|
|
testReadAndCompare(ctx, t, kvs, memStore, datas, stats, kvs[0].Key, memSizeLimit)
|
|
}
|
|
|
|
func TestReadAllOneFile(t *testing.T) {
|
|
seed := time.Now().Unix()
|
|
rand.Seed(uint64(seed))
|
|
t.Logf("seed: %d", seed)
|
|
ctx := context.Background()
|
|
memStore := storage.NewMemStorage()
|
|
memSizeLimit := (rand.Intn(10) + 1) * 400
|
|
|
|
w := NewWriterBuilder().
|
|
SetPropSizeDistance(100).
|
|
SetPropKeysDistance(2).
|
|
SetMemorySizeLimit(uint64(memSizeLimit)).
|
|
BuildOneFile(memStore, "/test", "0")
|
|
|
|
require.NoError(t, w.Init(ctx, int64(5*size.MB)))
|
|
|
|
kvCnt := rand.Intn(10) + 10000
|
|
kvs := make([]common.KvPair, kvCnt)
|
|
for i := 0; i < kvCnt; i++ {
|
|
kvs[i] = common.KvPair{
|
|
Key: []byte(fmt.Sprintf("key%05d", i)),
|
|
Val: []byte("56789"),
|
|
}
|
|
require.NoError(t, w.WriteRow(ctx, kvs[i].Key, kvs[i].Val))
|
|
}
|
|
|
|
require.NoError(t, w.Close(ctx))
|
|
|
|
slices.SortFunc(kvs, func(i, j common.KvPair) int {
|
|
return bytes.Compare(i.Key, j.Key)
|
|
})
|
|
|
|
datas, stats, err := GetAllFileNames(ctx, memStore, "")
|
|
require.NoError(t, err)
|
|
|
|
testReadAndCompare(ctx, t, kvs, memStore, datas, stats, kvs[0].Key, memSizeLimit)
|
|
}
|