Added tests for pickle import

neurodamus/utils/memory.py

@@ -299,6 +299,21 @@ def export_allocation_stats(rank_allocation, filename):
         logging.warning("Unable to export allocation stats: {}".format(e))
 
 
+@run_only_rank0
+def import_allocation_stats(filename):
+    """
+    Import allocation dictionary from serialized pickle file.
+    """
+    import pickle
+    try:
+        with open(filename, 'rb') as f:
+            rank_allocation = pickle.load(f)
+        return rank_allocation
+    except Exception as e:
+        logging.warning("Unable to import allocation stats: {}".format(e))
+        return None
+
+
 class SynapseMemoryUsage:
     ''' A small class that works as a lookup table
     for the memory used by each type of synapse.

tests/integration-e2e/test_dry_run_worflow.py

@@ -23,3 +23,13 @@ def test_dry_run_workflow(USECASE3):
     }
     assert nd._dry_run_stats.metype_counts == expected_items
     assert nd._dry_run_stats.suggest_nodes(0.3) > 0
+
+    # Test that the dry run mode works with a pickle file
+    import pickle
+    try:
+        with open((str(USECASE3 / "allocation.bin")), 'rb') as f:
+            rank_allocation = pickle.load(f)
+    except Exception as e:
+        print("Unable to import allocation stats: {}".format(e))
+    expected_items = {'NodeA': {0: [3]}, 'NodeB': {1: [2]}}
+    assert rank_allocation == expected_items