From 4177a974e00be0b2781a0d486d2d958bb1ceb2e9 Mon Sep 17 00:00:00 2001
From: "Gavin M. Roy" <gavinmroy@gmail.com>
Date: Wed, 13 May 2015 01:10:20 -0400
Subject: [PATCH] KV fixes

- Make get_item make a little more sense
- rstrip the / on a non-path key with a trailing slash
- Handle the response body encoding a little better for CAS
- Add additional acceptance tests for KV
---
 consulate/api/kv.py       | 14 ++++++-------
 tests/acceptance-tests.py | 43 ++++++++++++++++++++++++++++++++++++++-
 2 files changed, 49 insertions(+), 8 deletions(-)

diff --git a/consulate/api/kv.py b/consulate/api/kv.py
index 9890913..c35026d 100644
--- a/consulate/api/kv.py
+++ b/consulate/api/kv.py
@@ -126,7 +126,7 @@ def get(self, item, default=None, raw=False):
             return response.get('Value', default)
         return response or default
 
-    def get_record(self, item, default=None):
+    def get_record(self, item):
         """Get the full record from the Key/Value service, returning
         all fields including the flag.
 
@@ -135,10 +135,7 @@ def get_record(self, item, default=None):
         :raises: KeyError
 
         """
-        try:
-            return self._get_item(item)
-        except KeyError:
-            return default
+        return self._get_item(item)
 
     def find(self, prefix, separator=None):
         """Find all keys with the specified prefix, returning a dict of
@@ -337,7 +334,10 @@ def _get_modify_index(self, item, value, replace):
         index = 0
         if response.status_code == 200:
             index = response.body.get('ModifyIndex')
-            if response.body.get('Value') == value:
+            rvalue = response.body.get('Value')
+            if isinstance(rvalue, bytes):
+                rvalue = rvalue.encode('utf-8')
+            if rvalue == value:
                 return None
             if not replace:
                 return None
@@ -374,7 +374,7 @@ def _set_item(self, item, value, flags=None, replace=True):
         """
         value = self._prepare_value(value)
         if value and item.endswith('/'):
-            item = item.lstrip('/')
+            item = item.rstrip('/')
 
         index = self._get_modify_index(item, value, replace)
         if index is None:
diff --git a/tests/acceptance-tests.py b/tests/acceptance-tests.py
index 80776dc..2ade28d 100644
--- a/tests/acceptance-tests.py
+++ b/tests/acceptance-tests.py
@@ -117,11 +117,23 @@ def test_get_item_raises_key_error(self, key):
 
 class TestKVSet(BaseTestCase):
 
+    @generate_key
+    def test_set_item_del_item(self, key):
+        self.consul.kv[key] = 'foo'
+        del self.consul.kv[key]
+        self.assertNotIn(key, self.consul.kv)
+
     @generate_key
     def test_set_item_get_item_bool_value(self, key):
         self.consul.kv[key] = True
         self.assertTrue(self.consul.kv[key])
 
+    @generate_key
+    def test_set_path_with_value(self, key):
+        path = 'path/{0}/'.format(key)
+        self.consul.kv.set(path, 'bar')
+        self.assertEqual('bar', self.consul.kv[path[:-1]])
+
     @generate_key
     def test_set_item_get_item_int_value(self, key):
         self.consul.kv[key] = 128
@@ -152,6 +164,30 @@ def test_set_item_get_item_str_value(self, key):
         self.consul.kv.set(key, 'foo')
         self.assertEqual(self.consul.kv.get(key), 'foo')
 
+    @generate_key
+    def test_set_item_get_record(self, key):
+        self.consul.kv.set_record(key, 12, 'record')
+        record = self.consul.kv.get_record(key)
+        self.assertEqual('record', record['Value'])
+        self.assertEqual(12, record['Flags'])
+        self.assertIsInstance(record, dict)
+
+    @generate_key
+    def test_get_record_fail(self, key):
+        self.assertEqual(self.consul.kv.get_record(key), None)
+
+    @generate_key
+    def test_set_record_no_replace_get_item_str_value(self, key):
+        self.consul.kv.set(key, 'foo')
+        self.consul.kv.set_record(key, 0, 'foo', False)
+        self.assertEqual(self.consul.kv.get(key), 'foo')
+
+    @generate_key
+    def test_set_record_same_value_get_item_str_value(self, key):
+        self.consul.kv.set(key, 'foo')
+        self.consul.kv.set_record(key, 0, 'foo', True)
+        self.assertEqual(self.consul.kv.get(key), 'foo')
+
     @generate_key
     def test_set_item_get_item_dict_value(self, key):
         value = {'foo': 'bar'}
@@ -170,9 +206,14 @@ def test_set_item_get_item_unicode_value(self, key):
     def test_set_item_get_item_unicode_value(self, key):
         self.consul.kv.set(key, '✈')
         response = self.consul.kv.get(key)
-        print(repr(response))
         self.assertEqual(response, '✈')
 
+    @generate_key
+    def test_set_item_in_records(self, key):
+        self.consul.kv.set(key, 'zomg')
+        expectation = (key, 0, 'zomg')
+        self.assertIn(expectation, self.consul.kv.records())
+
 
 class TestSession(unittest.TestCase):