Fix test failures

usage_function/utils/usage.py

@@ -3,15 +3,15 @@
 import logging
 from datetime import datetime, timedelta
 from functools import lru_cache
-from typing import Dict, Optional, Iterable, Generator
+from typing import Dict, Generator, Iterable, Optional
 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 rctab_models import models
 
 from utils.auth import BearerAuth
 
@@ -82,9 +82,7 @@ def get_all_usage(
     return data
 
 
-def combine_items(
-    item_to_update: rctab_models.models.Usage, other_item: rctab_models.models.Usage
-) -> None:
+def combine_items(item_to_update: models.Usage, other_item: models.Usage) -> None:
     """Update one Usage with the cost, etc. of another Usage."""
     item_to_update.quantity = (item_to_update.quantity or 0) + (
         other_item.quantity or 0
@@ -107,7 +105,7 @@ def combine_items(
 
 def retrieve_usage(
     usage_data: Iterable[UsageDetailsListResult],
-) -> list[rctab_models.models.Usage]:
+) -> list[models.Usage]:
     """Retrieve usage data from Azure.
 
     Args:
@@ -118,14 +116,14 @@ def retrieve_usage(
     """
     logging.warning("Retrieve items")
 
-    all_items: Dict[str, rctab_models.models.Usage] = {}
+    all_items: Dict[str, models.Usage] = {}
     started_processing_at = datetime.now()
 
     for i, item in enumerate(usage_data):
         if i % 200 == 0:
             logging.warning("Requesting item %d", i)
 
-        usage_item = rctab_models.models.Usage(**vars(item))
+        usage_item = models.Usage(**vars(item))
 
         # When AmortizedCost metric is being used, the cost and effective_price values
         # for reserved instances are not zero, thus the cost value is moved to
@@ -196,11 +194,7 @@ def get_first_run_time() -> datetime:
 
     # Note that omitting the encoding appears to work but will
     # fail server-side with some characters, such as en-dash.
-    data = (
-        rctab_models.models.AllUsage(usage_list=all_item_list)
-        .model_dump_json()
-        .encode("utf-8")
-    )
+    data = models.AllUsage(usage_list=all_item_list).model_dump_json().encode("utf-8")
 
     for _ in range(2):
         resp = requests.post(