From 5672f50d12e6d3f402f3c1a59f04a354d031d94d Mon Sep 17 00:00:00 2001
From: Iain-S <25081046+Iain-S@users.noreply.github.com>
Date: Thu, 5 Sep 2024 17:22:48 +0100
Subject: [PATCH] Don't expect total_cost field from Azure API

---
 usage_function/tests/test_utils.py | 2 --
 usage_function/utils/usage.py      | 6 ++++--
 2 files changed, 4 insertions(+), 4 deletions(-)

diff --git a/usage_function/tests/test_utils.py b/usage_function/tests/test_utils.py
index 44dfc78..2c4b1dd 100644
--- a/usage_function/tests/test_utils.py
+++ b/usage_function/tests/test_utils.py
@@ -206,14 +206,12 @@ def test_retrieve_usage_1(self) -> None:
         datum_1 = DummyAzureUsage()
         datum_1.quantity = 1
         datum_1.cost = 1
-        datum_1.total_cost = 1
         datum_1.unit_price = 1
         datum_1.effective_price = 1
 
         datum_2 = DummyAzureUsage()
         datum_2.quantity = 1
         datum_2.cost = 1
-        datum_2.total_cost = 1
         datum_2.unit_price = 1
         datum_2.effective_price = 1
 
diff --git a/usage_function/utils/usage.py b/usage_function/utils/usage.py
index bdf7ff3..4dd4b65 100644
--- a/usage_function/utils/usage.py
+++ b/usage_function/utils/usage.py
@@ -123,12 +123,14 @@ def retrieve_usage(
         if i % 200 == 0:
             logging.warning("Requesting item %d", i)
 
-        usage_item = models.Usage(**vars(item))
+        item_dict = dict(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
         # amortised_cost
-        usage_item.total_cost = usage_item.cost
+        item_dict["total_cost"] = item_dict["cost"]
+
+        usage_item = models.Usage(**item_dict)
 
         if usage_item.reservation_id is not None:
             usage_item.amortised_cost = usage_item.cost