updates for tests

awslabs · Sep 26, 2024 · 08be0fb · 08be0fb
1 parent f82e0a4
commit 08be0fb
Show file tree

Hide file tree

Showing 3 changed files with 36 additions and 33 deletions.
diff --git a/omics/cli/run_analyzer/batch.py b/omics/cli/run_analyzer/batch.py
@@ -86,7 +86,11 @@ def _aggregate_resources(run_resources_list: list[list[dict]], task_name: str, e
     print(",".join([str(aggregate.get(h, "")) for h in hdrs]), file=out)
 
 
-def _do_aggregation(resources: list, resource_key: str, operation: str):
+def _do_aggregation(run_resources_list: list[list[dict]], resource_key: str, operation: str):
+
+    # flatten the list of lists into a single list
+    resources = [r for rs in run_resources_list for r in rs]
+
     if operation == "count":
         return len(resources)
     elif operation == "sum":

diff --git a/tests/cli/run_analyzer/unit/test_batch.py b/tests/cli/run_analyzer/unit/test_batch.py
@@ -6,7 +6,7 @@
 
 class TestRunAnalyzerBatch(unittest.TestCase):
     def test_do_aggregation(self):
-        resources = [{"key": 2}, {"key": 1}, {"key": 4}]
+        resources = [[{"key": 2}, {"key": 1}, {"key": 4}]]
         result = batch._do_aggregation(resources, "key", "count")
         self.assertEqual(result, 3)
         result = batch._do_aggregation(resources, "key", "sum")
@@ -19,35 +19,37 @@ def test_do_aggregation(self):
         self.assertEqual(result, 1.53)
 
     def test_do_aggregation_with_bad_operation(self):
-        resources = [{"key": 2}, {"key": 1}, {"key": 4}]
+        resources = [[{"key": 2}, {"key": 1}, {"key": 4}]]
         self.assertRaises(ValueError, batch._do_aggregation, resources, "key", "bad_operation")
 
     def test_aggregate_resources(self):
         resources = [
-            {
-                "type": "run",
-                "runningSeconds": 10.0,
-                "cpuUtilizationRatio": 1.0,
-                "memoryUtilizationRatio": 1.0,
-                "gpusReserved": 0,
-                "recommendedCpus": 4,
-                "recommendedMemoryGiB": 8,
-                "omicsInstanceTypeMinimum": "omics.c.large",
-                "estimatedUSD": 1.00,
-                "name": "foo-01-000",
-            },
-            {
-                "type": "run",
-                "runningSeconds": 20.0,
-                "cpuUtilizationRatio": 1.0,
-                "memoryUtilizationRatio": 1.0,
-                "gpusReserved": 0,
-                "recommendedCpus": 4,
-                "recommendedMemoryGiB": 8,
-                "omicsInstanceTypeMinimum": "omics.c.large",
-                "estimatedUSD": 1.00,
-                "name": "foo-02-000",
-            },
+            [
+                {
+                    "type": "run",
+                    "runningSeconds": 10.0,
+                    "cpuUtilizationRatio": 1.0,
+                    "memoryUtilizationRatio": 1.0,
+                    "gpusReserved": 0,
+                    "recommendedCpus": 4,
+                    "recommendedMemoryGiB": 8,
+                    "omicsInstanceTypeMinimum": "omics.c.large",
+                    "estimatedUSD": 1.00,
+                    "name": "foo-01-000",
+                },
+                {
+                    "type": "run",
+                    "runningSeconds": 20.0,
+                    "cpuUtilizationRatio": 1.0,
+                    "memoryUtilizationRatio": 1.0,
+                    "gpusReserved": 0,
+                    "recommendedCpus": 4,
+                    "recommendedMemoryGiB": 8,
+                    "omicsInstanceTypeMinimum": "omics.c.large",
+                    "estimatedUSD": 1.00,
+                    "name": "foo-02-000",
+                },
+            ]
         ]
         with io.StringIO() as result:
             batch._aggregate_resources(
@@ -117,12 +119,9 @@ def test_aggregate_and_print_resources(self):
             ],
         ]
         header_string = ",".join(batch.hdrs) + "\n"
-        expected = (
-            header_string
-            + "run, foo, 2, 15.0, 30.0, 7.07, 0.5, 0.5, 0, 4, 8, omics.c.large, 1.0, 1.0\n"
-        )
+        expected = header_string + "run,foo,4,20.0,30.0,8.16,1.0,1.0,0,4,8,omics.c.large,1.0,1.0\n"
         with io.StringIO() as result:
             batch.aggregate_and_print(
-                resources_list=resources_list, pricing={}, engine="WDL", out=result
+                run_resources_list=resources_list, pricing={}, engine="WDL", out=result
             )
             self.assertEqual(result.getvalue(), expected)
diff --git a/tests/cli/run_analyzer/unit/test_utils.py b/tests/cli/run_analyzer/unit/test_utils.py
@@ -1,6 +1,6 @@
 import unittest
 
-import botocore
+import botocore.session
 from botocore.stub import Stubber
 
 import omics.cli.run_analyzer.utils as utils