From 51fe7f096a735ab0833d4981795b9eb39b7e97f9 Mon Sep 17 00:00:00 2001
From: Sebastiaan Huber <mail@sphuber.net>
Date: Wed, 14 Sep 2022 11:00:26 +0200
Subject: [PATCH] Tests: Only reset database connection at end of suite

The `aiida_profile_clean` fixture is a function scoped fixture that
provides a fully configured profile with a storage backend that is
empty, with the exception of a single default user instance.

To empty the database before the start of the test, it calls the method
`clear_profile` on the `TestManager` which calls through to the
`ProfileManager` which calls `reset_profile` on the `StorageBackend`.
However, in addition to resetting the database, it also resets all
database connections in the session, which is not really necessary. In
addition, it can result in too many clients to the database being active
if the session is not cleaned up properly.

We remove the `reset_profile` call from `clear_profile` and instead move
it to the `destroy_all` call. This ensures that the database connections
are not completely reset in between tests anymore, however, all is
nicely cleaned up at the end of the test session.

Instead of the `reset_profile`, we still call `reset_communicator` and
`reset_runner` as these need to be reset for the next test run to
function properly.
---
 aiida/manage/manager.py           | 28 ++++++++++++++++++++++------
 aiida/manage/tests/main.py        |  9 ++++++---
 aiida/storage/psql_dos/backend.py |  1 +
 3 files changed, 29 insertions(+), 9 deletions(-)

diff --git a/aiida/manage/manager.py b/aiida/manage/manager.py
index 09f2232ee1..86c14ecfb9 100644
--- a/aiida/manage/manager.py
+++ b/aiida/manage/manager.py
@@ -37,7 +37,7 @@ def get_manager() -> 'Manager':
     return MANAGER
 
 
-class Manager:
+class Manager:  # pylint: disable=too-many-public-methods
     """Manager singleton for globally loaded resources.
 
     AiiDA can have the following global resources loaded:
@@ -143,17 +143,33 @@ def load_profile(self, profile: Union[None, str, 'Profile'] = None, allow_switch
 
     def reset_profile(self) -> None:
         """Close and reset any associated resources for the current profile."""
+        self.reset_profile_storage()
+        self.reset_communicator()
+        self.reset_runner()
+
+        self._daemon_client = None
+        self._persister = None
+
+    def reset_profile_storage(self) -> None:
+        """Reset the profile storage.
+
+        This will close any connections to the services used by the storage, such as database connections.
+        """
         if self._profile_storage is not None:
             self._profile_storage.close()
+        self._profile_storage = None
+
+    def reset_communicator(self) -> None:
+        """Reset the communicator."""
         if self._communicator is not None:
             self._communicator.close()
-        if self._runner is not None:
-            self._runner.stop()
-        self._profile_storage = None
         self._communicator = None
-        self._daemon_client = None
         self._process_controller = None
-        self._persister = None
+
+    def reset_runner(self) -> None:
+        """Reset the process runner."""
+        if self._runner is not None:
+            self._runner.close()
         self._runner = None
 
     def unload_profile(self) -> None:
diff --git a/aiida/manage/tests/main.py b/aiida/manage/tests/main.py
index 079f7bcd4b..fed186ebe4 100644
--- a/aiida/manage/tests/main.py
+++ b/aiida/manage/tests/main.py
@@ -163,14 +163,16 @@ def clear_profile():
         """Reset the global profile, clearing all its data and closing any open resources."""
         manager = get_manager()
         manager.get_profile_storage()._clear(recreate_user=True)  # pylint: disable=protected-access
-        manager.reset_profile()
         manager.get_profile_storage()  # reload the storage connection
+        manager.reset_communicator()
+        manager.reset_runner()
 
     def has_profile_open(self):
         return self._profile is not None
 
-    def destroy_all(self):
-        pass
+    def destroy_all(self):  # pylint: disable=no-self-use
+        manager = get_manager()
+        manager.reset_profile()
 
 
 class TemporaryProfileManager(ProfileManager):
@@ -396,6 +398,7 @@ def root_dir_ok(self):
 
     def destroy_all(self):
         """Remove all traces of the tests run"""
+        super().destroy_all()
         if self.root_dir:
             shutil.rmtree(self.root_dir)
             self.root_dir = None
diff --git a/aiida/storage/psql_dos/backend.py b/aiida/storage/psql_dos/backend.py
index eaecb88693..6dcb542561 100644
--- a/aiida/storage/psql_dos/backend.py
+++ b/aiida/storage/psql_dos/backend.py
@@ -172,6 +172,7 @@ def _clear(self, recreate_user: bool = True) -> None:
             ):
                 session.execute(table(table_name).delete())
             session.expunge_all()
+            self._default_user = None
 
             # restore the default user
             if recreate_user and default_user_kwargs: