Tests for routing strategies

Author: technige

neo4j/v1/routing.py

@@ -17,8 +17,7 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-
-
+from abc import abstractmethod
 from sys import maxsize
 from threading import Lock
 from time import clock
@@ -175,11 +174,12 @@ def build(cls, connection_pool, **config):
             return RoundRobinLoadBalancingStrategy()
         else:
             raise ValueError("Unknown load balancing strategy '%s'" % load_balancing_strategy)
-        pass
 
+    @abstractmethod
     def select_reader(self, known_readers):
         raise NotImplementedError()
 
+    @abstractmethod
     def select_writer(self, known_writers):
         raise NotImplementedError()
 
@@ -190,22 +190,20 @@ class RoundRobinLoadBalancingStrategy(LoadBalancingStrategy):
     _writers_offset = 0
 
     def select_reader(self, known_readers):
-        address = self.select(self._readers_offset, known_readers)
+        address = self._select(self._readers_offset, known_readers)
         self._readers_offset += 1
         return address
 
     def select_writer(self, known_writers):
-        address = self.select(self._writers_offset, known_writers)
+        address = self._select(self._writers_offset, known_writers)
         self._writers_offset += 1
         return address
 
-    def select(self, offset, addresses):
-        length = len(addresses)
-        if length == 0:
+    @classmethod
+    def _select(cls, offset, addresses):
+        if not addresses:
             return None
-        else:
-            index = offset % length
-            return addresses.get(index)
+        return addresses[offset % len(addresses)]
 
 
 class LeastConnectedLoadBalancingStrategy(LoadBalancingStrategy):
@@ -216,43 +214,42 @@ def __init__(self, connection_pool):
         self._connection_pool = connection_pool
 
     def select_reader(self, known_readers):
-        address = self.select(self._readers_offset, known_readers)
+        address = self._select(self._readers_offset, known_readers)
         self._readers_offset += 1
         return address
 
     def select_writer(self, known_writers):
-        address = self.select(self._writers_offset, known_writers)
+        address = self._select(self._writers_offset, known_writers)
         self._writers_offset += 1
         return address
 
-    def select(self, offset, addresses):
-        length = len(addresses)
-        if length == 0:
+    def _select(self, offset, addresses):
+        if not addresses:
             return None
-        else:
-            start_index = offset % length
-            index = start_index
+        num_addresses = len(addresses)
+        start_index = offset % num_addresses
+        index = start_index
 
-            least_connected_address = None
-            least_in_use_connections = maxsize
+        least_connected_address = None
+        least_in_use_connections = maxsize
 
-            while True:
-                address = addresses[index]
-                in_use_connections = self._connection_pool.in_use_connection_count(address)
+        while True:
+            address = addresses[index]
+            in_use_connections = self._connection_pool.in_use_connection_count(address)
 
-                if in_use_connections < least_in_use_connections:
-                    least_connected_address = address
-                    least_in_use_connections = in_use_connections
+            if in_use_connections < least_in_use_connections:
+                least_connected_address = address
+                least_in_use_connections = in_use_connections
 
-                if index == length - 1:
-                    index = 0
-                else:
-                    index += 1
+            if index == num_addresses - 1:
+                index = 0
+            else:
+                index += 1
 
-                if index == start_index:
-                    break
+            if index == start_index:
+                break
 
-            return least_connected_address
+        return least_connected_address
 
 
 class RoutingConnectionPool(ConnectionPool):

test/unit/test_routing.py

@@ -17,14 +17,15 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-
+from collections import OrderedDict
 from unittest import TestCase
 
 from neo4j.bolt import ProtocolError
 from neo4j.bolt.connection import connect
-from neo4j.v1.routing import OrderedSet, RoutingTable, RoutingConnectionPool
+from neo4j.v1.routing import OrderedSet, RoutingTable, RoutingConnectionPool, LeastConnectedLoadBalancingStrategy, \
+    RoundRobinLoadBalancingStrategy
 from neo4j.v1.security import basic_auth
