Browse Source

integration: support watch apis in cluster_proxy build

Anthony Romano 9 years ago
parent
commit
1cad722a6d
1 changed files with 9 additions and 1 deletions
  1. 9 1
      integration/cluster_proxy.go

+ 9 - 1
integration/cluster_proxy.go

@@ -22,12 +22,17 @@ import (
 	"github.com/coreos/etcd/proxy/grpcproxy"
 )
 
+var proxies map[*clientv3.Client]grpcAPI = make(map[*clientv3.Client]grpcAPI)
+
 func toGRPC(c *clientv3.Client) grpcAPI {
+	if v, ok := proxies[c]; ok {
+		return v
+	}
 	return grpcAPI{
 		pb.NewClusterClient(c.ActiveConnection()),
 		grpcproxy.KvServerToKvClient(grpcproxy.NewKvProxy(c)),
 		pb.NewLeaseClient(c.ActiveConnection()),
-		pb.NewWatchClient(c.ActiveConnection()),
+		grpcproxy.WatchServerToWatchClient(grpcproxy.NewWatchProxy(c)),
 		pb.NewMaintenanceClient(c.ActiveConnection()),
 	}
 }
@@ -37,6 +42,9 @@ func newClientV3(cfg clientv3.Config) (*clientv3.Client, error) {
 	if err != nil {
 		return nil, err
 	}
+
+	proxies[c] = toGRPC(c)
 	c.KV = clientv3.NewKVFromKVClient(grpcproxy.KvServerToKvClient(grpcproxy.NewKvProxy(c)))
+	c.Watcher = clientv3.NewWatchFromWatchClient(grpcproxy.WatchServerToWatchClient(grpcproxy.NewWatchProxy(c)))
 	return c, nil
 }