Fix cm unit test failure

alan-turing-institute · May 17, 2024 · 16a8cd0 · 16a8cd0
1 parent 2d0375e
commit 16a8cd0
Show file tree

Hide file tree

Showing 3 changed files with 19 additions and 12 deletions.
diff --git a/usage_function/costmanagement/__init__.py b/usage_function/costmanagement/__init__.py
@@ -8,13 +8,13 @@
 from azure.identity import DefaultAzureCredential
 from azure.mgmt.costmanagement import CostManagementClient
 from azure.mgmt.costmanagement.models import (
-    QueryDefinition,
-    QueryGrouping,
-    QueryDataset,
-    TimeframeType,
     ExportType,
     QueryAggregation,
+    QueryDataset,
+    QueryDefinition,
+    QueryGrouping,
     QueryTimePeriod,
+    TimeframeType,
 )
 
 import utils.settings
@@ -27,7 +27,7 @@
     format="%(levelname)s %(asctime)s: %(name)s - %(message)s",
     datefmt="%d/%m/%Y %I:%M:%S %p",
 )
-logger: Final = logging.getLogger(__name__)
+logger = logging.getLogger(__name__)
 
 RETRY_ATTEMPTS: Final = 5
 # We should only need one set of credentials

diff --git a/usage_function/tests/test_function_app.py b/usage_function/tests/test_function_app.py
@@ -255,19 +255,23 @@ def test_get_all_usage(self) -> None:
             )
             parameters2 = QueryDefinition(
                 time_period=QueryTimePeriod(
-                    from_property=start_datetime + timedelta(364),
+                    from_property=start_datetime + timedelta(365),
                     to=end_datetime,
                 ),
                 dataset=query_dataset,
                 type=query_type,
                 timeframe=query_timeframe,
             )
             scope = "/providers/Microsoft.Management/managementGroups/ea"
-            mock_list_func.assert_has_calls(
-                [
-                    call(scope=scope, parameters=parameters1),
-                    call(scope=scope, parameters=parameters2),
-                ]
+            self.assertEqual(mock_list_func.call_args_list[0].kwargs["scope"], scope)
+            self.assertEqual(
+                mock_list_func.call_args_list[0].kwargs["parameters"].serialize(),
+                parameters1.serialize(),
+            )
+            self.assertEqual(mock_list_func.call_args_list[1].kwargs["scope"], scope)
+            self.assertEqual(
+                mock_list_func.call_args_list[1].kwargs["parameters"].serialize(),
+                parameters2.serialize(),
             )
             self.assertEqual(mock_list_func.call_count, 2)
 

diff --git a/usage_function/utils/auth.py b/usage_function/utils/auth.py
@@ -1,4 +1,5 @@
 """Authentication between usage package and API web app."""
+
 from datetime import datetime, timedelta
 
 import jwt
@@ -36,7 +37,9 @@ def create_access_token(self):
         expire = datetime.utcnow() + access_token_expires
         token_claims = {"sub": "usage-app", "exp": expire}
 
-        return jwt.encode(token_claims, self.private_key, algorithm="RS256")  # type: ignore
+        return jwt.encode(
+            token_claims, self.private_key, algorithm="RS256"  # type: ignore
+        )
 
     def __call__(self, r: requests.PreparedRequest) -> requests.PreparedRequest:
         """Add the bearer token to the request."""