remove irods session context manager and multiple sessions

cubi_tk/irods_common.py

@@ -1,4 +1,3 @@
-from contextlib import contextmanager
 import getpass
 import os.path
 from pathlib import Path
@@ -22,8 +21,6 @@
 formatter = logzero.LogFormatter(fmt="%(message)s")
 output_logger = logzero.setup_logger(formatter=formatter)
 
-NUM_PARALLEL_SESSIONS = 4
-
 
 @attrs.frozen(auto_attribs=True)
 class TransferJob:
@@ -139,17 +136,6 @@ def _save_irods_token(self, token: str):
         else:
             logger.warning("No token found to be saved.")
 
-    @contextmanager
-    def _get_irods_sessions(self, count=NUM_PARALLEL_SESSIONS):
-        if count < 1:
-            count = 1
-        irods_sessions = [self._init_irods() for _ in range(count)]
-        try:
-            yield irods_sessions
-        finally:
-            for irods in irods_sessions:
-                irods.cleanup()
-
     @property
     def session(self):
         return self._init_irods()
@@ -183,8 +169,8 @@ def destinations(self):
 
     def _create_collections(self, job: TransferJob):
         collection = str(Path(job.path_remote).parent)
-        with self._get_irods_sessions(1) as session:
-            session[0].collections.create(collection)
+        with self.session as session:
+            session.collections.create(collection)
 
     def put(self, recursive: bool = False, sync: bool = False):
         # Double tqdm for currently transferred file info
@@ -201,13 +187,13 @@ def put(self, recursive: bool = False, sync: bool = False):
                     f"File [{n + 1}/{len(self.__jobs)}]: {Path(job.path_local).name}"
                 )
                 try:
-                    with self._get_irods_sessions(1) as session:
+                    with self.session as session:
                         if recursive:
                             self._create_collections(job)
-                        if sync and session[0].data_objects.exists(job.path_remote):
+                        if sync and session.data_objects.exists(job.path_remote):
                             t.update(job.bytes)
                             continue
-                        session[0].data_objects.put(job.path_local, job.path_remote)
+                        session.data_objects.put(job.path_local, job.path_remote)
                         t.update(job.bytes)
                 except Exception as e:  # pragma: no cover
                     logger.error(f"Problem during transfer of {job.path_local}")
@@ -222,8 +208,8 @@ def chksum(self):
             if not job.path_local.endswith(".md5"):
                 output_logger.info(Path(job.path_remote).relative_to(common_prefix))
                 try:
-                    with self._get_irods_sessions(1) as session:
-                        data_object = session[0].data_objects.get(job.path_remote)
+                    with self.session as session:
+                        data_object = session.data_objects.get(job.path_remote)
                         if not data_object.checksum:
                             data_object.chksum()
                 except Exception as e:  # pragma: no cover
@@ -232,9 +218,9 @@ def chksum(self):
 
     def get(self):
         """Download files from SODAR."""
-        with self._get_irods_sessions(1) as session:
+        with self.session as session:
             self.__jobs = [
-                attrs.evolve(job, bytes=session[0].data_objects.get(job.path_remote).size)
+                attrs.evolve(job, bytes=session.data_objects.get(job.path_remote).size)
                 for job in self.__jobs
             ]
         self.__total_bytes = sum([job.bytes for job in self.__jobs])
@@ -252,8 +238,8 @@ def get(self):
                     f"File [{n + 1}/{len(self.__jobs)}]: {Path(job.path_local).name}"
                 )
                 try:
-                    with self._get_irods_sessions(1) as session:
-                        session[0].data_objects.get(job.path_remote, job.path_local)
+                    with self.session as session:
+                        session.data_objects.get(job.path_remote, job.path_local)
                     t.update(job.bytes)
                 except FileNotFoundError:  # pragma: no cover
                     raise

tests/test_irods_common.py

@@ -66,14 +66,6 @@ def test_save_irods_token(mocksession, mockencode, fs):
     mockencode.assert_called_with("secure")
 
 
-@patch("cubi_tk.irods_common.iRODSSession")
-def test_get_irods_sessions(mocksession):
-    with iRODSCommon()._get_irods_sessions(count=3) as sessions:
-        assert len(sessions) == 3
-    with iRODSCommon()._get_irods_sessions(count=-1) as sessions:
-        assert len(sessions) == 1
-
-
 # Test iRODSTransfer #########
 @pytest.fixture
 def jobs():
@@ -95,15 +87,17 @@ def test_irods_transfer_init(jobs):
         assert itransfer.destinations == [job.path_remote for job in jobs]
 
 
-@patch("cubi_tk.irods_common.iRODSSession")
+@patch("cubi_tk.irods_common.iRODSTransfer._init_irods")
 @patch("cubi_tk.irods_common.iRODSTransfer._create_collections")
 def test_irods_transfer_put(mockrecursive, mocksession, jobs):
     mockput = MagicMock()
-    mockexists = MagicMock()
+    mockexists = MagicMock(return_value=True)
     mockobj = MagicMock()
     mockobj.put = mockput
     mockobj.exists = mockexists
-    mocksession.return_value.data_objects = mockobj
+
+    # fit for context management
+    mocksession.return_value.__enter__.return_value.data_objects = mockobj
     itransfer = iRODSTransfer(jobs)
 
     # put
@@ -123,25 +117,25 @@ def test_irods_transfer_put(mockrecursive, mocksession, jobs):
     mockexists.assert_called()
 
 
-@patch("cubi_tk.irods_common.iRODSSession")
+@patch("cubi_tk.irods_common.iRODSTransfer._init_irods")
 def test_create_collections(mocksession, jobs):
     mockcreate = MagicMock()
     mockcoll = MagicMock()
     mockcoll.create = mockcreate
-    mocksession.return_value.collections = mockcoll
+    mocksession.return_value.__enter__.return_value.collections = mockcoll
     itransfer = iRODSTransfer(jobs)
 
     itransfer._create_collections(itransfer.jobs[1])
     coll_path = str(Path(itransfer.jobs[1].path_remote).parent)
     mockcreate.assert_called_with(coll_path)
 
 
-@patch("cubi_tk.irods_common.iRODSSession")
+@patch("cubi_tk.irods_common.iRODSTransfer._init_irods")
 def test_irods_transfer_chksum(mocksession, jobs):
     mockget = MagicMock()
     mockobj = MagicMock()
     mockobj.get = mockget
-    mocksession.return_value.data_objects = mockobj
+    mocksession.return_value.__enter__.return_value.data_objects = mockobj
 
     mock_data_object = MagicMock()
     mock_data_object.checksum = None
@@ -156,12 +150,12 @@ def test_irods_transfer_chksum(mocksession, jobs):
         mockget.assert_any_call(path)
 
 
-@patch("cubi_tk.irods_common.iRODSSession")
+@patch("cubi_tk.irods_common.iRODSTransfer._init_irods")
 def test_irods_transfer_get(mocksession, jobs):
     mockget = MagicMock()
     mockobj = MagicMock()
     mockobj.get = mockget
-    mocksession.return_value.data_objects = mockobj
+    mocksession.return_value.__enter__.return_value.data_objects = mockobj
     itransfer = iRODSTransfer(jobs)
 
     mockget.return_value.size = 111