RoundRobinSet->OrderedSet plus tests

Author: technige

neo4j/v1/routing.py

@@ -39,7 +39,7 @@
 LOAD_BALANCING_STRATEGY_DEFAULT = LOAD_BALANCING_STRATEGY_LEAST_CONNECTED
 
 
-class RoundRobinSet(MutableSet):
+class OrderedSet(MutableSet):
 
     def __init__(self, elements=()):
         self._elements = OrderedDict.fromkeys(elements)
@@ -51,22 +51,15 @@ def __repr__(self):
     def __contains__(self, element):
         return element in self._elements
 
-    def __next__(self):
-        current = None
-        if self._elements:
-            if self._current is None:
-                self._current = 0
-            else:
-                self._current = (self._current + 1) % len(self._elements)
-            current = self.get(self._current)
-        return current
-
     def __iter__(self):
         return iter(self._elements)
 
     def __len__(self):
         return len(self._elements)
 
+    def __getitem__(self, index):
+        return list(self._elements.keys())[index]
+
     def add(self, element):
         self._elements[element] = None
 
@@ -79,9 +72,6 @@ def discard(self, element):
         except KeyError:
             pass
 
-    def next(self):
-        return self.__next__()
-
     def remove(self, element):
         try:
             del self._elements[element]
@@ -96,9 +86,6 @@ def replace(self, elements=()):
         e.clear()
         e.update(OrderedDict.fromkeys(elements))
 
-    def get(self, index):
-        return list(self._elements.keys())[index]
-
 
 class RoutingTable(object):
 
@@ -135,9 +122,9 @@ def parse_routing_info(cls, records):
             return cls(routers, readers, writers, ttl)
 
     def __init__(self, routers=(), readers=(), writers=(), ttl=0):
-        self.routers = RoundRobinSet(routers)
-        self.readers = RoundRobinSet(readers)
-        self.writers = RoundRobinSet(writers)
+        self.routers = OrderedSet(routers)
+        self.readers = OrderedSet(readers)
+        self.writers = OrderedSet(writers)
         self.last_updated_time = self.timer()
         self.ttl = ttl
 
@@ -250,7 +237,7 @@ def select(self, offset, addresses):
             least_in_use_connections = maxsize
 
             while True:
-                address = addresses.get(index)
+                address = addresses[index]
                 in_use_connections = self._connection_pool.in_use_connection_count(address)
 
                 if in_use_connections < least_in_use_connections:

test/unit/test_routing.py

@@ -22,7 +22,7 @@
 
 from neo4j.bolt import ProtocolError
 from neo4j.bolt.connection import connect
-from neo4j.v1.routing import RoundRobinSet, RoutingTable, RoutingConnectionPool
+from neo4j.v1.routing import OrderedSet, RoutingTable, RoutingConnectionPool
 from neo4j.v1.security import basic_auth
 from neo4j.v1.api import Driver, READ_ACCESS, WRITE_ACCESS
 
@@ -58,87 +58,80 @@ def connector(address):
 class RoundRobinSetTestCase(TestCase):
 
     def test_should_repr_as_set(self):
-        rrs = RoundRobinSet([1, 2, 3])
-        assert repr(rrs) == "{1, 2, 3}"
+        s = OrderedSet([1, 2, 3])
+        assert repr(s) == "{1, 2, 3}"
 
     def test_should_contain_element(self):
-        rrs = RoundRobinSet([1, 2, 3])
-        assert 2 in rrs
+        s = OrderedSet([1, 2, 3])
+        assert 2 in s
 
     def test_should_not_contain_non_element(self):
