diff --git a/usage_function/costmanagement/__init__.py b/usage_function/costmanagement/__init__.py index 2feeb35..700ba2d 100644 --- 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 index 0865ba4..cd3c7d9 100644 --- a/usage_function/tests/test_function_app.py +++ b/usage_function/tests/test_function_app.py @@ -255,7 +255,7 @@ 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, @@ -263,11 +263,15 @@ def test_get_all_usage(self) -> None: 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 index db03291..6e50abf 100644 --- 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."""