github.com/decred/dcrlnd@v0.7.6/kvdb/etcd/fixture.go (about)

     1  //go:build kvdb_etcd
     2  // +build kvdb_etcd
     3  
     4  package etcd
     5  
     6  import (
     7  	"context"
     8  	"io/ioutil"
     9  	"os"
    10  	"testing"
    11  	"time"
    12  
    13  	"github.com/btcsuite/btcwallet/walletdb"
    14  	"github.com/stretchr/testify/require"
    15  	clientv3 "go.etcd.io/etcd/client/v3"
    16  	"go.etcd.io/etcd/client/v3/namespace"
    17  )
    18  
    19  const (
    20  	// testEtcdTimeout is used for all RPC calls initiated by the test fixture.
    21  	testEtcdTimeout = 5 * time.Second
    22  )
    23  
    24  // EtcdTestFixture holds internal state of the etcd test fixture.
    25  type EtcdTestFixture struct {
    26  	t       *testing.T
    27  	cli     *clientv3.Client
    28  	config  *Config
    29  	cleanup func()
    30  }
    31  
    32  // NewTestEtcdInstance creates an embedded etcd instance for testing, listening
    33  // on random open ports. Returns the connection config and a cleanup func that
    34  // will stop the etcd instance.
    35  func NewTestEtcdInstance(t *testing.T, path string) (*Config, func()) {
    36  	t.Helper()
    37  
    38  	config, cleanup, err := NewEmbeddedEtcdInstance(path, 0, 0, "")
    39  	if err != nil {
    40  		t.Fatalf("error while staring embedded etcd instance: %v", err)
    41  	}
    42  
    43  	return config, cleanup
    44  }
    45  
    46  // NewTestEtcdTestFixture creates a new etcd-test fixture. This is helper
    47  // object to facilitate etcd tests and ensure pre and post conditions.
    48  func NewEtcdTestFixture(t *testing.T) *EtcdTestFixture {
    49  	tmpDir, err := ioutil.TempDir("", "etcd")
    50  	if err != nil {
    51  		t.Fatalf("unable to create temp dir: %v", err)
    52  	}
    53  
    54  	config, etcdCleanup := NewTestEtcdInstance(t, tmpDir)
    55  
    56  	cli, err := clientv3.New(clientv3.Config{
    57  		Endpoints: []string{config.Host},
    58  		Username:  config.User,
    59  		Password:  config.Pass,
    60  	})
    61  	if err != nil {
    62  		os.RemoveAll(tmpDir)
    63  		t.Fatalf("unable to create etcd test fixture: %v", err)
    64  	}
    65  
    66  	// Apply the default namespace (since that's what we use in tests).
    67  	cli.KV = namespace.NewKV(cli.KV, defaultNamespace)
    68  	cli.Watcher = namespace.NewWatcher(cli.Watcher, defaultNamespace)
    69  	cli.Lease = namespace.NewLease(cli.Lease, defaultNamespace)
    70  
    71  	return &EtcdTestFixture{
    72  		t:      t,
    73  		cli:    cli,
    74  		config: config,
    75  		cleanup: func() {
    76  			etcdCleanup()
    77  			os.RemoveAll(tmpDir)
    78  		},
    79  	}
    80  }
    81  
    82  func (f *EtcdTestFixture) NewBackend(singleWriter bool) walletdb.DB {
    83  	cfg := f.BackendConfig()
    84  	if singleWriter {
    85  		cfg.SingleWriter = true
    86  	}
    87  
    88  	db, err := newEtcdBackend(context.TODO(), cfg)
    89  	require.NoError(f.t, err)
    90  
    91  	return db
    92  }
    93  
    94  // Put puts a string key/value into the test etcd database.
    95  func (f *EtcdTestFixture) Put(key, value string) {
    96  	ctx, cancel := context.WithTimeout(context.TODO(), testEtcdTimeout)
    97  	defer cancel()
    98  
    99  	_, err := f.cli.Put(ctx, key, value)
   100  	if err != nil {
   101  		f.t.Fatalf("etcd test fixture failed to put: %v", err)
   102  	}
   103  }
   104  
   105  // Get queries a key and returns the stored value from the test etcd database.
   106  func (f *EtcdTestFixture) Get(key string) string {
   107  	ctx, cancel := context.WithTimeout(context.TODO(), testEtcdTimeout)
   108  	defer cancel()
   109  
   110  	resp, err := f.cli.Get(ctx, key)
   111  	if err != nil {
   112  		f.t.Fatalf("etcd test fixture failed to get: %v", err)
   113  	}
   114  
   115  	if len(resp.Kvs) > 0 {
   116  		return string(resp.Kvs[0].Value)
   117  	}
   118  
   119  	return ""
   120  }
   121  
   122  // Dump scans and returns all key/values from the test etcd database.
   123  func (f *EtcdTestFixture) Dump() map[string]string {
   124  	ctx, cancel := context.WithTimeout(context.TODO(), testEtcdTimeout)
   125  	defer cancel()
   126  
   127  	resp, err := f.cli.Get(ctx, "\x00", clientv3.WithFromKey())
   128  	if err != nil {
   129  		f.t.Fatalf("etcd test fixture failed to get: %v", err)
   130  	}
   131  
   132  	result := make(map[string]string)
   133  	for _, kv := range resp.Kvs {
   134  		result[string(kv.Key)] = string(kv.Value)
   135  	}
   136  
   137  	return result
   138  }
   139  
   140  // BackendConfig returns the backend config for connecting to theembedded
   141  // etcd instance.
   142  func (f *EtcdTestFixture) BackendConfig() Config {
   143  	return *f.config
   144  }
   145  
   146  // Cleanup should be called at test fixture teardown to stop the embedded
   147  // etcd instance and remove all temp db files form the filesystem.
   148  func (f *EtcdTestFixture) Cleanup() {
   149  	f.cleanup()
   150  }