From 95b8953bac5fde460473761295539ca34f581bd2 Mon Sep 17 00:00:00 2001
From: congqixia <congqi.xia@zilliz.com>
Date: Fri, 29 Oct 2021 00:16:38 +0800
Subject: [PATCH] Add kv Value interface (#10469)

Signed-off-by: Congqi Xia <congqi.xia@zilliz.com>
---
 internal/kv/kv.go      | 34 ++++++++++++++++++++++++
 internal/kv/kv_test.go | 59 ++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 93 insertions(+)
 create mode 100644 internal/kv/kv_test.go

diff --git a/internal/kv/kv.go b/internal/kv/kv.go
index 075dc7b08a..3f5d65ddc2 100644
--- a/internal/kv/kv.go
+++ b/internal/kv/kv.go
@@ -16,6 +16,40 @@ import (
 	clientv3 "go.etcd.io/etcd/client/v3"
 )
 
+// Value is interface for kv-value, needed to support string and byte slice
+type Value interface {
+	Serialize() []byte
+	String() string
+}
+
+// StringValue type alias for string to implement Value
+type StringValue string
+
+func (s StringValue) Serialize() []byte {
+	return []byte(s)
+}
+
+func (s StringValue) String() string {
+	return string(s)
+}
+
+// BytesValue type alias for byte slice to implement value
+type BytesValue []byte
+
+func (s BytesValue) Serialize() []byte {
+	return s
+}
+
+func (s BytesValue) String() string {
+	return string(s)
+}
+
+// ValueKV example usage to change KV interface to
+type ValueKV interface {
+	Save(key string, value Value) error
+	Load(key string) (Value, error)
+}
+
 // BaseKV contains base operations of kv. Include save, load and remove.
 type BaseKV interface {
 	Load(key string) (string, error)
diff --git a/internal/kv/kv_test.go b/internal/kv/kv_test.go
new file mode 100644
index 0000000000..482886a33b
--- /dev/null
+++ b/internal/kv/kv_test.go
@@ -0,0 +1,59 @@
+// Copyright (C) 2019-2020 Zilliz. All rights reserved.
+//
+// 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 kv
+
+import (
+	"errors"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+type mockTestKV struct {
+	data map[string][]byte
+}
+
+func newMockTestKV() ValueKV {
+	return &mockTestKV{
+		data: make(map[string][]byte),
+	}
+}
+
+func (m *mockTestKV) Save(key string, value Value) error {
+	m.data[key] = value.Serialize()
+	return nil
+}
+
+func (m *mockTestKV) Load(key string) (Value, error) {
+	d, ok := m.data[key]
+	if !ok {
+		return nil, errors.New("not found")
+	}
+	return BytesValue(d), nil
+}
+
+func TestTKV(t *testing.T) {
+	tkv := newMockTestKV()
+	bv := BytesValue([]byte{1, 2, 3})
+	sv := StringValue("test")
+	tkv.Save("1", bv)
+	tkv.Save("2", sv)
+
+	v, err := tkv.Load("1")
+	assert.NoError(t, err)
+
+	assert.EqualValues(t, bv.Serialize(), v.Serialize())
+
+	v, err = tkv.Load("2")
+	assert.NoError(t, err)
+	assert.EqualValues(t, sv.String(), v.String())
+}