95 lines
3.0 KiB
Python
95 lines
3.0 KiB
Python
import unittest
|
|
from consistent_hash import ConsistentHashRing
|
|
|
|
|
|
class TestConsistentHash(unittest.TestCase):
|
|
def test_empty_ring_returns_none(self):
|
|
r = ConsistentHashRing()
|
|
self.assertIsNone(r.get_node("k"))
|
|
|
|
def test_single_node_always(self):
|
|
r = ConsistentHashRing()
|
|
r.add_node("A")
|
|
for k in ["x", "y", "z", 1, 2]:
|
|
self.assertEqual(r.get_node(k), "A")
|
|
|
|
def test_add_remove(self):
|
|
r = ConsistentHashRing()
|
|
r.add_node("A"); r.add_node("B")
|
|
self.assertEqual(len(r), 2)
|
|
r.remove_node("A")
|
|
self.assertEqual(len(r), 1)
|
|
self.assertNotIn("A", r)
|
|
for k in ["k1", "k2", "k3"]:
|
|
self.assertEqual(r.get_node(k), "B")
|
|
|
|
def test_duplicate_add_idempotent(self):
|
|
r = ConsistentHashRing(num_vnodes=5)
|
|
r.add_node("A"); r.add_node("A")
|
|
self.assertEqual(len(r), 1)
|
|
self.assertEqual(len(r._ring), 5)
|
|
|
|
def test_remove_nonexistent_noop(self):
|
|
r = ConsistentHashRing()
|
|
r.add_node("A")
|
|
r.remove_node("Z")
|
|
self.assertEqual(len(r), 1)
|
|
|
|
def test_determinism(self):
|
|
r1 = ConsistentHashRing(num_vnodes=10)
|
|
r2 = ConsistentHashRing(num_vnodes=10)
|
|
for n in ["n1", "n2", "n3"]:
|
|
r1.add_node(n); r2.add_node(n)
|
|
for k in ["a", "b", "c", "d"]:
|
|
self.assertEqual(r1.get_node(k), r2.get_node(k))
|
|
|
|
def test_minimal_rebalancing(self):
|
|
# adding a node should move only a fraction of keys
|
|
r = ConsistentHashRing(num_vnodes=50)
|
|
for n in ["n1", "n2"]:
|
|
r.add_node(n)
|
|
before = {k: r.get_node(k) for k in [f"k{i}" for i in range(200)]}
|
|
r.add_node("n3")
|
|
moved = sum(1 for k in before if r.get_node(k) != before[k])
|
|
# expect roughly 1/3 moved, definitely not all
|
|
self.assertLess(moved, 200)
|
|
self.assertGreater(moved, 0)
|
|
|
|
def test_get_nodes_replication(self):
|
|
r = ConsistentHashRing(num_vnodes=20)
|
|
for n in ["a", "b", "c"]:
|
|
r.add_node(n)
|
|
ns = r.get_nodes("key", 2)
|
|
self.assertEqual(len(ns), 2)
|
|
self.assertEqual(len(set(ns)), 2)
|
|
# n larger than node count
|
|
ns2 = r.get_nodes("key", 10)
|
|
self.assertEqual(len(ns2), 3)
|
|
self.assertEqual(len(set(ns2)), 3)
|
|
|
|
def test_get_nodes_empty(self):
|
|
r = ConsistentHashRing()
|
|
self.assertEqual(r.get_nodes("k", 3), [])
|
|
|
|
def test_invalid_vnodes(self):
|
|
with self.assertRaises(ValueError):
|
|
ConsistentHashRing(0)
|
|
with self.assertRaises(ValueError):
|
|
ConsistentHashRing(-5)
|
|
|
|
def test_invalid_n(self):
|
|
r = ConsistentHashRing()
|
|
r.add_node("A")
|
|
with self.assertRaises(ValueError):
|
|
r.get_nodes("k", 0)
|
|
|
|
def test_nodes_sorted(self):
|
|
r = ConsistentHashRing()
|
|
for n in ["z", "a", "m"]:
|
|
r.add_node(n)
|
|
self.assertEqual(r.nodes, ["a", "m", "z"])
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|