move connections to sub-modules

WIPACrepo · Jun 16, 2023 · dc65181 · dc65181
1 parent 9edbbf2
commit dc65181
Show file tree

Hide file tree

Showing 3 changed files with 43 additions and 41 deletions.
diff --git a/clientmanager/clientmanager.py b/clientmanager/clientmanager.py
@@ -5,8 +5,6 @@
 import time
 from pathlib import Path
 
-import htcondor  # type: ignore[import]
-import kubernetes  # type: ignore[import]
 from wipac_dev_tools import argparse_tools, logging_tools
 
 from . import condor, k8s
@@ -56,33 +54,11 @@ def main() -> None:
     # Go!
     match args.orchestrator:
         case "condor":
-            # condor auth & go
-            with htcondor.SecMan() as secman:
-                secman.setToken(htcondor.Token(ENV.CONDOR_TOKEN))
-                schedd_obj = condor.condor_tools.get_schedd_obj(
-                    args.collector, args.schedd
-                )
-                condor.act(args, schedd_obj)
+            condor.act(args)
         case "k8s":
-            # Creating K8S cluster client
-            k8s_client_config = kubernetes.client.Configuration()
-            if args.host == "local":
-                # use *this* pod's service account
-                kubernetes.config.load_incluster_config(k8s_client_config)
-            else:
-                # connect to remote host
-                if args.cluster_config:
-                    kubernetes.config.load_kube_config(
-                        config_file=args.cluster_config,
-                        client_configuration=k8s_client_config,
-                    )
-                k8s_client_config.host = args.host
-                k8s_client_config.api_key["authorization"] = ENV.WORKER_K8S_TOKEN
-            # connect & go
-            with kubernetes.client.ApiClient(k8s_client_config) as k8s_api_client:
-                k8s.act(args, k8s_api_client)
+            k8s.act(args)
         case other:
-            raise RuntimeError(f"Not supported orchestrator: {other}")
+            raise RuntimeError(f"Orchestrator not supported: {other}")
 
 
 class OrchestratorArgs:

diff --git a/clientmanager/condor/act.py b/clientmanager/condor/act.py
@@ -6,15 +6,20 @@
 import htcondor  # type: ignore[import]
 
 from .. import utils
-from ..config import LOGGER
-from . import starter, stopper
+from ..config import ENV, LOGGER
+from . import condor_tools, starter, stopper
 
 
-def act(
-    args: argparse.Namespace,
-    schedd_obj: htcondor.Schedd,
-) -> None:
+def act(args: argparse.Namespace) -> None:
     """Do the action."""
+    # condor auth & go
+    with htcondor.SecMan() as secman:
+        secman.setToken(htcondor.Token(ENV.CONDOR_TOKEN))
+        schedd_obj = condor_tools.get_schedd_obj(args.collector, args.schedd)
+        _act(args, schedd_obj)
+
+
+def _act(args: argparse.Namespace, schedd_obj: htcondor.Schedd) -> None:
     match args.action:
         case "start":
             LOGGER.info(

diff --git a/clientmanager/k8s/act.py b/clientmanager/k8s/act.py
@@ -11,16 +11,37 @@
 from . import starter, stopper
 
 
-def act(args: argparse.Namespace, k8s_client: kubernetes.client.ApiClient) -> None:
+def act(args: argparse.Namespace) -> None:
     """Do the action."""
-    try:
-        LOGGER.debug("testing k8s credentials")
-        api_response = kubernetes.client.BatchV1Api(k8s_client).get_api_resources()
-        LOGGER.debug(api_response)
-    except kubernetes.client.rest.ApiException as e:
-        LOGGER.exception(e)
-        raise
 
+    # Creating K8S cluster client
+    k8s_client_config = kubernetes.client.Configuration()
+    if args.host == "local":
+        # use *this* pod's service account
+        kubernetes.config.load_incluster_config(k8s_client_config)
+    else:
+        # connect to remote host
+        if args.cluster_config:
+            kubernetes.config.load_kube_config(
+                config_file=args.cluster_config,
+                client_configuration=k8s_client_config,
+            )
+        k8s_client_config.host = args.host
+        k8s_client_config.api_key["authorization"] = ENV.WORKER_K8S_TOKEN
+
+    # connect & go
+    with kubernetes.client.ApiClient(k8s_client_config) as k8s_client:
+        try:
+            LOGGER.debug("testing k8s credentials")
+            api_response = kubernetes.client.BatchV1Api(k8s_client).get_api_resources()
+            LOGGER.debug(api_response)
+        except kubernetes.client.rest.ApiException as e:
+            LOGGER.exception(e)
+            raise
+        _act(args, k8s_client)
+
+
+def _act(args: argparse.Namespace, k8s_client: kubernetes.client.ApiClient) -> None:
     match args.action:
         case "start":
             LOGGER.info(