From 4e5c015fe958b78fbf184b7033677126bcb62bf9 Mon Sep 17 00:00:00 2001
From: Brian Waldon <brian@waldon.cc>
Date: Wed, 21 Jan 2015 15:10:26 -0800
Subject: [PATCH] client: add Update method

---
 client/keys.go              | 20 ++++++++++++++++++++
 discovery/discovery_test.go | 16 ++++++++++------
 2 files changed, 30 insertions(+), 6 deletions(-)

diff --git a/client/keys.go b/client/keys.go
index 7537f3340..0c391de94 100644
--- a/client/keys.go
+++ b/client/keys.go
@@ -62,6 +62,8 @@ func NewDiscoveryKeysAPI(c HTTPClient) KeysAPI {
 
 type KeysAPI interface {
 	Create(ctx context.Context, key, value string) (*Response, error)
+	Update(ctx context.Context, key, value string) (*Response, error)
+
 	Get(ctx context.Context, key string) (*Response, error)
 
 	Watch(key string, idx uint64) Watcher
@@ -119,6 +121,24 @@ func (k *httpKeysAPI) Create(ctx context.Context, key, val string) (*Response, e
 	return unmarshalHTTPResponse(resp.StatusCode, resp.Header, body)
 }
 
+func (k *httpKeysAPI) Update(ctx context.Context, key, val string) (*Response, error) {
+	act := &setAction{
+		Prefix: k.prefix,
+		Key:    key,
+		Value:  val,
+		Options: SetOptions{
+			PrevExist: PrevExist,
+		},
+	}
+
+	resp, body, err := k.client.Do(ctx, act)
+	if err != nil {
+		return nil, err
+	}
+
+	return unmarshalHTTPResponse(resp.StatusCode, resp.Header, body)
+}
+
 func (k *httpKeysAPI) Get(ctx context.Context, key string) (*Response, error) {
 	get := &getAction{
 		Prefix:    k.prefix,
diff --git a/discovery/discovery_test.go b/discovery/discovery_test.go
index 01a916fca..a4cbf0ede 100644
--- a/discovery/discovery_test.go
+++ b/discovery/discovery_test.go
@@ -27,6 +27,7 @@ import (
 
 	"github.com/coreos/etcd/Godeps/_workspace/src/github.com/jonboulle/clockwork"
 	"github.com/coreos/etcd/Godeps/_workspace/src/golang.org/x/net/context"
+
 	"github.com/coreos/etcd/client"
 )
 
@@ -262,7 +263,7 @@ func TestWaitNodes(t *testing.T) {
 
 	for i, tt := range tests {
 		// Basic case
-		c := &clientWithResp{nil, &watcherWithResp{tt.rs}}
+		c := &clientWithResp{rs: nil, w: &watcherWithResp{rs: tt.rs}}
 		dBase := &discovery{cluster: "1000", c: c}
 
 		// Retry case
@@ -312,12 +313,12 @@ func TestWaitNodes(t *testing.T) {
 func TestCreateSelf(t *testing.T) {
 	rs := []*client.Response{{Node: &client.Node{Key: "1000/1", CreatedIndex: 2}}}
 
-	w := &watcherWithResp{rs}
-	errw := &watcherWithErr{errors.New("watch err")}
+	w := &watcherWithResp{rs: rs}
+	errw := &watcherWithErr{err: errors.New("watch err")}
 
-	c := &clientWithResp{rs, w}
-	errc := &clientWithErr{errors.New("create err"), w}
-	errwc := &clientWithResp{rs, errw}
+	c := &clientWithResp{rs: rs, w: w}
+	errc := &clientWithErr{err: errors.New("create err"), w: w}
+	errwc := &clientWithResp{rs: rs, w: errw}
 
 	tests := []struct {
 		c    client.KeysAPI
@@ -409,6 +410,7 @@ func TestRetryFailure(t *testing.T) {
 type clientWithResp struct {
 	rs []*client.Response
 	w  client.Watcher
+	client.KeysAPI
 }
 
 func (c *clientWithResp) Create(ctx context.Context, key string, value string) (*client.Response, error) {
@@ -440,6 +442,7 @@ func (c *clientWithResp) RecursiveWatch(key string, waitIndex uint64) client.Wat
 type clientWithErr struct {
 	err error
 	w   client.Watcher
+	client.KeysAPI
 }
 
 func (c *clientWithErr) Create(ctx context.Context, key string, value string) (*client.Response, error) {
@@ -459,6 +462,7 @@ func (c *clientWithErr) RecursiveWatch(key string, waitIndex uint64) client.Watc
 }
 
 type watcherWithResp struct {
+	client.KeysAPI
 	rs []*client.Response
 }