150 lines
3.7 KiB
Go
150 lines
3.7 KiB
Go
// +build kvdb_etcd
|
|
|
|
package etcd
|
|
|
|
import (
|
|
"context"
|
|
"io/ioutil"
|
|
"os"
|
|
"testing"
|
|
"time"
|
|
|
|
"github.com/btcsuite/btcwallet/walletdb"
|
|
"github.com/stretchr/testify/require"
|
|
"go.etcd.io/etcd/clientv3"
|
|
"go.etcd.io/etcd/clientv3/namespace"
|
|
)
|
|
|
|
const (
|
|
// testEtcdTimeout is used for all RPC calls initiated by the test fixture.
|
|
testEtcdTimeout = 5 * time.Second
|
|
)
|
|
|
|
// EtcdTestFixture holds internal state of the etcd test fixture.
|
|
type EtcdTestFixture struct {
|
|
t *testing.T
|
|
cli *clientv3.Client
|
|
config *Config
|
|
cleanup func()
|
|
}
|
|
|
|
// NewTestEtcdInstance creates an embedded etcd instance for testing, listening
|
|
// on random open ports. Returns the connection config and a cleanup func that
|
|
// will stop the etcd instance.
|
|
func NewTestEtcdInstance(t *testing.T, path string) (*Config, func()) {
|
|
t.Helper()
|
|
|
|
config, cleanup, err := NewEmbeddedEtcdInstance(path, 0, 0)
|
|
if err != nil {
|
|
t.Fatalf("error while staring embedded etcd instance: %v", err)
|
|
}
|
|
|
|
return config, cleanup
|
|
}
|
|
|
|
// NewTestEtcdTestFixture creates a new etcd-test fixture. This is helper
|
|
// object to facilitate etcd tests and ensure pre and post conditions.
|
|
func NewEtcdTestFixture(t *testing.T) *EtcdTestFixture {
|
|
tmpDir, err := ioutil.TempDir("", "etcd")
|
|
if err != nil {
|
|
t.Fatalf("unable to create temp dir: %v", err)
|
|
}
|
|
|
|
config, etcdCleanup := NewTestEtcdInstance(t, tmpDir)
|
|
|
|
cli, err := clientv3.New(clientv3.Config{
|
|
Endpoints: []string{config.Host},
|
|
Username: config.User,
|
|
Password: config.Pass,
|
|
})
|
|
if err != nil {
|
|
os.RemoveAll(tmpDir)
|
|
t.Fatalf("unable to create etcd test fixture: %v", err)
|
|
}
|
|
|
|
// Apply the default namespace (since that's what we use in tests).
|
|
cli.KV = namespace.NewKV(cli.KV, defaultNamespace)
|
|
cli.Watcher = namespace.NewWatcher(cli.Watcher, defaultNamespace)
|
|
cli.Lease = namespace.NewLease(cli.Lease, defaultNamespace)
|
|
|
|
return &EtcdTestFixture{
|
|
t: t,
|
|
cli: cli,
|
|
config: config,
|
|
cleanup: func() {
|
|
etcdCleanup()
|
|
os.RemoveAll(tmpDir)
|
|
},
|
|
}
|
|
}
|
|
|
|
func (f *EtcdTestFixture) NewBackend(singleWriter bool) walletdb.DB {
|
|
cfg := f.BackendConfig()
|
|
if singleWriter {
|
|
cfg.SingleWriter = true
|
|
}
|
|
|
|
db, err := newEtcdBackend(context.TODO(), cfg)
|
|
require.NoError(f.t, err)
|
|
|
|
return db
|
|
}
|
|
|
|
// Put puts a string key/value into the test etcd database.
|
|
func (f *EtcdTestFixture) Put(key, value string) {
|
|
ctx, cancel := context.WithTimeout(context.TODO(), testEtcdTimeout)
|
|
defer cancel()
|
|
|
|
_, err := f.cli.Put(ctx, key, value)
|
|
if err != nil {
|
|
f.t.Fatalf("etcd test fixture failed to put: %v", err)
|
|
}
|
|
}
|
|
|
|
// Get queries a key and returns the stored value from the test etcd database.
|
|
func (f *EtcdTestFixture) Get(key string) string {
|
|
ctx, cancel := context.WithTimeout(context.TODO(), testEtcdTimeout)
|
|
defer cancel()
|
|
|
|
resp, err := f.cli.Get(ctx, key)
|
|
if err != nil {
|
|
f.t.Fatalf("etcd test fixture failed to get: %v", err)
|
|
}
|
|
|
|
if len(resp.Kvs) > 0 {
|
|
return string(resp.Kvs[0].Value)
|
|
}
|
|
|
|
return ""
|
|
}
|
|
|
|
// Dump scans and returns all key/values from the test etcd database.
|
|
func (f *EtcdTestFixture) Dump() map[string]string {
|
|
ctx, cancel := context.WithTimeout(context.TODO(), testEtcdTimeout)
|
|
defer cancel()
|
|
|
|
resp, err := f.cli.Get(ctx, "\x00", clientv3.WithFromKey())
|
|
if err != nil {
|
|
f.t.Fatalf("etcd test fixture failed to get: %v", err)
|
|
}
|
|
|
|
result := make(map[string]string)
|
|
for _, kv := range resp.Kvs {
|
|
result[string(kv.Key)] = string(kv.Value)
|
|
}
|
|
|
|
return result
|
|
}
|
|
|
|
// BackendConfig returns the backend config for connecting to theembedded
|
|
// etcd instance.
|
|
func (f *EtcdTestFixture) BackendConfig() Config {
|
|
return *f.config
|
|
}
|
|
|
|
// Cleanup should be called at test fixture teardown to stop the embedded
|
|
// etcd instance and remove all temp db files form the filesystem.
|
|
func (f *EtcdTestFixture) Cleanup() {
|
|
f.cleanup()
|
|
}
|