Adding the update() method that allows atomic updates on results

lavagetto · lavagetto · commit eee2785663f0 · 2014-03-16T18:18:16.000+01:00
diff --git a/src/etcd/client.py b/src/etcd/client.py
@@ -50,7 +50,7 @@ def __init__(
             read_timeout (int):  max seconds to wait for a read.
 
             allow_redirect (bool): allow the client to connect to other nodes.
-
++
             protocol (str):  Protocol used to connect to etcd.
 
             cert (mixed):   If a string, the whole ssl client certificate;
@@ -207,6 +207,7 @@ def _sanitize_key(self, key):
             key = "/{}".format(key)
         return key
 
+
     def write(self, key, value, ttl=None, dir=False, append=False, **kwdargs):
         """
         Writes the value for a key, possibly doing atomit Compare-and-Swap
@@ -268,6 +269,29 @@ def write(self, key, value, ttl=None, dir=False, append=False, **kwdargs):
         response = self.api_execute(path, method, params=params)
         return self._result_from_response(response)
 
+    def update(self, obj):
+        """
+        Updates the value for a key atomically. Typical usage would be:
+
+        c = etcd.Client()
+        o = c.read("/somekey")
+        o.value += 1
+        c.update(o)
+
+        Args:
+            obj (etcd.EtcdResult):  The object that needs updating.
+
+        """
+        return self.write(
+            obj.key,
+            obj.value,
+            ttl=obj.ttl,
+            dir=obj.dir,
+            prevExist=True,
+            prevIndex=obj.modifiedIndex)
+
+
+
     def read(self, key, **kwdargs):
         """
         Returns the value of the key 'key'.
diff --git a/src/etcd/tests/integration/test_simple.py b/src/etcd/tests/integration/test_simple.py
@@ -103,6 +103,16 @@ def test_get_set_delete(self):
         except KeyError as e:
             pass
 
+    def test_update(self):
+        """INTEGRATION: update a value"""
+        self.client.set('/foo', 3)
+        c = self.client.get('/foo')
+        c.value = int(c.value) + 3
+        self.client.update(c)
+        newres = self.client.get('/foo')
+        self.assertEquals(newres.value, u'6')
+        self.assertRaises(ValueError, self.client.update, c)
+
     def test_retrieve_subkeys(self):
         """ INTEGRATION: retrieve multiple subkeys """
         set_result = self.client.write('/subtree/test_set', 'test-key1')
diff --git a/src/etcd/tests/unit/test_request.py b/src/etcd/tests/unit/test_request.py
@@ -113,6 +113,26 @@ def test_set_plain(self):
         res = self.client.write('/testkey', 'test')
         self.assertEquals(res, etcd.EtcdResult(**d))
 
+
+    def test_update(self):
+        """Can update a result."""
+        d = {u'action': u'set',
+             u'node': {
+                 u'expiration': u'2013-09-14T00:56:59.316195568+02:00',
+                 u'modifiedIndex': 6,
+                 u'key': u'/testkey',
+                 u'ttl': 19,
+                 u'value': u'test'
+                 }
+             }
+        self._mock_api(200,d)
+        res = self.client.get('/testkey')
+        res.value = 'ciao'
+        d['node']['value'] = 'ciao'
+        self._mock_api(200,d)
+        newres = self.client.update(res)
+        self.assertEquals(newres.value, 'ciao')
+
     def test_newkey(self):
         """ Can set a new value """
         d = {
@@ -130,6 +150,8 @@ def test_newkey(self):
         d['node']['newKey'] = True
         self.assertEquals(res, etcd.EtcdResult(**d))
 
+
+
     def test_not_found_response(self):
         """ Can handle server not found response """
         self._mock_api(404, 'Not found')
@@ -444,3 +466,25 @@ def test_not_in(self):
 
     def test_in(self):
         pass
+
+    def test_update_fails(self):
+        """ Non-atomic updates fail """
+        d = {u'action': u'set',
+             u'node': {
+                u'expiration': u'2013-09-14T00:56:59.316195568+02:00',
+                u'modifiedIndex': 6,
+                u'key': u'/testkey',
+                u'ttl': 19,
+                u'value': u'test'
+                }
+             }
+        res = etcd.EtcdResult(**d)
+
+        error = {
+            "errorCode":101,
+            "message":"Compare failed",
+            "cause":"[ != bar] [7 != 6]",
+            "index":6}
+        self._mock_api(412, error)
+        res.value = 'bar'
+        self.assertRaises(ValueError, self.client.update, res)