-from neo4j.v1.api import Driver, READ_ACCESS, WRITE_ACCESS
+from neo4j.v1.api import READ_ACCESS, WRITE_ACCESS
 
 
 VALID_ROUTING_RECORD = {
@@ -56,7 +57,6 @@ def connector(address):
 
 
 class RoundRobinSetTestCase(TestCase):
-
     def test_should_repr_as_set(self):
         s = OrderedSet([1, 2, 3])
         assert repr(s) == "{1, 2, 3}"
@@ -135,15 +135,13 @@ def test_should_be_able_to_replace(self):
 
 
 class RoutingTableConstructionTestCase(TestCase):
-
     def test_should_be_initially_stale(self):
         table = RoutingTable()
         assert not table.is_fresh(READ_ACCESS)
         assert not table.is_fresh(WRITE_ACCESS)
 
 
 class RoutingTableParseRoutingInfoTestCase(TestCase):
-
     def test_should_return_routing_table_on_valid_record(self):
         table = RoutingTable.parse_routing_info([VALID_ROUTING_RECORD])
         assert table.routers == {('127.0.0.1', 9001), ('127.0.0.1', 9002), ('127.0.0.1', 9003)}
@@ -172,7 +170,6 @@ def test_should_fail_on_multiple_records(self):
 
 
 class RoutingTableFreshnessTestCase(TestCase):
-
     def test_should_be_fresh_after_update(self):
         table = RoutingTable.parse_routing_info([VALID_ROUTING_RECORD])
         assert table.is_fresh(READ_ACCESS)
@@ -198,7 +195,6 @@ def test_should_become_stale_if_no_writers(self):
 
 
 class RoutingTableUpdateTestCase(TestCase):
-
     def setUp(self):
         self.table = RoutingTable(
             [("192.168.1.1", 7687), ("192.168.1.2", 7687)], [("192.168.1.3", 7687)], [], 0)
@@ -224,9 +220,119 @@ def test_update_should_replace_ttl(self):
 
 
 class RoutingConnectionPoolConstructionTestCase(TestCase):
-
     def test_should_populate_initial_router(self):
         initial_router = ("127.0.0.1", 9001)
         router = ("127.0.0.1", 9002)
         with RoutingConnectionPool(connector, initial_router, {}, router) as pool:
             assert pool.routing_table.routers == {("127.0.0.1", 9002)}
+
+
+class FakeConnectionPool(object):
+
+    def __init__(self, addresses):
+        self._addresses = addresses
+
+    def in_use_connection_count(self, address):
+        return self._addresses.get(address, 0)
+
+
+class RoundRobinLoadBalancingStrategyTestCase(TestCase):
+
+    def test_simple_reader_selection(self):
+        strategy = RoundRobinLoadBalancingStrategy()
+        self.assertEqual(strategy.select_reader(["0.0.0.0", "1.1.1.1", "2.2.2.2"]), "0.0.0.0")
+        self.assertEqual(strategy.select_reader(["0.0.0.0", "1.1.1.1", "2.2.2.2"]), "1.1.1.1")
+        self.assertEqual(strategy.select_reader(["0.0.0.0", "1.1.1.1", "2.2.2.2"]), "2.2.2.2")
+        self.assertEqual(strategy.select_reader(["0.0.0.0", "1.1.1.1", "2.2.2.2"]), "0.0.0.0")
+
+    def test_empty_reader_selection(self):
+        strategy = RoundRobinLoadBalancingStrategy()
+        self.assertIsNone(strategy.select_reader([]))
+
+    def test_simple_writer_selection(self):
+        strategy = RoundRobinLoadBalancingStrategy()
+        self.assertEqual(strategy.select_writer(["0.0.0.0", "1.1.1.1", "2.2.2.2"]), "0.0.0.0")
+        self.assertEqual(strategy.select_writer(["0.0.0.0", "1.1.1.1", "2.2.2.2"]), "1.1.1.1")
+        self.assertEqual(strategy.select_writer(["0.0.0.0", "1.1.1.1", "2.2.2.2"]), "2.2.2.2")
+        self.assertEqual(strategy.select_writer(["0.0.0.0", "1.1.1.1", "2.2.2.2"]), "0.0.0.0")
+
+    def test_empty_writer_selection(self):
+        strategy = RoundRobinLoadBalancingStrategy()
+        self.assertIsNone(strategy.select_writer([]))
+
+
+class LeastConnectedLoadBalancingStrategyTestCase(TestCase):
+
+    def test_simple_reader_selection(self):
+        strategy = LeastConnectedLoadBalancingStrategy(FakeConnectionPool(OrderedDict([
+            ("0.0.0.0", 2),
+            ("1.1.1.1", 1),
+            ("2.2.2.2", 0),
+        ])))
+        self.assertEqual(strategy.select_reader(["0.0.0.0", "1.1.1.1", "2.2.2.2"]), "2.2.2.2")
+
+    def test_reader_selection_with_clash(self):
+        strategy = LeastConnectedLoadBalancingStrategy(FakeConnectionPool(OrderedDict([
+            ("0.0.0.0", 0),
+            ("0.0.0.1", 0),
+            ("1.1.1.1", 1),
+        ])))
+        self.assertEqual(strategy.select_reader(["0.0.0.0", "0.0.0.1", "1.1.1.1"]), "0.0.0.0")
+        self.assertEqual(strategy.select_reader(["0.0.0.0", "0.0.0.1", "1.1.1.1"]), "0.0.0.1")
+
+    def test_empty_reader_selection(self):
+        strategy = LeastConnectedLoadBalancingStrategy(FakeConnectionPool(OrderedDict([
+        ])))
+        self.assertIsNone(strategy.select_reader([]))
+
+    def test_not_in_pool_reader_selection(self):
+        strategy = LeastConnectedLoadBalancingStrategy(FakeConnectionPool(OrderedDict([
+            ("1.1.1.1", 1),
+            ("2.2.2.2", 2),
+        ])))
+        self.assertEqual(strategy.select_reader(["2.2.2.2", "3.3.3.3"]), "3.3.3.3")
+
+    def test_partially_in_pool_reader_selection(self):
+        strategy = LeastConnectedLoadBalancingStrategy(FakeConnectionPool(OrderedDict([
+            ("1.1.1.1", 1),
+            ("2.2.2.2", 0),
+        ])))
+        self.assertEqual(strategy.select_reader(["2.2.2.2", "3.3.3.3"]), "2.2.2.2")
+        self.assertEqual(strategy.select_reader(["2.2.2.2", "3.3.3.3"]), "3.3.3.3")
+
+    def test_simple_writer_selection(self):
+        strategy = LeastConnectedLoadBalancingStrategy(FakeConnectionPool(OrderedDict([
+            ("0.0.0.0", 2),
+            ("1.1.1.1", 1),
+            ("2.2.2.2", 0),
+        ])))
+        self.assertEqual(strategy.select_writer(["0.0.0.0", "1.1.1.1", "2.2.2.2"]), "2.2.2.2")
+
+    def test_writer_selection_with_clash(self):
+        strategy = LeastConnectedLoadBalancingStrategy(FakeConnectionPool(OrderedDict([
+            ("0.0.0.0", 0),
+            ("0.0.0.1", 0),
+            ("1.1.1.1", 1),
+        ])))
+        self.assertEqual(strategy.select_writer(["0.0.0.0", "0.0.0.1", "1.1.1.1"]), "0.0.0.0")
+        self.assertEqual(strategy.select_writer(["0.0.0.0", "0.0.0.1", "1.1.1.1"]), "0.0.0.1")
+
+    def test_empty_writer_selection(self):
+        strategy = LeastConnectedLoadBalancingStrategy(FakeConnectionPool(OrderedDict([
+        ])))
+        self.assertIsNone(strategy.select_writer([]))
+
+    def test_not_in_pool_writer_selection(self):
+        strategy = LeastConnectedLoadBalancingStrategy(FakeConnectionPool(OrderedDict([
+            ("1.1.1.1", 1),
+            ("2.2.2.2", 2),
+        ])))
+        self.assertEqual(strategy.select_writer(["2.2.2.2", "3.3.3.3"]), "3.3.3.3")
+
+    def test_partially_in_pool_writer_selection(self):
+        strategy = LeastConnectedLoadBalancingStrategy(FakeConnectionPool(OrderedDict([
+            ("1.1.1.1", 1),
+            ("2.2.2.2", 0),
+        ])))
+        self.assertEqual(strategy.select_writer(["2.2.2.2", "3.3.3.3"]), "2.2.2.2")
+        self.assertEqual(strategy.select_writer(["2.2.2.2", "3.3.3.3"]), "3.3.3.3")