From 1f7cb9bebe2b7826a65aa9c7c4e74fac47d57aeb Mon Sep 17 00:00:00 2001
From: Mark Schreiber <mrschre@amazon.com>
Date: Thu, 19 Sep 2024 18:02:56 -0400
Subject: [PATCH] adds mocked tests for get_engine function

---
 omics/cli/run_analyzer/utils.py           |  7 ++---
 tests/cli/run_analyzer/unit/test_utils.py | 38 +++++++++++++++++++++++
 2 files changed, 41 insertions(+), 4 deletions(-)

diff --git a/omics/cli/run_analyzer/utils.py b/omics/cli/run_analyzer/utils.py
index 1af57b4..6c0c5e2 100644
--- a/omics/cli/run_analyzer/utils.py
+++ b/omics/cli/run_analyzer/utils.py
@@ -7,11 +7,10 @@
 _cwl_task_regex = r"^(^\D+)(_\d+)?$"
 
 
-def get_engine(workflow_arn, session) -> str:
-    """Get the engine name for the workflow_arn"""
-    omics = session.client("omics")
+def get_engine(workflow_arn: str, client) -> str:
+    """Get the engine name for the workflow_arn using the omics client"""
     id = workflow_arn.split("/")[-1]
-    return omics.get_workflow(id)["engine"]
+    return client.get_workflow(id=id)["engine"]
 
 
 def task_base_name(name: str, engine: str) -> str:
diff --git a/tests/cli/run_analyzer/unit/test_utils.py b/tests/cli/run_analyzer/unit/test_utils.py
index 34901ec..6e8a8e6 100644
--- a/tests/cli/run_analyzer/unit/test_utils.py
+++ b/tests/cli/run_analyzer/unit/test_utils.py
@@ -1,5 +1,8 @@
 import unittest
 
+import botocore
+from botocore.stub import Stubber
+
 import omics.cli.run_analyzer.utils as utils
 
 
@@ -35,3 +38,38 @@ def _weight(instance):
         self.assertTrue(_weight("omics.m.4xlarge") < _weight("omics.r.4xlarge"))
         self.assertTrue(_weight("omics.r.4xlarge") < _weight("omics.g4dn.4xlarge"))
         self.assertTrue(_weight("omics.r.4xlarge") < _weight("omics.g5.4xlarge"))
+
+    def test_get_engine_with_workflow_arn(self):
+        session = botocore.session.get_session()
+        region = "us-west-2"
+        omics = session.create_client(
+            "omics",
+            region,
+            aws_access_key_id="foo",
+            aws_secret_access_key="bar",
+        )
+        stubber = Stubber(omics)
+        workflow_arn = "arn:aws:omics:us-east-1:123456789012:workflow/9876"
+        stubber.add_response(
+            "get_workflow",
+            {
+                "arn": workflow_arn,
+                "id": "9876",
+                "status": "ACTIVE",
+                "type": "PRIVATE",
+                "name": "hello",
+                "engine": "WDL",
+                "main": "main.wdl",
+                "digest": "sha256:367f76a49c1e6f412a6fb319fcc7061d78ad612d06a9b8ef5b5e5f2e17a32e6f",
+                "parameterTemplate": {
+                    "param": {"description": "desc"},
+                },
+                "creationTime": "2024-04-19T14:38:56.492330+00:00",
+                "statusMessage": "status",
+                "tags": {},
+            },
+            {"id": "9876"},
+        )
+        stubber.activate()
+        self.assertEqual(utils.get_engine(workflow_arn, client=omics), "WDL")
+        stubber.deactivate()