Factorise acceptIdentification checks.

neo/client/handlers/master.py

@@ -28,17 +28,13 @@ class PrimaryBootstrapHandler(AnswerBaseHandler):
         app = self.app
         app.trying_master_node = None
 
-    def acceptIdentification(self, conn, node_type, uuid, num_partitions,
+    def _acceptIdentification(self, node, uuid, num_partitions,
             num_replicas, your_uuid, primary_uuid, known_master_list):
         app = self.app
-        # this must be a master node
-        if node_type != NodeTypes.MASTER:
-            conn.close()
-            return
 
         # Register new master nodes.
         found = False
-        conn_address = conn.getAddress()
+        conn_address = node.getAddress()
         for node_address, node_uuid in known_master_list:
             if node_address == conn_address:
                 assert uuid == node_uuid, (dump(uuid), dump(node_uuid))
@@ -48,8 +44,9 @@ class PrimaryBootstrapHandler(AnswerBaseHandler):
                 n = app.nm.createMaster(address=node_address)
             if node_uuid is not None and n.getUUID() != node_uuid:
                 n.setUUID(node_uuid)
-        assert found, (conn, dump(uuid), known_master_list)
+        assert found, (node, dump(uuid), known_master_list)
 
+        conn = node.getConnection()
         if primary_uuid is not None:
             primary_node = app.nm.getByUUID(primary_uuid)
             if primary_node is None:

neo/client/handlers/storage.py

@@ -47,20 +47,12 @@ class StorageBootstrapHandler(AnswerBaseHandler):
     def notReady(self, conn, message):
         raise NodeNotReady(message)
 
-    def acceptIdentification(self, conn, node_type,
+    def _acceptIdentification(self, node,
            uuid, num_partitions, num_replicas, your_uuid, primary_uuid,
            master_list):
         assert primary_uuid == self.app.primary_master_node.getUUID(), (
             dump(primary_uuid), dump(self.app.primary_master_node.getUUID()))
-        # this must be a storage node
-        if node_type != NodeTypes.STORAGE:
-            conn.close()
-            return
-
-        node = self.app.nm.getByAddress(conn.getAddress())
-        assert node is not None, conn.getAddress()
         node.setUUID(uuid)
-        assert node.getConnection() is conn, (node.getConnection(), conn)
 
 class StorageAnswersHandler(AnswerBaseHandler):
     """ Handle all messages related to ZODB operations """

neo/lib/bootstrap.py

@@ -89,16 +89,16 @@ class BootstrapManager(EventHandler):
         """
         conn.close()
 
-    def acceptIdentification(self, conn, node_type, uuid, num_partitions,
+    def _acceptIdentification(self, node, uuid, num_partitions,
             num_replicas, your_uuid, primary_uuid, known_master_list):
         nm  = self.app.nm
 
         # Register new master nodes.
         for address, uuid in known_master_list:
-            node = nm.getByAddress(address)
-            if node is None:
-                node = nm.createMaster(address=address)
-            node.setUUID(uuid)
+            master_node = nm.getByAddress(address)
+            if master_node is None:
+                master_node = nm.createMaster(address=address)
+            master_node.setUUID(uuid)
 
         self.primary = nm.getByUUID(primary_uuid)
         if self.primary is None or self.current is not self.primary:
@@ -106,7 +106,7 @@ class BootstrapManager(EventHandler):
             # - something goes wrong (unknown UUID)
             # - this master doesn't know who's the primary
             # - got the primary's uuid, so cut here
-            conn.close()
+            node.getConnection().close()
             return
 
         neo.lib.logging.info('connected to a primary master node')

neo/master/handlers/election.py

@@ -79,15 +79,9 @@ class ClientElectionHandler(BaseElectionHandler):
         self.app.unconnected_master_node_set.add(addr)
         self.app.negotiating_master_node_set.discard(addr)
 
-    def acceptIdentification(self, conn, node_type, peer_uuid, num_partitions,
+    def _acceptIdentification(self, node, peer_uuid, num_partitions,
             num_replicas, your_uuid, primary_uuid, known_master_list):
         app = self.app
-        if node_type != NodeTypes.MASTER:
-            # The peer is not a master node!
-            neo.lib.logging.error('%r is not a master node', conn)
-            app.nm.remove(app.nm.getByAddress(conn.getAddress()))
-            conn.close()
-            return
 
         if your_uuid != app.uuid:
             # uuid conflict happened, accept the new one and restart election
@@ -96,7 +90,7 @@ class ClientElectionHandler(BaseElectionHandler):
                             dump(your_uuid))
             raise ElectionFailure, 'new uuid supplied'
 
-        conn.setUUID(peer_uuid)
+        node.setUUID(peer_uuid)
 
         # Register new master nodes.
         for address, uuid in known_master_list:
@@ -135,7 +129,7 @@ class ClientElectionHandler(BaseElectionHandler):
                 app.negotiating_master_node_set.clear()
                 return
 
-        elect(app, peer_uuid, conn.getAddress())
+        elect(app, peer_uuid, node.getAddress())
 
 
 class ServerElectionHandler(BaseElectionHandler, MasterHandler):

neo/master/handlers/secondary.py

@@ -84,18 +84,15 @@ class PrimaryHandler(EventHandler):
                     if n.getUUID() is None:
                         n.setUUID(uuid)
 
-    def acceptIdentification(self, conn, node_type, uuid, num_partitions,
+    def _acceptIdentification(self, node, uuid, num_partitions,
             num_replicas, your_uuid, primary_uuid, known_master_list):
         app = self.app
         if primary_uuid != app.primary_master_node.getUUID():
             raise PrimaryFailure('unexpected primary uuid')
-        node = app.nm.getByAddress(conn.getAddress())
-        assert node_type == NodeTypes.MASTER
 
         if your_uuid != app.uuid:
             # uuid conflict happened, accept the new one
             app.uuid = your_uuid
 
-        conn.setUUID(uuid)
         node.setUUID(uuid)
 

neo/tests/__init__.py

@@ -333,13 +333,15 @@ class NeoUnitTestBase(NeoTestBase):
         """ ensure no UUID was set on the connection """
         self.assertEqual(len(conn.mockGetNamedCalls('setUUID')), 0)
 
-    def checkUUIDSet(self, conn, uuid=None):
-        """ ensure no UUID was set on the connection """
+    def checkUUIDSet(self, conn, uuid=None, check_intermediate=True):
+        """ ensure UUID was set on the connection """
         calls = conn.mockGetNamedCalls('setUUID')
-        self.assertEqual(len(calls), 1)
-        call = calls.pop()
+        found_uuid = calls.pop().getParam(0)
+        if check_intermediate:
+            for call in calls:
+                self.assertEqual(found_uuid, call.getParam(0))
         if uuid is not None:
-            self.assertEqual(call.getParam(0), uuid)
+            self.assertEqual(found_uuid, uuid)
 
     # in check(Ask|Answer|Notify)Packet we return the packet so it can be used
     # in tests if more accurates checks are required

neo/tests/client/testMasterHandler.py

@@ -32,6 +32,16 @@ class MasterHandlerTests(NeoUnitTestBase):
         self.app = Mock({'getDB': self.db})
         self.app.nm = NodeManager()
         self.app.dispatcher = Mock()
+        self._next_port = 3000
+
+    def getKnownMaster(self):
+        node = self.app.nm.createMaster(address=(
+            self.local_ip, self._next_port),
+        )
+        self._next_port += 1
+        conn = self.getFakeConnection(address=node.getAddress())
+        node.setConnection(conn)
+        return node, conn
 
 class MasterBootstrapHandlerTests(MasterHandlerTests):
 
@@ -51,18 +61,15 @@ class MasterBootstrapHandlerTests(MasterHandlerTests):
 
     def test_acceptIdentification1(self):
         """ Non-master node """
-        conn = self.getFakeConnection()
-        uuid = self.getNewUUID()
+        node, conn = self.getKnownMaster()
         self.handler.acceptIdentification(conn, NodeTypes.CLIENT,
-            uuid, 100, 0, None, None, [])
+            node.getUUID(), 100, 0, None, None, [])
         self.checkClosed(conn)
 
     def test_acceptIdentification2(self):
         """ No UUID supplied """
-        conn = self.getFakeConnection()
+        node, conn = self.getKnownMaster()
         uuid = self.getNewUUID()
-        node = Mock()
-        self.app.nm = Mock({'getByAddress': node, 'getByUUID': node})
         self.checkProtocolErrorRaised(self.handler.acceptIdentification,
             conn, NodeTypes.MASTER, uuid, 100, 0, None,
             uuid, [(conn.getAddress(), uuid)],
@@ -70,17 +77,13 @@ class MasterBootstrapHandlerTests(MasterHandlerTests):
 
     def test_acceptIdentification3(self):
         """ identification accepted  """
-        node = Mock()
-        conn = self.getFakeConnection()
+        node, conn = self.getKnownMaster()
         uuid = self.getNewUUID()
         your_uuid = self.getNewUUID()
-        partitions = 100
-        replicas = 2
-        self.app.nm = Mock({'getByAddress': node, 'getByUUID': node})
         self.handler.acceptIdentification(conn, NodeTypes.MASTER, uuid,
-            partitions, replicas, your_uuid, uuid, [(conn.getAddress(), uuid)])
+            100, 2, your_uuid, uuid, [(conn.getAddress(), uuid)])
         self.assertEqual(self.app.uuid, your_uuid)
-        self.checkUUIDSet(node, uuid)
+        self.assertEqual(node.getUUID(), uuid)
         self.assertTrue(isinstance(self.app.pt, PartitionTable))
 
     def _getMasterList(self, uuid_list):

neo/tests/client/testStorageHandler.py

@@ -24,6 +24,7 @@ from neo.client.exception import NEOStorageError, NEOStorageNotFoundError
 from neo.client.exception import NEOStorageDoesNotExistError
 from ZODB.POSException import ConflictError
 from neo.lib.exception import NodeNotReady
+from neo.lib.node import NodeManager
 from ZODB.TimeStamp import TimeStamp
 
 MARKER = []
@@ -33,7 +34,24 @@ class StorageBootstrapHandlerTests(NeoUnitTestBase):
     def setUp(self):
         super(NeoUnitTestBase, self).setUp()
         self.app = Mock()
+        self.app.nm = NodeManager()
         self.handler = StorageBootstrapHandler(self.app)
+        self.app.primary_master_node = node = Mock({
+            'getConnection': self.getFakeConnection(),
+            'getUUID': self.getNewUUID(),
+        })
+        self._next_port = 3000
+
+    def getKnownStorage(self):
+        node = self.app.nm.createStorage(
+            uuid=self.getNewUUID(),
+            address=(self.local_ip, self._next_port),
+        )
+        self._next_port += 1
+        conn = self.getFakeConnection(address=node.getAddress(),
+            uuid=node.getUUID())
+        node.setConnection(conn)
+        return node, conn
 
     def test_notReady(self):
         conn = self.getFakeConnection()
@@ -41,26 +59,18 @@ class StorageBootstrapHandlerTests(NeoUnitTestBase):
 
     def test_acceptIdentification1(self):
         """ Not a storage node """
-        uuid = self.getNewUUID()
-        node_uuid = self.getNewUUID()
-        conn = self.getFakeConnection()
-        self.app.primary_master_node = node = Mock({'getUUID': node_uuid})
-        self.app.nm = Mock({'getByAddress': node})
-        self.handler.acceptIdentification(conn, NodeTypes.CLIENT, uuid,
-            10, 0, None, node_uuid, [])
+        node, conn = self.getKnownStorage()
+        self.handler.acceptIdentification(conn, NodeTypes.CLIENT,
+            node.getUUID(),
+            10, 0, None, self.app.primary_master_node.getUUID(), [])
         self.checkClosed(conn)
 
     def test_acceptIdentification2(self):
-        uuid = self.getNewUUID()
-        node_uuid = self.getNewUUID()
-        conn = self.getFakeConnection()
-        self.app.primary_master_node = node = Mock({'getConnection': conn,
-            'getUUID': node_uuid})
-        self.app.nm = Mock({'getByAddress': node})
-        self.handler.acceptIdentification(conn, NodeTypes.STORAGE, uuid,
-            10, 0, None, node_uuid, [])
-        self.checkUUIDSet(node, uuid)
-
+        node, conn = self.getKnownStorage()
+        self.handler.acceptIdentification(conn, NodeTypes.STORAGE,
+            node.getUUID(),
+            10, 0, None, self.app.primary_master_node.getUUID(), [])
+        self.checkNotClosed(conn)
 
 class StorageAnswerHandlerTests(NeoUnitTestBase):
 

neo/tests/master/testElectionHandler.py

@@ -31,7 +31,26 @@ def _addPacket(self, packet):
     if self.connector is not None:
         self.connector._addPacket(packet)
 
-class MasterClientElectionTests(NeoUnitTestBase):
+class MasterClientElectionTestBase(NeoUnitTestBase):
+
+    def setUp(self):
+        super(MasterClientElectionTestBase, self).setUp()
+        self._master_port = 3001
+
+    def identifyToMasterNode(self, uuid=True):
+        if uuid is True:
+            uuid = self.getNewUUID()
+        node = self.app.nm.createMaster(uuid=uuid)
+        node.setAddress((self.local_ip, self._master_port))
+        self._master_port += 1
+        conn = self.getFakeConnection(
+                uuid=node.getUUID(),
+                address=node.getAddress(),
+        )
+        node.setConnection(conn)
+        return (node, conn)
+
+class MasterClientElectionTests(MasterClientElectionTestBase):
 
     def setUp(self):
         NeoUnitTestBase.setUp(self)
@@ -49,19 +68,13 @@ class MasterClientElectionTests(NeoUnitTestBase):
         # apply monkey patches
         self._addPacket = ClientConnection._addPacket
         ClientConnection._addPacket = _addPacket
+        super(MasterClientElectionTests, self).setUp()
 
     def tearDown(self):
         # restore patched methods
         ClientConnection._addPacket = self._addPacket
         NeoUnitTestBase.tearDown(self)
 
-    def identifyToMasterNode(self):
-        node = self.app.nm.getMasterList()[0]
-        node.setUUID(self.getNewUUID())
-        conn = self.getFakeConnection(uuid=node.getUUID(),
-                address=node.getAddress())
-        return (node, conn)
-
     def _checkUnconnected(self, node):
         addr = node.getAddress()
         self.assertFalse(addr in self.app.negotiating_master_node_set)
@@ -107,9 +120,8 @@ class MasterClientElectionTests(NeoUnitTestBase):
         self.checkClosed(conn)
 
     def test_acceptIdentificationDoesNotKnowPrimary(self):
-        master1_uuid = self.getNewUUID()
-        master1_address = ('127.0.0.1', 2001)
-        master1_conn = self.getFakeConnection(address=master1_address)
+        master1, master1_conn = self.identifyToMasterNode()
+        master1_uuid = master1.getUUID()
         self.election.acceptIdentification(
             master1_conn,
             NodeTypes.MASTER,
@@ -118,14 +130,13 @@ class MasterClientElectionTests(NeoUnitTestBase):
             0,
             self.app.uuid,
             None,
-            [(master1_address, master1_uuid)],
+            [(master1.getAddress(), master1_uuid)],
         )
         self.assertEqual(self.app.primary_master_node, None)
 
     def test_acceptIdentificationKnowsPrimary(self):
-        master1_uuid = self.getNewUUID()
-        master1_address = ('127.0.0.1', 2001)
-        master1_conn = self.getFakeConnection(address=master1_address)
+        master1, master1_conn = self.identifyToMasterNode()
+        master1_uuid = master1.getUUID()
         self.election.acceptIdentification(
             master1_conn,
             NodeTypes.MASTER,
@@ -134,19 +145,20 @@ class MasterClientElectionTests(NeoUnitTestBase):
             0,
             self.app.uuid,
             master1_uuid,
-            [(master1_address, master1_uuid)],
+            [(master1.getAddress(), master1_uuid)],
         )
         self.assertNotEqual(self.app.primary_master_node, None)
 
     def test_acceptIdentificationMultiplePrimaries(self):
-        master1_uuid = self.getNewUUID()
-        master2_uuid = self.getNewUUID()
-        master3_uuid = self.getNewUUID()
-        master1_address = ('127.0.0.1', 2001)
-        master2_address = ('127.0.0.1', 2002)
-        master3_address = ('127.0.0.1', 2003)
-        master1_conn = self.getFakeConnection(address=master1_address)
-        master2_conn = self.getFakeConnection(address=master2_address)
+        master1, master1_conn = self.identifyToMasterNode()
+        master2, master2_conn = self.identifyToMasterNode()
+        master3, _ = self.identifyToMasterNode()
+        master1_uuid = master1.getUUID()
+        master2_uuid = master2.getUUID()
+        master3_uuid = master3.getUUID()
+        master1_address = master1.getAddress()
+        master2_address = master2.getAddress()
+        master3_address = master3.getAddress()
         self.election.acceptIdentification(
             master1_conn,
             NodeTypes.MASTER,
@@ -197,7 +209,7 @@ class MasterClientElectionTests(NeoUnitTestBase):
         return [(x.getAddress(), x.getUUID()) for x in master_list]
 
 
-class MasterServerElectionTests(NeoUnitTestBase):
+class MasterServerElectionTests(MasterClientElectionTestBase):
 
     def setUp(self):
         NeoUnitTestBase.setUp(self)
@@ -219,26 +231,13 @@ class MasterServerElectionTests(NeoUnitTestBase):
         # apply monkey patches
         self._addPacket = ClientConnection._addPacket
         ClientConnection._addPacket = _addPacket
+        super(MasterServerElectionTests, self).setUp()
 
     def tearDown(self):
         NeoUnitTestBase.tearDown(self)
         # restore environnement
         ClientConnection._addPacket = self._addPacket
 
-    def identifyToMasterNode(self, uuid=True):
-        node = self.app.nm.getMasterList()[0]
-        if uuid is True:
-            uuid = self.getNewUUID()
-        node.setUUID(uuid)
-        conn = self.getFakeConnection(
-                uuid=node.getUUID(),
-                address=node.getAddress(),
-        )
-        return (node, conn)
-
-
-    # Tests
-
     def test_requestIdentification1(self):
         """ A non-master node request identification """
         node, conn = self.identifyToMasterNode()
@@ -275,7 +274,7 @@ class MasterServerElectionTests(NeoUnitTestBase):
         args = (self.app.uuid, node.getAddress(), self.app.name)
         self.election.requestIdentification(conn,
             NodeTypes.MASTER, *args)
-        self.checkUUIDSet(conn)
+        self.checkUUIDSet(conn, check_intermediate=False)
         args = self.checkAcceptIdentification(conn, decode=True)
         (node_type, uuid, partitions, replicas, new_uuid, primary_uuid,
             master_list) = args