From db3593649345b9e83160dc566a8a83206fea6eee Mon Sep 17 00:00:00 2001
From: Julien Danjou <julien@danjou.info>
Date: Wed, 8 Mar 2017 16:49:41 +0100
Subject: [PATCH] hashring: allow to use bytes as node name

Change-Id: I8d7c59086c191275ea0ee4b9b1be4822dae0be12
---
 tooz/hashring.py            |  5 +++--
 tooz/tests/test_hashring.py | 20 ++++++++++++++++++++
 2 files changed, 23 insertions(+), 2 deletions(-)

diff --git a/tooz/hashring.py b/tooz/hashring.py
index 94c93dcc..0bc1970b 100644
--- a/tooz/hashring.py
+++ b/tooz/hashring.py
@@ -19,6 +19,7 @@ import hashlib
 import six
 
 import tooz
+from tooz import utils
 
 
 class UnknownNode(tooz.ToozError):
@@ -68,7 +69,7 @@ class HashRing(object):
         will each handle 1/6, 1/3 and 1/2 of the resources, respectively.
         """
         for node in nodes:
-            key = six.text_type(node).encode('utf-8')
+            key = utils.to_binary(node, 'utf-8')
             key_hash = hashlib.md5(key)
             for r in six.moves.range(self._partition_number * weight):
                 key_hash.update(key)
@@ -90,7 +91,7 @@ class HashRing(object):
         except KeyError:
             raise UnknownNode(node)
 
-        key = six.text_type(node).encode('utf-8')
+        key = utils.to_binary(node, 'utf-8')
         key_hash = hashlib.md5(key)
         for r in six.moves.range(self._partition_number * weight):
             key_hash.update(key)
diff --git a/tooz/tests/test_hashring.py b/tooz/tests/test_hashring.py
index 5e7dba72..a5a9cbb2 100644
--- a/tooz/tests/test_hashring.py
+++ b/tooz/tests/test_hashring.py
@@ -61,6 +61,16 @@ class HashRingTestCase(testcase.TestCase):
         self.assertEqual(nodes, set(ring.nodes.keys()))
         self.assertEqual(2 ** 5 * len(nodes), len(ring))
 
+    def test_add_node_bytes(self):
+        nodes = {'foo', 'bar'}
+        ring = hashring.HashRing(nodes)
+        self.assertEqual(nodes, set(ring.nodes.keys()))
+        self.assertEqual(2 ** 5 * len(nodes), len(ring))
+        nodes.add(b'Z\xe2\xfa\x90\x17EC\xac\xae\x88\xa7[\xa1}:E')
+        ring.add_node(b'Z\xe2\xfa\x90\x17EC\xac\xae\x88\xa7[\xa1}:E')
+        self.assertEqual(nodes, set(ring.nodes.keys()))
+        self.assertEqual(2 ** 5 * len(nodes), len(ring))
+
     def test_add_node_unicode(self):
         nodes = {'foo', 'bar'}
         ring = hashring.HashRing(nodes)
@@ -91,6 +101,16 @@ class HashRingTestCase(testcase.TestCase):
         self.assertEqual(nodes, set(ring.nodes.keys()))
         self.assertEqual(2 ** 5 * len(nodes), len(ring))
 
+    def test_remove_node_bytes(self):
+        nodes = {'foo', b'Z\xe2\xfa\x90\x17EC\xac\xae\x88\xa7[\xa1}:E'}
+        ring = hashring.HashRing(nodes)
+        self.assertEqual(nodes, set(ring.nodes.keys()))
+        self.assertEqual(2 ** 5 * len(nodes), len(ring))
+        nodes.discard(b'Z\xe2\xfa\x90\x17EC\xac\xae\x88\xa7[\xa1}:E')
+        ring.remove_node(b'Z\xe2\xfa\x90\x17EC\xac\xae\x88\xa7[\xa1}:E')
+        self.assertEqual(nodes, set(ring.nodes.keys()))
+        self.assertEqual(2 ** 5 * len(nodes), len(ring))
+
     def test_remove_node_unknown(self):
         nodes = ['foo', 'bar']
         ring = hashring.HashRing(nodes)