From 68a829119ea457e1324ee2ec91106a36804b9a57 Mon Sep 17 00:00:00 2001
From: Armon Dadgar <armon.dadgar@gmail.com>
Date: Wed, 20 Aug 2014 15:18:08 -0700
Subject: [PATCH] watch: test key watch

---
 watch/funcs.go     |  3 +++
 watch/plan.go      |  7 +++++--
 watch/plan_test.go | 11 ++++++++---
 3 files changed, 16 insertions(+), 5 deletions(-)

diff --git a/watch/funcs.go b/watch/funcs.go
index 7593a468fb..5b39320f71 100644
--- a/watch/funcs.go
+++ b/watch/funcs.go
@@ -35,6 +35,9 @@ func keyWatch(params map[string][]string) (WatchFunc, error) {
 		if err != nil {
 			return 0, nil, err
 		}
+		if pair == nil {
+			return meta.LastIndex, nil, err
+		}
 		return meta.LastIndex, pair, err
 	}
 	return fn, nil
diff --git a/watch/plan.go b/watch/plan.go
index a43849fc30..2faea27d85 100644
--- a/watch/plan.go
+++ b/watch/plan.go
@@ -34,6 +34,7 @@ func (p *WatchPlan) Run(address string) error {
 
 	// Loop until we are canceled
 	failures := 0
+OUTER:
 	for !p.shouldStop() {
 		// Invoke the handler
 		index, result, err := p.Func(p)
@@ -56,7 +57,7 @@ func (p *WatchPlan) Run(address string) error {
 			}
 			select {
 			case <-time.After(retry):
-				continue
+				continue OUTER
 			case <-p.stopCh:
 				return nil
 			}
@@ -78,7 +79,9 @@ func (p *WatchPlan) Run(address string) error {
 
 		// Handle the updated result
 		p.lastResult = result
-		p.Handler(index, result)
+		if p.Handler != nil {
+			p.Handler(index, result)
+		}
 	}
 	return nil
 }
diff --git a/watch/plan_test.go b/watch/plan_test.go
index 99cae3df53..d898495005 100644
--- a/watch/plan_test.go
+++ b/watch/plan_test.go
@@ -17,11 +17,16 @@ func noopWatch(params map[string][]string) (WatchFunc, error) {
 	return fn, nil
 }
 
-func TestRun_Stop(t *testing.T) {
-	plan, err := Parse("type:noop")
+func mustParse(t *testing.T, q string) *WatchPlan {
+	plan, err := Parse(q)
 	if err != nil {
 		t.Fatalf("err: %v", err)
 	}
+	return plan
+}
+
+func TestRun_Stop(t *testing.T) {
+	plan := mustParse(t, "type:noop")
 	var expect uint64 = 1
 	plan.Handler = func(idx uint64, val interface{}) {
 		if idx != expect {
@@ -37,7 +42,7 @@ func TestRun_Stop(t *testing.T) {
 		plan.Stop()
 	})
 
-	err = plan.Run("127.0.0.1:8500")
+	err := plan.Run("127.0.0.1:8500")
 	if err != nil {
 		t.Fatalf("err: %v", err)
 	}