Unit test snapshot and restore

store/store_test.go

@@ -4,14 +4,27 @@ import (
 	"encoding/json"
 	"io/ioutil"
 	"os"
+	"path/filepath"
 	"testing"
 	"time"
 
 	sql "github.com/otoolep/rqlite/db"
 )
 
+type mockSnapshotSink struct {
+	*os.File
+}
+
+func (m *mockSnapshotSink) ID() string {
+	return "1"
+}
+
+func (m *mockSnapshotSink) Cancel() error {
+	return nil
+}
+
 func Test_OpenStoreSingleNode(t *testing.T) {
-	s := mustNewStore()
+	s := mustNewStore(true)
 	defer os.RemoveAll(s.Path())
 
 	if err := s.Open(true); err != nil {
@@ -20,7 +33,7 @@ func Test_OpenStoreSingleNode(t *testing.T) {
 }
 
 func Test_OpenStoreCloseSingleNode(t *testing.T) {
-	s := mustNewStore()
+	s := mustNewStore(true)
 	defer os.RemoveAll(s.Path())
 
 	if err := s.Open(true); err != nil {
@@ -32,7 +45,7 @@ func Test_OpenStoreCloseSingleNode(t *testing.T) {
 }
 
 func Test_SingleNodeExecuteQuery(t *testing.T) {
-	s := mustNewStore()
+	s := mustNewStore(true)
 	defer os.RemoveAll(s.Path())
 
 	if err := s.Open(true); err != nil {
@@ -70,7 +83,7 @@ func Test_SingleNodeExecuteQuery(t *testing.T) {
 }
 
 func Test_SingleNodeExecuteQueryTx(t *testing.T) {
-	s := mustNewStore()
+	s := mustNewStore(true)
 	defer os.RemoveAll(s.Path())
 
 	if err := s.Open(true); err != nil {
@@ -112,7 +125,7 @@ func Test_SingleNodeExecuteQueryTx(t *testing.T) {
 }
 
 func Test_MultiNodeExecuteQuery(t *testing.T) {
-	s0 := mustNewStore()
+	s0 := mustNewStore(true)
 	defer os.RemoveAll(s0.Path())
 	if err := s0.Open(true); err != nil {
 		t.Fatalf("failed to open node for multi-node test: %s", err.Error())
@@ -120,7 +133,7 @@ func Test_MultiNodeExecuteQuery(t *testing.T) {
 	defer s0.Close(true)
 	s0.WaitForLeader(10 * time.Second)
 
-	s1 := mustNewStore()
+	s1 := mustNewStore(true)
 	defer os.RemoveAll(s1.Path())
 	if err := s1.Open(false); err != nil {
 		t.Fatalf("failed to open node for multi-node test: %s", err.Error())
@@ -176,11 +189,75 @@ func Test_MultiNodeExecuteQuery(t *testing.T) {
 	}
 }
 
-func mustNewStore() *Store {
+func Test_SingleNodeSnapshot(t *testing.T) {
+	s := mustNewStore(false)
+	defer os.RemoveAll(s.Path())
+
+	if err := s.Open(true); err != nil {
+		t.Fatalf("failed to open single-node store: %s", err.Error())
+	}
+	defer s.Close(true)
+	s.WaitForLeader(10 * time.Second)
+
+	queries := []string{
+		`CREATE TABLE foo (id INTEGER NOT NULL PRIMARY KEY, name TEXT)`,
+		`INSERT INTO foo(id, name) VALUES(1, "fiona")`,
+	}
+	_, err := s.Execute(queries, false, false)
+	if err != nil {
+		t.Fatalf("failed to execute on single node: %s", err.Error())
+	}
+	_, err = s.Query([]string{`SELECT * FROM foo`}, false, false, None)
+	if err != nil {
+		t.Fatalf("failed to query single node: %s", err.Error())
+	}
+
+	// Snap the node and write to disk.
+	f, err := s.Snapshot()
+	if err != nil {
+		t.Fatalf("failed to snapshot node: %s", err.Error())
+	}
+
+	snapDir := mustTempDir()
+	defer os.RemoveAll(snapDir)
+	snapFile, err := os.Create(filepath.Join(snapDir, "snapshot"))
+	if err != nil {
+		t.Fatalf("failed to create snapshot file: %s", err.Error())
+	}
+	sink := &mockSnapshotSink{snapFile}
+	if err := f.Persist(sink); err != nil {
+		t.Fatalf("failed to persist snapshot to disk: %s", err.Error())
+	}
+
+	// Check restoration.
+	snapFile, err = os.Open(filepath.Join(snapDir, "snapshot"))
+	if err != nil {
+		t.Fatalf("failed to open snapshot file: %s", err.Error())
+	}
+	if err := s.Restore(snapFile); err != nil {
+		t.Fatalf("failed to restore snapshot from disk: %s", err.Error())
+	}
+
+	// Ensure database is back in the correct state.
+	r, err := s.Query([]string{`SELECT * FROM foo`}, false, false, None)
+	if err != nil {
+		t.Fatalf("failed to query single node: %s", err.Error())
+	}
+	if exp, got := `["id","name"]`, asJSON(r[0].Columns); exp != got {
+		t.Fatalf("unexpected results for query\nexp: %s\ngot: %s", exp, got)
+	}
+	if exp, got := `[[1,"fiona"]]`, asJSON(r[0].Values); exp != got {
+		t.Fatalf("unexpected results for query\nexp: %s\ngot: %s", exp, got)
+	}
+}
+
+func mustNewStore(inmem bool) *Store {
 	path := mustTempDir()
 	defer os.RemoveAll(path)
 
-	s := New(newInMemoryConfig(), path, "localhost:0")
+	cfg := sql.NewConfig()
+	cfg.Memory = inmem
+	s := New(cfg, path, "localhost:0")
 	if s == nil {
 		panic("failed to create new store")
 	}
@@ -196,12 +273,6 @@ func mustTempDir() string {
 	return path
 }
 
-func newInMemoryConfig() *sql.Config {
-	c := sql.NewConfig()
-	c.Memory = true
-	return c
-}
-
 func asJSON(v interface{}) string {
 	b, err := json.Marshal(v)
 	if err != nil {