Add ttl parameter to CreateObj

duda · Sep 25, 2014 · 4d4ebcd · 4d4ebcd
1 parent 3ea383d
commit 4d4ebcd
Show file tree

Hide file tree

Showing 4 changed files with 30 additions and 6 deletions.
diff --git a/pkg/registry/etcd/etcd.go b/pkg/registry/etcd/etcd.go
@@ -119,7 +119,7 @@ func (r *Registry) CreatePod(pod *api.Pod) error {
 	// DesiredState.Host == "" is a signal to the scheduler that this pod needs scheduling.
 	pod.DesiredState.Status = api.PodRunning
 	pod.DesiredState.Host = ""
-	err := r.CreateObj(makePodKey(pod.ID), pod)
+	err := r.CreateObj(makePodKey(pod.ID), pod, 0)
 	return etcderr.InterpretCreateError(err, "pod", pod.ID)
 }
 
@@ -253,7 +253,7 @@ func (r *Registry) GetController(controllerID string) (*api.ReplicationControlle
 
 // CreateController creates a new ReplicationController.
 func (r *Registry) CreateController(controller *api.ReplicationController) error {
-	err := r.CreateObj(makeControllerKey(controller.ID), controller)
+	err := r.CreateObj(makeControllerKey(controller.ID), controller, 0)
 	return etcderr.InterpretCreateError(err, "replicationController", controller.ID)
 }
 
@@ -283,7 +283,7 @@ func (r *Registry) ListServices() (*api.ServiceList, error) {
 
 // CreateService creates a new Service.
 func (r *Registry) CreateService(svc *api.Service) error {
-	err := r.CreateObj(makeServiceKey(svc.ID), svc)
+	err := r.CreateObj(makeServiceKey(svc.ID), svc, 0)
 	return etcderr.InterpretCreateError(err, "service", svc.ID)
 }
 

diff --git a/pkg/tools/etcd_tools.go b/pkg/tools/etcd_tools.go
@@ -207,8 +207,9 @@ func (h *EtcdHelper) bodyAndExtractObj(key string, objPtr runtime.Object, ignore
 	return body, response.Node.ModifiedIndex, err
 }
 
-// CreateObj adds a new object at a key unless it already exists.
-func (h *EtcdHelper) CreateObj(key string, obj runtime.Object) error {
+// CreateObj adds a new object at a key unless it already exists. 'ttl' is time-to-live in seconds,
+// and 0 means forever.
+func (h *EtcdHelper) CreateObj(key string, obj runtime.Object, ttl uint64) error {
 	data, err := h.Codec.Encode(obj)
 	if err != nil {
 		return err
@@ -219,7 +220,7 @@ func (h *EtcdHelper) CreateObj(key string, obj runtime.Object) error {
 		}
 	}
 
-	_, err = h.Client.Create(key, string(data), 0)
+	_, err = h.Client.Create(key, string(data), ttl)
 	return err
 }
 

diff --git a/pkg/tools/etcd_tools_test.go b/pkg/tools/etcd_tools_test.go
@@ -163,6 +163,27 @@ func TestExtractObjNotFoundErr(t *testing.T) {
 	try("/some/key3")
 }
 
+func TestCreateObj(t *testing.T) {
+	obj := &api.Pod{JSONBase: api.JSONBase{ID: "foo"}}
+	fakeClient := NewFakeEtcdClient(t)
+	helper := EtcdHelper{fakeClient, latest.Codec, versioner}
+	err := helper.CreateObj("/some/key", obj, 5)
+	if err != nil {
+		t.Errorf("Unexpected error %#v", err)
+	}
+	data, err := latest.Codec.Encode(obj)
+	if err != nil {
+		t.Errorf("Unexpected error %#v", err)
+	}
+	node := fakeClient.Data["/some/key"].R.Node
+	if e, a := string(data), node.Value; e != a {
+		t.Errorf("Wanted %v, got %v", e, a)
+	}
+	if e, a := uint64(5), fakeClient.LastSetTTL; e != a {
+		t.Errorf("Wanted %v, got %v", e, a)
+	}
+}
+
 func TestSetObj(t *testing.T) {
 	obj := &api.Pod{JSONBase: api.JSONBase{ID: "foo"}}
 	fakeClient := NewFakeEtcdClient(t)

diff --git a/pkg/tools/fake_etcd_client.go b/pkg/tools/fake_etcd_client.go
@@ -49,6 +49,7 @@ type FakeEtcdClient struct {
 	Ix          int
 	TestIndex   bool
 	ChangeIndex uint64
+	LastSetTTL  uint64
 
 	// Will become valid after Watch is called; tester may write to it. Tester may
 	// also read from it to verify that it's closed after injecting an error.
@@ -135,6 +136,7 @@ func (f *FakeEtcdClient) nodeExists(key string) bool {
 }
 
 func (f *FakeEtcdClient) setLocked(key, value string, ttl uint64) (*etcd.Response, error) {
+	f.LastSetTTL = ttl
 	if f.Err != nil {
 		return nil, f.Err
 	}