github.com/vescale/zgraph@v0.0.0-20230410094002-959c02d50f95/storage/batch_getter_test.go (about)

     1  // Copyright 2022 zGraph Authors. All rights reserved.
     2  //
     3  // Copyright 2016 PingCAP, Inc.
     4  //
     5  // Licensed under the Apache License, Version 2.0 (the "License");
     6  // you may not use this file except in compliance with the License.
     7  // You may obtain a copy of the License at
     8  //
     9  //     http://www.apache.org/licenses/LICENSE-2.0
    10  //
    11  // Unless required by applicable law or agreed to in writing, software
    12  // distributed under the License is distributed on an "AS IS" BASIS,
    13  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    14  // See the License for the specific language governing permissions and
    15  // limitations under the License.
    16  
    17  package storage
    18  
    19  import (
    20  	"context"
    21  	"testing"
    22  
    23  	"github.com/stretchr/testify/assert"
    24  	"github.com/vescale/zgraph/storage/kv"
    25  )
    26  
    27  func TestBufferBatchGetter(t *testing.T) {
    28  	snap := newMockStore()
    29  	ka := []byte("a")
    30  	kb := []byte("b")
    31  	kc := []byte("c")
    32  	kd := []byte("d")
    33  	snap.Set(ka, ka)
    34  	snap.Set(kb, kb)
    35  	snap.Set(kc, kc)
    36  	snap.Set(kd, kd)
    37  
    38  	buffer := newMockStore()
    39  	buffer.Set(ka, []byte("a1"))
    40  	buffer.Delete(kb)
    41  
    42  	batchGetter := NewBufferBatchGetter(buffer, snap)
    43  	result, err := batchGetter.BatchGet(context.Background(), []kv.Key{ka, kb, kc, kd})
    44  	assert.NoError(t, err)
    45  	assert.Equal(t, 3, len(result))
    46  	assert.Equal(t, "a1", string(result[string(ka)]))
    47  	assert.Equal(t, "c", string(result[string(kc)]))
    48  	assert.Equal(t, "d", string(result[string(kd)]))
    49  }
    50  
    51  type mockBatchGetterStore struct {
    52  	index [][]byte
    53  	value [][]byte
    54  }
    55  
    56  func newMockStore() *mockBatchGetterStore {
    57  	return &mockBatchGetterStore{
    58  		index: make([][]byte, 0),
    59  		value: make([][]byte, 0),
    60  	}
    61  }
    62  
    63  func (s *mockBatchGetterStore) Len() int {
    64  	return len(s.index)
    65  }
    66  
    67  func (s *mockBatchGetterStore) Get(_ context.Context, k kv.Key) ([]byte, error) {
    68  	for i, key := range s.index {
    69  		if kv.CmpKey(key, k) == 0 {
    70  			return s.value[i], nil
    71  		}
    72  	}
    73  	return nil, kv.ErrNotExist
    74  }
    75  
    76  func (s *mockBatchGetterStore) BatchGet(ctx context.Context, keys []kv.Key) (map[string][]byte, error) {
    77  	m := make(map[string][]byte)
    78  	for _, k := range keys {
    79  		v, err := s.Get(ctx, k)
    80  		if err == nil {
    81  			m[string(k)] = v
    82  			continue
    83  		}
    84  		if kv.IsErrNotFound(err) {
    85  			continue
    86  		}
    87  		return m, err
    88  	}
    89  	return m, nil
    90  }
    91  
    92  func (s *mockBatchGetterStore) Set(k []byte, v []byte) error {
    93  	for i, key := range s.index {
    94  		if kv.CmpKey(key, k) == 0 {
    95  			s.value[i] = v
    96  			return nil
    97  		}
    98  	}
    99  	s.index = append(s.index, k)
   100  	s.value = append(s.value, v)
   101  	return nil
   102  }
   103  
   104  func (s *mockBatchGetterStore) Delete(k []byte) error {
   105  	s.Set(k, []byte{})
   106  	return nil
   107  }