Merge branch 'main' into rctab-models

usage_function/tests/mypy.ini

@@ -1,6 +1,6 @@
 [mypy]
 ignore_missing_imports = True
-check_untyped_defs = True
+disallow_untyped_defs = True
 plugins = pydantic.mypy
 
 [pydantic-mypy]

usage_function/tests/test_function_app.py

@@ -316,7 +316,7 @@ def test_send_usage(self) -> None:
 
                 with patch("costmanagement.logger.warning") as mock_log:
 
-                    def send():
+                    def send() -> None:
                         costmanagement.send_usage(
                             "https://123.234.345.456",
                             local_usage,

usage_function/tests/test_utils.py

@@ -30,10 +30,12 @@ def test_get_all_usage(self) -> None:
             mock_list_func.return_value = expected
 
             jan_tenth = datetime(2021, 1, 10, 1, 1, 1, 1)
-            actual = utils.usage.get_all_usage(
-                jan_tenth - timedelta(days=5),
-                jan_tenth,
-                mgmt_group="some-mgmt-group",
+            actual = list(
+                utils.usage.get_all_usage(
+                    jan_tenth - timedelta(days=5),
+                    jan_tenth,
+                    mgmt_group="some-mgmt-group",
+                )
             )
 
             mock_client.assert_called_once_with(
@@ -123,7 +125,7 @@ def test_retrieve_and_send_usage(self) -> None:
                     with self.assertRaises(RuntimeError):
                         utils.usage.retrieve_and_send_usage(
                             HTTP_ADAPTER.validate_python("https://123.123.123.123"),
-                            [example_usage_detail],
+                            [example_usage_detail],  # type: ignore
                         )
 
                     usage = rctab_models.models.Usage(**usage_dict)
@@ -157,7 +159,7 @@ def test_retrieve_and_send_usage(self) -> None:
                 with patch("usage.logging.warning"):
                     utils.usage.retrieve_and_send_usage(
                         HTTP_ADAPTER.validate_python("https://123.123.123.123"),
-                        [example_usage_detail],
+                        [example_usage_detail],  # type: ignore
                     )
 
                     usage = rctab_models.models.Usage(**usage_dict)

usage_function/utils/auth.py

@@ -31,7 +31,7 @@ def __init__(self) -> None:
             private_key_txt.encode(), password=b""
         )
 
-    def create_access_token(self):
+    def create_access_token(self) -> str:
         """Create an access token."""
         access_token_expires = timedelta(minutes=ACCESS_TOKEN_EXPIRE_MINUTES)
         expire = datetime.utcnow() + access_token_expires

usage_function/utils/usage.py

@@ -3,21 +3,25 @@
 import logging
 from datetime import datetime, timedelta
 from functools import lru_cache
-from typing import Dict, Optional
+from typing import Dict, Optional, Iterable, Generator
 from uuid import UUID
 
 import rctab_models.models
 import requests
 from azure.identity import DefaultAzureCredential
 from azure.mgmt.consumption import ConsumptionManagementClient
+from azure.mgmt.consumption.models import UsageDetailsListResult
+from pydantic_core import Url
 
 from utils.auth import BearerAuth
 
 # We should only need one set of credentials
 CREDENTIALS = DefaultAzureCredential(exclude_shared_token_cache_credential=True)
 
 
-def date_range(start_date, end_date):
+def date_range(
+    start_date: datetime, end_date: datetime
+) -> Generator[datetime, None, None]:
     """Yield a datetime day for each day between start_date and end_date (inclusive).
 
     Args:
@@ -33,7 +37,7 @@ def get_all_usage(
     end_time: datetime,
     billing_account_id: Optional[str] = None,
     mgmt_group: Optional[str] = None,
-):
+) -> Iterable[UsageDetailsListResult]:
     """Get Azure usage data for a subscription between start_time and end_time.
 
     Args:
@@ -101,7 +105,8 @@ def combine_items(
     item_to_update.cost += other_item.cost
 
 
-def retrieve_usage(usage_data) -> list[rctab_models.models.Usage]:
+def retrieve_usage(usage_data: Iterable[UsageDetailsListResult]) -> list[rctab_models.models.Usage]:
+
     """Retrieve usage data from Azure.
 
     Args:
@@ -154,7 +159,9 @@ def retrieve_usage(usage_data) -> list[rctab_models.models.Usage]:
     return list(all_items.values())
 
 
-def retrieve_and_send_usage(hostname_or_ip, usage_data):
+def retrieve_and_send_usage(
+    hostname_or_ip: Url, usage_data: Iterable[UsageDetailsListResult]
+) -> None:
     """Retrieve usage data from Azure and send it to the API.
 
     Args:
@@ -166,11 +173,15 @@ def retrieve_and_send_usage(hostname_or_ip, usage_data):
     send_usage(hostname_or_ip, usage_list)
 
 
-def send_usage(hostname_or_ip, all_item_list, monthly_usage_upload=False):
+def send_usage(
+    hostname_or_ip: Url,
+    all_item_list: list[models.Usage],
+    monthly_usage_upload: bool = False,
+) -> None:
     """Post each item of usage_data to a route."""
 
     @lru_cache(1)
-    def get_first_run_time():
+    def get_first_run_time() -> datetime:
         return datetime.now()
 
     started_processing_at = datetime.now()