From 22c95472f41ac863fc695748087c03cb8f1fa710 Mon Sep 17 00:00:00 2001
From: brettlangdon <brett@blangdon.com>
Date: Tue, 10 Sep 2013 08:55:24 -0400
Subject: [PATCH] added add_serializer/add_deserializer and
 serialize/deserialize methods also removed setup_serializer and favored
 self.serialize/self.deserialize in code

---
 riakcached/client.py            | 34 ++++++++++++++++++++++-----------
 riakcached/tests/test_client.py | 21 ++++++++++++--------
 2 files changed, 36 insertions(+), 19 deletions(-)

diff --git a/riakcached/client.py b/riakcached/client.py
index 49479e1..69e8b97 100644
--- a/riakcached/client.py
+++ b/riakcached/client.py
@@ -37,13 +37,30 @@ class RiakClient(object):
         if auto_connect:
             self._connect()
 
-    def setup_serializer(self, content_type, serializer, deserializer):
+    def add_serializer(self, content_type, serializer):
         """
         """
         content_type = content_type.lower()
         self._serializers[content_type] = serializer
+
+    def add_deserializer(self, content_type, deserializer):
+        """
+        """
+        content_type = content_type.lower()
         self._deserializers[content_type] = deserializer
 
+    def serialize(self, data, content_type):
+        """
+        """
+        serializer = self._serializers.get(content_type, str)
+        return serializer(data)
+
+    def deserialize(self, data, content_type):
+        """
+        """
+        deserializer = self._deserializers.get(content_type, str)
+        return deserializer(data)
+
     def close(self):
         """
         """
@@ -65,8 +82,7 @@ class RiakClient(object):
         if response.status not in (200, 300, 304):
             return None
 
-        deserializer = self._deserializers.get(response.getheader("content-type"), str)
-        return deserializer(response.data)
+        return self.deserialize(response.data, response.getheader("content-type"))
 
     def get_many(self, keys):
         """
@@ -82,8 +98,7 @@ class RiakClient(object):
     def set(self, key, value, content_type="text/plain"):
         """
         """
-        serializer = self._serializers.get(content_type.lower(), str)
-        value = serializer(value)
+        value = self.serialize(value, content_type)
 
         response = self._request(
             method="POST",
@@ -137,8 +152,7 @@ class RiakClient(object):
             url="%s/stats" % self.url,
         )
         if response.status == 200:
-            deserializer = self._deserializers.get("application/json", json.loads)
-            return deserializer(response.data)
+            return self.deserialize(response.data, "application/json")
         return None
 
     def props(self):
@@ -153,11 +167,10 @@ class RiakClient(object):
         return None
 
     def set_props(self, props):
-        serializer = self._serializers.get("application/json", json.dumps)
         response = self._request(
             method="PUT",
             url="%s/buckets/%s/props" % (self.url, self.bucket),
-            body=serializer(props),
+            body=self.serialize(props, "application/json"),
             headers={
                 "Content-Type": "application/json",
             }
@@ -170,8 +183,7 @@ class RiakClient(object):
             url="%s/buckets/%s/keys?keys=true" % (self.url, self.bucket),
         )
         if response.status == 200:
-            deserializer = self._deserializers.get("application/json", json.loads)
-            return deserializer(response.data)
+            return self.deserialize(response.data, "application/json")
         return None
 
     def ping(self):
diff --git a/riakcached/tests/test_client.py b/riakcached/tests/test_client.py
index e0640f9..8f86cfc 100644
--- a/riakcached/tests/test_client.py
+++ b/riakcached/tests/test_client.py
@@ -21,17 +21,22 @@ class TestRiakClient(unittest2.TestCase):
         client = RiakClient("test_bucket", url="http://127.0.0.1:8098/", auto_connect=False)
         self.assertEqual(client.url, "http://127.0.0.1:8098")
 
-    def test_client_adds_serializer_and_deserializer(self):
+    def test_client_adds_serializer(self):
         serializer = mock.Mock()
-        deserializer = mock.Mock()
 
         client = RiakClient("test_bucket", auto_connect=False)
-        client.setup_serializer("application/test", serializer, deserializer)
+        client.add_serializer("application/test", serializer)
         self.assertEqual(client._serializers["application/test"], serializer)
-        self.assertEqual(client._deserializers["application/test"], deserializer)
-        client._serializers["application/test"]()
-        client._deserializers["application/test"]()
+        client.serialize("test", "application/test")
         serializer.assert_called()
+
+    def test_client_adds_deserializer(self):
+        deserializer = mock.Mock()
+
+        client = RiakClient("test_bucket", auto_connect=False)
+        client.add_deserializer("application/test", deserializer)
+        self.assertEqual(client._deserializers["application/test"], deserializer)
+        client.deserialize("test", "application/test")
         deserializer.assert_called()
 
     def test_client_calls_pool_close(self):
@@ -140,7 +145,7 @@ class TestGet(unittest2.TestCase):
             return "deserialized"
 
         client = RiakClient("test_bucket")
-        client.setup_serializer("application/test", None, deserializer)
+        client.add_deserializer("application/test", deserializer)
         client._pool.urlopen.return_value = InlineClass({
             "status": 200,
             "data": "some data",
@@ -247,7 +252,7 @@ class TestSet(unittest2.TestCase):
             return "serialized"
 
         client = RiakClient("test_bucket")
-        client.setup_serializer("application/test", serializer, None)
+        client.add_serializer("application/test", serializer)
         client._pool.urlopen.return_value = InlineClass({
             "status": 204,
             "data": "",