From 5799cf848944207cbf21d7183b75c78f678bbb38 Mon Sep 17 00:00:00 2001
From: Dimitri RODARIE <d.rodarie@gmail.com>
Date: Thu, 12 Dec 2024 18:05:49 +0100
Subject: [PATCH] fix: postprocessing (#901)

* fix: after hooks postprocessing for multi threads

* fix: postprocessing tests

* fix: remove repository cloning to EBRAINS

Only the metadata package (dbbs-lab/bsb) should be deployed on EBRAINS

---------

Co-authored-by: Filippo <52816133+filimarc@users.noreply.github.com>
---
 .github/workflows/main.yml   | 13 ------
 bsb/postprocessing.py        | 16 +++----
 tests/test_postprocessing.py | 87 ++++++++++++++++++++++++++++++++++++
 3 files changed, 95 insertions(+), 21 deletions(-)
 create mode 100644 tests/test_postprocessing.py

diff --git a/.github/workflows/main.yml b/.github/workflows/main.yml
index 2054aaf2..d45934f3 100644
--- a/.github/workflows/main.yml
+++ b/.github/workflows/main.yml
@@ -157,16 +157,3 @@ jobs:
 
       - name: Publish package distributions to PyPI
         uses: pypa/gh-action-pypi-publish@release/v1
-
-  deploy_ebrains:
-    runs-on: ubuntu-latest
-    needs: release
-    environment: EBRAINS
-    steps:
-      - name: Mirror to EBRAINS Gitlab
-        uses: wei/git-sync@v3
-        with:
-          source_repo: dbbs-lab/bsb-core
-          source_branch: main
-          destination_repo: https://push:${{ secrets.EBRAINS_GITLAB_ACCESS_TOKEN }}@gitlab.ebrains.eu/robinde/bsb-core.git
-          destination_branch: main
diff --git a/bsb/postprocessing.py b/bsb/postprocessing.py
index 60b90519..4249f694 100644
--- a/bsb/postprocessing.py
+++ b/bsb/postprocessing.py
@@ -13,10 +13,10 @@ class AfterPlacementHook(abc.ABC):
     name: str = config.attr(key=True)
 
     def queue(self, pool):
-        pool.queue(
-            lambda scaffold: scaffold.after_placement[self.name].postprocess(),
-            submitter=self,
-        )
+        def static_function(scaffold, name):
+            return scaffold.after_placement[name].postprocess()
+
+        pool.queue(static_function, (self.name,), submitter=self)
 
     @abc.abstractmethod
     def postprocess(self):
@@ -28,10 +28,10 @@ class AfterConnectivityHook(abc.ABC):
     name: str = config.attr(key=True)
 
     def queue(self, pool):
-        pool.queue(
-            lambda scaffold: scaffold.after_connectivity[self.name].postprocess(),
-            submitter=self,
-        )
+        def static_function(scaffold, name):
+            return scaffold.after_connectivity[name].postprocess()
+
+        pool.queue(static_function, (self.name,), submitter=self)
 
     @abc.abstractmethod
     def postprocess(self):
diff --git a/tests/test_postprocessing.py b/tests/test_postprocessing.py
new file mode 100644
index 00000000..c3504dc6
--- /dev/null
+++ b/tests/test_postprocessing.py
@@ -0,0 +1,87 @@
+import os
+import unittest
+
+from bsb_test import RandomStorageFixture
+
+from bsb import (
+    MPI,
+    AfterConnectivityHook,
+    AfterPlacementHook,
+    Configuration,
+    Scaffold,
+    config,
+)
+
+
+class TestAfterConnectivityHook(
+    RandomStorageFixture, unittest.TestCase, engine_name="hdf5"
+):
+    def setUp(self):
+        super().setUp()
+
+        @config.node
+        class TestAfterConn(AfterConnectivityHook):
+            def postprocess(self):
+                with open(f"test_after_conn_{MPI.get_rank()}.txt", "a") as f:
+                    # make sure we have access to the scaffold context
+                    f.write(f"{self.scaffold.configuration.name}\n")
+
+        self.network = Scaffold(
+            config=Configuration.default(
+                name="Test config",
+                after_connectivity={"test_after_conn": TestAfterConn()},
+            ),
+            storage=self.storage,
+        )
+
+    def test_after_connectivity_job(self):
+        self.network.compile()
+        if MPI.get_rank() == 0:
+            count_files = 0
+            for filename in os.listdir():
+                if filename.startswith("test_after_conn_"):
+                    count_files += 1
+                    with open(filename, "r") as f:
+                        lines = f.readlines()
+                        self.assertEqual(
+                            len(lines), 1, "The postprocess should be called only once."
+                        )
+                        self.assertEqual(lines[0], "Test config\n")
+                    os.remove(filename)
+            self.assertEqual(count_files, 1)
+
+
+class TestAfterPlacementHook(RandomStorageFixture, unittest.TestCase, engine_name="hdf5"):
+    def setUp(self):
+        super().setUp()
+
+        @config.node
+        class TestAfterPlace(AfterPlacementHook):
+            def postprocess(self):
+                with open(f"test_after_place_{MPI.get_rank()}.txt", "a") as f:
+                    # make sure we have access to the scaffold context
+                    f.write(f"{self.scaffold.configuration.name}\n")
+
+        self.network = Scaffold(
+            config=Configuration.default(
+                name="Test config",
+                after_placement={"test_after_placement": TestAfterPlace()},
+            ),
+            storage=self.storage,
+        )
+
+    def test_after_placement_job(self):
+        self.network.compile()
+        if MPI.get_rank() == 0:
+            count_files = 0
+            for filename in os.listdir():
+                if filename.startswith("test_after_place_"):
+                    count_files += 1
+                    with open(filename, "r") as f:
+                        lines = f.readlines()
+                        self.assertEqual(
+                            len(lines), 1, "The postprocess should be called only once."
+                        )
+                        self.assertEqual(lines[0], "Test config\n")
+                    os.remove(filename)
+            self.assertEqual(count_files, 1)