Add getNodeMap to partition table.

git-svn-id: https://svn.erp5.org/repos/neo/trunk@2549 71dcc9de-d417-0410-9af5-da40c76e7ee4

neo/pt.py

@@ -318,6 +318,13 @@ class PartitionTable(object):
         getRow = self.getRow
         return [(x, getRow(x)) for x in xrange(self.np)]
 
+    def getNodeMap(self):
+        """ Return a list of 2-tuple: (uuid, partition_list) """
+        uuid_map = {}
+        for index, row in enumerate(self.partition_list):
+            for cell in row:
+                uuid_map.setdefault(cell.getNode(), []).append(index)
+        return uuid_map
 
 def thread_safe(method):
     def wrapper(self, *args, **kwargs):
@@ -367,3 +374,7 @@ class MTPartitionTable(PartitionTable):
     def getNodeList(self, *args, **kwargs):
         return PartitionTable.getNodeList(self, *args, **kwargs)
 
+    @thread_safe
+    def getNodeMap(self, *args, **kwargs):
+        return PartitionTable.getNodeMap(self, *args, **kwargs)
+

neo/tests/testPT.py

@@ -421,6 +421,35 @@ class PartitionTableTests(NeoUnitTestBase):
         # unknwon row
         self.assertRaises(IndexError,  pt.getRow, 5)
 
+    def test_getNodeMap(self):
+        num_partitions = 5
+        num_replicas = 2
+        pt = PartitionTable(num_partitions, num_replicas)
+        uuid1 = self.getNewUUID()
+        uuid2 = self.getNewUUID()
+        uuid3 = self.getNewUUID()
+        sn1 = StorageNode(Mock(),("127.0.0.1", 19001) , uuid1)
+        pt.setCell(0, sn1, CellStates.UP_TO_DATE)
+        pt.setCell(1, sn1, CellStates.UP_TO_DATE)
+        pt.setCell(2, sn1, CellStates.UP_TO_DATE)
+        self.assertEqual(pt.getNodeMap(), {
+            sn1: [0, 1, 2],
+        })
+        sn2 = StorageNode(Mock(), ("127.0.0.2", 19001), uuid2)
+        pt.setCell(0, sn2, CellStates.UP_TO_DATE)
+        pt.setCell(1, sn2, CellStates.UP_TO_DATE)
+        self.assertEqual(pt.getNodeMap(), {
+            sn1: [0, 1, 2],
+            sn2: [0, 1],
+        })
+        sn3 = StorageNode(Mock(), ("127.0.0.3", 19001), uuid3)
+        pt.setCell(0, sn3, CellStates.UP_TO_DATE)
+        self.assertEqual(pt.getNodeMap(), {
+            sn1: [0, 1, 2],
+            sn2: [0, 1],
+            sn3: [0],
+        })
+
 if __name__ == '__main__':
     unittest.main()