refactor: Use iRODSCommon as parent class

cubi_tk/irods/check.py

@@ -1,7 +1,6 @@
 """``cubi-tk irods check``: Check target iRODS collection (all md5 files? metadata md5 consistent? enough replicas?)."""
 
 import argparse
-from contextlib import contextmanager
 import json
 from multiprocessing.pool import ThreadPool
 import os
@@ -13,10 +12,11 @@
 from irods.data_object import iRODSDataObject
 from irods.models import Collection as CollectionModel
 from irods.models import DataObject as DataObjectModel
-from irods.session import iRODSSession
 from logzero import logger
 import tqdm
 
+from cubi_tk.irods_common import iRODSCommon
+
 MIN_NUM_REPLICAS = 2
 NUM_PARALLEL_TESTS = 4
 NUM_DISPLAY_FILES = 20
@@ -27,43 +27,20 @@
 DEFAULT_HASH_SCHEME = "MD5"
 
 
-class IrodsCheckCommand:
+class IrodsCheckCommand(iRODSCommon):
     """Implementation of iRDOS check command."""
 
     command_name = "check"
 
     def __init__(self, args):
+        super().__init__()
+
         #: Command line arguments.
         self.args = args
 
-        #: Path to iRODS environment file
-        self.irods_env_path = os.path.join(
-            os.path.expanduser("~"), ".irods", "irods_environment.json"
-        )
-
         #: iRODS environment
         self.irods_env = None
 
-    def _init_irods(self):
-        """Connect to iRODS."""
-        try:
-            return iRODSSession(irods_env_file=self.irods_env_path)
-        except Exception as e:
-            logger.error("iRODS connection failed: %s", self.get_irods_error(e))
-            logger.error("Are you logged in? try 'iinit'")
-            raise
-
-    @contextmanager
-    def _get_irods_sessions(self, count=NUM_PARALLEL_TESTS):
-        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()
-
     @classmethod
     def setup_argparse(cls, parser: argparse.ArgumentParser) -> None:
         parser.add_argument(
@@ -100,12 +77,6 @@ def setup_argparse(cls, parser: argparse.ArgumentParser) -> None:
         )
         parser.add_argument("irods_path", help="Path to an iRODS collection.")
 
-    @classmethod
-    def get_irods_error(cls, e: Exception):
-        """Return logger friendly iRODS exception."""
-        es = str(e)
-        return es if es != "None" else e.__class__.__name__
-
     def get_data_objs(
         self, root_coll: iRODSCollection
     ) -> typing.Dict[
@@ -162,12 +133,12 @@ def execute(self):
         if res:  # pragma: nocover
             return res
         logger.info("Starting cubi-tk irods %s", self.command_name)
-        logger.info("Args: %s", self.args)
+        logger.debug("Args: %s", self.args)
 
         # Load iRODS environment
         with open(self.irods_env_path, "r", encoding="utf-8") as f:
             irods_env = json.load(f)
-        logger.info("iRODS environment: %s", irods_env)
+        logger.debug("iRODS environment: %s", irods_env)
 
         # Connect to iRODS
         with self._get_irods_sessions(self.args.num_parallel_tests) as irods_sessions:
@@ -186,7 +157,6 @@ def execute(self):
             logger.info("Querying for data objects")
             data_objs = self.get_data_objs(root_coll)
             self.run_checks(data_objs)
-            logger.info("All done")
 
     def run_checks(self, data_objs: dict):
         """Run checks on files, in parallel if enabled."""
@@ -205,7 +175,6 @@ def run_checks(self, data_objs: dict):
             lst_files,
         )
 
-        # counter = Value(c_ulonglong, 0)
         with tqdm.tqdm(total=num_files, unit="files", unit_scale=False) as t:
             if self.args.num_parallel_tests < 2:
                 for obj in data_objs["files"]:

cubi_tk/irods_common.py

@@ -1,3 +1,4 @@
+from contextlib import contextmanager
 import getpass
 import os.path
 from pathlib import Path
@@ -135,6 +136,17 @@ def _save_irods_token(self, token: str):
         else:
             logger.warning("No token found to be saved.")
 
+    @contextmanager
+    def _get_irods_sessions(self, count: int):
+        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()

tests/test_irods_common.py

@@ -39,6 +39,18 @@ def test_init_irods(mocksession, fs):
     mocksession.assert_called()
 
 
+@patch("cubi_tk.irods_common.iRODSCommon._init_irods")
+def test_get_irods_sessions(mockinit):
+    with iRODSCommon()._get_irods_sessions(count=4) as sessions:
+        [s for s in sessions]
+    assert mockinit.call_count == 4
+
+    mockinit.reset_mock()
+    with iRODSCommon()._get_irods_sessions(count=-1) as sessions:
+        [s for s in sessions]
+    assert mockinit.call_count == 1
+
+
 @patch("getpass.getpass")
 @patch("cubi_tk.irods_common.iRODSSession")
 def test_irods_login(mocksession, mockpass, fs):