-        rrs = RoundRobinSet([1, 2, 3])
-        assert 4 not in rrs
-
-    def test_should_be_able_to_get_next_if_empty(self):
-        rrs = RoundRobinSet([])
-        assert next(rrs) is None
-
-    def test_should_be_able_to_get_next_repeatedly(self):
-        rrs = RoundRobinSet([1, 2, 3])
-        assert next(rrs) == 1
-        assert next(rrs) == 2
-        assert next(rrs) == 3
-        assert next(rrs) == 1
-
-    def test_should_be_able_to_get_next_repeatedly_via_old_method(self):
-        rrs = RoundRobinSet([1, 2, 3])
-        assert rrs.next() == 1
-        assert rrs.next() == 2
-        assert rrs.next() == 3
-        assert rrs.next() == 1
+        s = OrderedSet([1, 2, 3])
+        assert 4 not in s
+
+    def test_should_be_able_to_get_item_if_empty(self):
+        s = OrderedSet([])
+        with self.assertRaises(IndexError):
+            _ = s[0]
+
+    def test_should_be_able_to_get_items_by_index(self):
+        s = OrderedSet([1, 2, 3])
+        self.assertEqual(s[0], 1)
+        self.assertEqual(s[1], 2)
+        self.assertEqual(s[2], 3)
 
     def test_should_be_iterable(self):
-        rrs = RoundRobinSet([1, 2, 3])
-        assert list(iter(rrs)) == [1, 2, 3]
+        s = OrderedSet([1, 2, 3])
+        assert list(iter(s)) == [1, 2, 3]
 
     def test_should_have_length(self):
-        rrs = RoundRobinSet([1, 2, 3])
-        assert len(rrs) == 3
+        s = OrderedSet([1, 2, 3])
+        assert len(s) == 3
 
     def test_should_be_able_to_add_new(self):
-        rrs = RoundRobinSet([1, 2, 3])
-        rrs.add(4)
-        assert list(rrs) == [1, 2, 3, 4]
+        s = OrderedSet([1, 2, 3])
+        s.add(4)
+        assert list(s) == [1, 2, 3, 4]
 
     def test_should_be_able_to_add_existing(self):
-        rrs = RoundRobinSet([1, 2, 3])
-        rrs.add(2)
-        assert list(rrs) == [1, 2, 3]
+        s = OrderedSet([1, 2, 3])
+        s.add(2)
+        assert list(s) == [1, 2, 3]
 
     def test_should_be_able_to_clear(self):
-        rrs = RoundRobinSet([1, 2, 3])
-        rrs.clear()
-        assert list(rrs) == []
+        s = OrderedSet([1, 2, 3])
+        s.clear()
+        assert list(s) == []
 
     def test_should_be_able_to_discard_existing(self):
-        rrs = RoundRobinSet([1, 2, 3])
-        rrs.discard(2)
-        assert list(rrs) == [1, 3]
+        s = OrderedSet([1, 2, 3])
+        s.discard(2)
+        assert list(s) == [1, 3]
 
     def test_should_be_able_to_discard_non_existing(self):
-        rrs = RoundRobinSet([1, 2, 3])
-        rrs.discard(4)
-        assert list(rrs) == [1, 2, 3]
+        s = OrderedSet([1, 2, 3])
+        s.discard(4)
+        assert list(s) == [1, 2, 3]
 
     def test_should_be_able_to_remove_existing(self):
-        rrs = RoundRobinSet([1, 2, 3])
-        rrs.remove(2)
-        assert list(rrs) == [1, 3]
+        s = OrderedSet([1, 2, 3])
+        s.remove(2)
+        assert list(s) == [1, 3]
 
     def test_should_not_be_able_to_remove_non_existing(self):
-        rrs = RoundRobinSet([1, 2, 3])
+        s = OrderedSet([1, 2, 3])
         with self.assertRaises(ValueError):
-            rrs.remove(4)
+            s.remove(4)
 
     def test_should_be_able_to_update(self):
-        rrs = RoundRobinSet([1, 2, 3])
-        rrs.update([3, 4, 5])
-        assert list(rrs) == [1, 2, 3, 4, 5]
+        s = OrderedSet([1, 2, 3])
+        s.update([3, 4, 5])
+        assert list(s) == [1, 2, 3, 4, 5]
 
     def test_should_be_able_to_replace(self):
-        rrs = RoundRobinSet([1, 2, 3])
-        rrs.replace([3, 4, 5])
-        assert list(rrs) == [3, 4, 5]
+        s = OrderedSet([1, 2, 3])
+        s.replace([3, 4, 5])
+        assert list(s) == [3, 4, 5]
 
 
 class RoutingTableConstructionTestCase(TestCase):