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

     1  // Copyright 2022 zGraph Authors. All rights reserved.
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  //     http://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    12  // See the License for the specific language governing permissions and
    13  // limitations under the License.
    14  
    15  package storage
    16  
    17  import (
    18  	"bytes"
    19  	"context"
    20  	"reflect"
    21  	"testing"
    22  
    23  	"github.com/cockroachdb/pebble"
    24  	"github.com/stretchr/testify/assert"
    25  	"github.com/vescale/zgraph/storage/kv"
    26  	"github.com/vescale/zgraph/storage/mvcc"
    27  )
    28  
    29  func TestSnapshot_Get(t *testing.T) {
    30  	s, err := Open(t.TempDir())
    31  	assert.Nil(t, err)
    32  	assert.NotNil(t, s)
    33  
    34  	db := s.(*mvccStorage).db
    35  	writes := db.NewBatch()
    36  
    37  	// Mock data
    38  	data := []struct {
    39  		key []byte
    40  		val []byte
    41  		ver kv.Version
    42  	}{
    43  		{key: []byte("test"), val: []byte("test1"), ver: 100},
    44  		{key: []byte("test"), val: []byte("test3"), ver: 300},
    45  		{key: []byte("test1"), val: []byte("test5"), ver: 500},
    46  		{key: []byte("test1"), val: []byte("test7"), ver: 700},
    47  		{key: []byte("test2"), val: []byte("test9"), ver: 900},
    48  	}
    49  	wo := &pebble.WriteOptions{}
    50  	for _, d := range data {
    51  		v := mvcc.Value{
    52  			Type:      mvcc.ValueTypePut,
    53  			StartVer:  d.ver,
    54  			CommitVer: d.ver + 10,
    55  			Value:     d.val,
    56  		}
    57  		val, err := v.MarshalBinary()
    58  		assert.Nil(t, err)
    59  		assert.NotNil(t, val)
    60  		err = writes.Set(mvcc.Encode(d.key, d.ver+1), val, wo)
    61  		assert.Nil(t, err)
    62  	}
    63  
    64  	err = db.Apply(writes, wo)
    65  	assert.Nil(t, err)
    66  
    67  	expected := []struct {
    68  		key string
    69  		val string
    70  		ver kv.Version
    71  	}{
    72  		{"test", "test1", 200},
    73  		{"test", "test1", 300},
    74  		{"test", "test3", 310},
    75  		{"test", "test3", 400},
    76  		{"test1", "test5", 600},
    77  		{"test1", "test7", 750},
    78  	}
    79  	for _, e := range expected {
    80  		snapshot, err := s.Snapshot(e.ver)
    81  		assert.Nil(t, err)
    82  		val, err := snapshot.Get(context.Background(), []byte(e.key))
    83  		assert.Nil(t, err)
    84  		assert.True(t, bytes.Equal(val, []byte(e.val)), "expected: %v, got: %v (key: %s, ver:%d)",
    85  			e.val, string(val), e.key, e.ver)
    86  	}
    87  }
    88  
    89  func TestSnapshot_BatchGet(t *testing.T) {
    90  	s, err := Open(t.TempDir())
    91  	assert.Nil(t, err)
    92  	assert.NotNil(t, s)
    93  
    94  	db := s.(*mvccStorage).db
    95  	writes := db.NewBatch()
    96  
    97  	// Mock data
    98  	data := []struct {
    99  		key []byte
   100  		val []byte
   101  		ver kv.Version
   102  	}{
   103  		{key: []byte("test"), val: []byte("test1"), ver: 100},
   104  		{key: []byte("test"), val: []byte("test3"), ver: 300},
   105  		{key: []byte("test1"), val: []byte("test5"), ver: 500},
   106  		{key: []byte("test1"), val: []byte("test7"), ver: 700},
   107  		{key: []byte("test2"), val: []byte("test9"), ver: 900},
   108  	}
   109  	wo := &pebble.WriteOptions{}
   110  	for _, d := range data {
   111  		v := mvcc.Value{
   112  			Type:      mvcc.ValueTypePut,
   113  			StartVer:  d.ver,
   114  			CommitVer: d.ver + 10,
   115  			Value:     d.val,
   116  		}
   117  		val, err := v.MarshalBinary()
   118  		assert.Nil(t, err)
   119  		assert.NotNil(t, val)
   120  		err = writes.Set(mvcc.Encode(d.key, d.ver+1), val, wo)
   121  		assert.Nil(t, err)
   122  	}
   123  
   124  	err = db.Apply(writes, wo)
   125  	assert.Nil(t, err)
   126  
   127  	expected := []struct {
   128  		keys   []string
   129  		values map[string][]byte
   130  		ver    kv.Version
   131  	}{
   132  		{
   133  			keys: []string{"test", "test1", "test2"},
   134  			values: map[string][]byte{
   135  				"test":  []byte("test1"),
   136  				"test1": nil,
   137  				"test2": nil,
   138  			},
   139  			ver: 200,
   140  		},
   141  		{
   142  			keys: []string{"test", "test1", "test2"},
   143  			values: map[string][]byte{
   144  				"test":  []byte("test3"),
   145  				"test1": []byte("test5"),
   146  				"test2": nil,
   147  			},
   148  			ver: 510,
   149  		},
   150  		{
   151  			keys: []string{"test", "test1", "test2"},
   152  			values: map[string][]byte{
   153  				"test":  []byte("test3"),
   154  				"test1": []byte("test7"),
   155  				"test2": []byte("test9"),
   156  			},
   157  			ver: 1000,
   158  		},
   159  	}
   160  	for _, e := range expected {
   161  		snapshot, err := s.Snapshot(e.ver)
   162  		assert.Nil(t, err)
   163  		var keys []kv.Key
   164  		for _, k := range e.keys {
   165  			keys = append(keys, kv.Key(k))
   166  		}
   167  		values, err := snapshot.BatchGet(context.Background(), keys)
   168  		assert.Nil(t, err)
   169  		assert.True(t, reflect.DeepEqual(values, e.values), "expected: %v, got: %v (key: %s, ver:%d)",
   170  			e.values, values, e.keys, e.ver)
   171  	}
   172  }