From 4e2a6ab61ef11a3ac7b19aab881cc3709805e8fd Mon Sep 17 00:00:00 2001
From: trevorb1 <trevor_barnes@sfu.ca>
Date: Fri, 8 Nov 2024 14:07:10 -0800
Subject: [PATCH] correct total cost for missing storage

---
 src/otoole/results/result_package.py  | 28 ++++++++++++++++++---------
 tests/results/test_results_package.py | 22 +++++++++++++++++++++
 2 files changed, 41 insertions(+), 9 deletions(-)

diff --git a/src/otoole/results/result_package.py b/src/otoole/results/result_package.py
index 9a961cb..acf1b56 100644
--- a/src/otoole/results/result_package.py
+++ b/src/otoole/results/result_package.py
@@ -952,8 +952,6 @@ def total_discounted_cost(self) -> pd.DataFrame:
         """
         try:
             discounted_cost_by_technology = self["DiscountedCostByTechnology"]
-            discounted_cost_by_storage = self["DiscountedCostByStorage"]
-
         except KeyError as ex:
             raise KeyError(self._msg("TotalDiscountedCost", str(ex)))
 
@@ -963,14 +961,26 @@ def total_discounted_cost(self) -> pd.DataFrame:
             .groupby(["REGION", "YEAR"])
             .sum()
         )
-        discounted_storage = (
-            discounted_cost_by_storage.droplevel("STORAGE")
-            .reset_index()
-            .groupby(["REGION", "YEAR"])
-            .sum()
-        )
 
-        total_discounted_cost = discounted_tech.add(discounted_storage, fill_value=0)
+        try:
+            discounted_cost_by_storage = self["DiscountedCostByStorage"]
+
+            discounted_storage = (
+                discounted_cost_by_storage.droplevel("STORAGE")
+                .reset_index()
+                .groupby(["REGION", "YEAR"])
+                .sum()
+            )
+        except KeyError as ex:  # storage not always included
+            LOGGER.debug(ex)
+
+            discounted_storage = pd.DataFrame(
+                columns=["REGION", "YEAR", "VALUE"]
+            ).set_index(["REGION", "YEAR"])
+
+        total_discounted_cost = discounted_tech.add(
+            discounted_storage, fill_value=0
+        ).astype(float)
 
         data = total_discounted_cost
 
diff --git a/tests/results/test_results_package.py b/tests/results/test_results_package.py
index d3e7619..8b90d5c 100644
--- a/tests/results/test_results_package.py
+++ b/tests/results/test_results_package.py
@@ -1238,6 +1238,28 @@ def test_calculate_total_discounted_cost(
 
         assert_frame_equal(actual, expected)
 
+    def test_calculate_total_discounted_cost_no_storage(
+        self, discounted_technology_cost
+    ):
+        """Situations where NewStorageCapacity not available"""
+
+        results = {
+            "DiscountedCostByTechnology": discounted_technology_cost,
+        }
+
+        package = ResultsPackage(results)
+        actual = package.total_discounted_cost()
+        expected = pd.DataFrame(
+            data=[
+                ["SIMPLICITY", 2014, 555.0],
+                ["SIMPLICITY", 2015, 777.0],
+                ["SIMPLICITY", 2016, 999.0],
+            ],
+            columns=["REGION", "YEAR", "VALUE"],
+        ).set_index(["REGION", "YEAR"])
+
+        assert_frame_equal(actual, expected)
+
     def test_null(self, null: ResultsPackage):
         """ """
         package = null