From d6879dae37c8f875f8e310f47e64991215d6164f Mon Sep 17 00:00:00 2001 From: Iain Stenson <25081046+Iain-S@users.noreply.github.com> Date: Fri, 9 Aug 2024 10:15:10 +0100 Subject: [PATCH] Add Type Hints to Usage Function (#33) * WIP * Add typing to usage function --- usage_function/tests/mypy.ini | 2 +- usage_function/tests/test_function_app.py | 2 +- usage_function/tests/test_utils.py | 14 +++++++------ usage_function/utils/auth.py | 2 +- usage_function/utils/usage.py | 24 ++++++++++++++++------- 5 files changed, 28 insertions(+), 16 deletions(-) diff --git a/usage_function/tests/mypy.ini b/usage_function/tests/mypy.ini index 193eac3..140a17e 100644 --- a/usage_function/tests/mypy.ini +++ b/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] diff --git a/usage_function/tests/test_function_app.py b/usage_function/tests/test_function_app.py index 16a525c..d18cd8f 100644 --- a/usage_function/tests/test_function_app.py +++ b/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, diff --git a/usage_function/tests/test_utils.py b/usage_function/tests/test_utils.py index d53d45f..702b2f7 100644 --- a/usage_function/tests/test_utils.py +++ b/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 = utils.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 = utils.usage.models.Usage(**usage_dict) diff --git a/usage_function/utils/auth.py b/usage_function/utils/auth.py index 6e50abf..9bcc9f7 100644 --- a/usage_function/utils/auth.py +++ b/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 diff --git a/usage_function/utils/usage.py b/usage_function/utils/usage.py index 96c453a..000e6dc 100644 --- a/usage_function/utils/usage.py +++ b/usage_function/utils/usage.py @@ -2,12 +2,14 @@ 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 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 import models from utils.auth import BearerAuth @@ -16,7 +18,9 @@ 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: @@ -32,7 +36,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: @@ -98,7 +102,7 @@ def combine_items(item_to_update: models.Usage, other_item: models.Usage) -> Non item_to_update.cost += other_item.cost -def retrieve_usage(usage_data) -> list[models.Usage]: +def retrieve_usage(usage_data: Iterable[UsageDetailsListResult]) -> list[models.Usage]: """Retrieve usage data from Azure. Args: @@ -151,7 +155,9 @@ def retrieve_usage(usage_data) -> list[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: @@ -163,11 +169,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()