From ccd1bb17805ad61a291aab66e70737ef8eeeccb9 Mon Sep 17 00:00:00 2001
From: Anthony Romano <anthony.romano@coreos.com>
Date: Thu, 10 Aug 2017 08:38:19 -0700
Subject: [PATCH] mvcc: test keys gauge is reloaded correctly on restore

---
 mvcc/kv_test.go | 20 ++++++++++++++++++++
 1 file changed, 20 insertions(+)

diff --git a/mvcc/kv_test.go b/mvcc/kv_test.go
index 162dac40b..d6f49ee14 100644
--- a/mvcc/kv_test.go
+++ b/mvcc/kv_test.go
@@ -25,6 +25,9 @@ import (
 	"github.com/coreos/etcd/mvcc/backend"
 	"github.com/coreos/etcd/mvcc/mvccpb"
 	"github.com/coreos/etcd/pkg/testutil"
+
+	"github.com/prometheus/client_golang/prometheus"
+	dto "github.com/prometheus/client_model/go"
 )
 
 // Functional tests for features implemented in v3 store. It treats v3 store
@@ -612,6 +615,7 @@ func TestKVRestore(t *testing.T) {
 			kv.Put([]byte("foo"), []byte("bar0"), 1)
 			kv.Put([]byte("foo"), []byte("bar1"), 2)
 			kv.Put([]byte("foo"), []byte("bar2"), 3)
+			kv.Put([]byte("foo2"), []byte("bar0"), 1)
 		},
 		func(kv KV) {
 			kv.Put([]byte("foo"), []byte("bar0"), 1)
@@ -633,10 +637,17 @@ func TestKVRestore(t *testing.T) {
 			r, _ := s.Range([]byte("a"), []byte("z"), RangeOptions{Rev: k})
 			kvss = append(kvss, r.KVs)
 		}
+
+		keysBefore := readGaugeInt(&keysGauge)
 		s.Close()
 
 		// ns should recover the the previous state from backend.
 		ns := NewStore(b, &lease.FakeLessor{}, nil)
+
+		if keysRestore := readGaugeInt(&keysGauge); keysBefore != keysRestore {
+			t.Errorf("#%d: got %d key count, expected %d", i, keysRestore, keysBefore)
+		}
+
 		// wait for possible compaction to finish
 		testutil.WaitSchedule()
 		var nkvss [][]mvccpb.KeyValue
@@ -652,6 +663,15 @@ func TestKVRestore(t *testing.T) {
 	}
 }
 
+func readGaugeInt(g *prometheus.Gauge) int {
+	ch := make(chan prometheus.Metric, 1)
+	keysGauge.Collect(ch)
+	m := <-ch
+	mm := &dto.Metric{}
+	m.Write(mm)
+	return int(mm.GetGauge().GetValue())
+}
+
 func TestKVSnapshot(t *testing.T) {
 	b, tmpPath := backend.NewDefaultTmpBackend()
 	s := NewStore(b, &lease.FakeLessor{}, nil)