From 7740c31874d500f93ecaaee53f4a626f7c1b60a6 Mon Sep 17 00:00:00 2001
From: Khalim Conn-Kowlessar <kconnkowlessar@gmail.com>
Date: Fri, 5 Jan 2024 16:37:05 +0000
Subject: [PATCH] set up tests for TestSolarPhotoSupply

---
 backend/app/plan/router.py                 |  2 +-
 etl/solar/SolarPhotoSupply.py              |  6 +++++-
 etl/solar/tests/test_solar_photo_supply.py | 25 ++++++++++++++++++++++
 3 files changed, 31 insertions(+), 2 deletions(-)
 create mode 100644 etl/solar/tests/test_solar_photo_supply.py

diff --git a/backend/app/plan/router.py b/backend/app/plan/router.py
index 1a499d27..217672fd 100644
--- a/backend/app/plan/router.py
+++ b/backend/app/plan/router.py
@@ -35,7 +35,7 @@ from recommendations.optimiser.GainOptimiser import GainOptimiser
 from recommendations.optimiser.optimiser_functions import prepare_input_measures
 from recommendations.Recommendations import Recommendations
 from utils.logger import setup_logger
-from utils.s3 import read_dataframe_from_s3_parquet, read_dataframe_from_s3_parquet
+from utils.s3 import read_dataframe_from_s3_parquet
 from backend.ml_models.Valuation import PropertyValuation
 from backend.ml_models.AnnualBillSavings import AnnualBillSavings
 
diff --git a/etl/solar/SolarPhotoSupply.py b/etl/solar/SolarPhotoSupply.py
index 9fad1831..6a225b5a 100644
--- a/etl/solar/SolarPhotoSupply.py
+++ b/etl/solar/SolarPhotoSupply.py
@@ -25,7 +25,7 @@ class SolarPhotoSupply:
         self.results = []
         self.decile_thresholds = None
 
-        self.roof_lookup = pd.DataFrame(cleaned_lookup["roof-description"])
+        self.roof_lookup = pd.DataFrame(cleaned_lookup.get("roof-description"))
 
         self.photo_supply_lookup = pd.DataFrame()
         self.floor_area_decile_thresholds = pd.DataFrame()
@@ -35,6 +35,10 @@ class SolarPhotoSupply:
         Create a dataset from the provided file directories. This method processes the data files,
         applies transformations, and aggregates data into a useful format.
         """
+
+        if self.roof_lookup.empty:
+            raise ValueError("No roof lookup data")
+
         results = []
 
         logger.info("Creating solar photo supply dataset")
diff --git a/etl/solar/tests/test_solar_photo_supply.py b/etl/solar/tests/test_solar_photo_supply.py
new file mode 100644
index 00000000..79d39fc4
--- /dev/null
+++ b/etl/solar/tests/test_solar_photo_supply.py
@@ -0,0 +1,25 @@
+import unittest
+from etl.solar.SolarPhotoSupply import SolarPhotoSupply
+
+
+class TestSolarPhotoSupply(unittest.TestCase):
+    def test_classify_floor_area(self):
+        # Setup
+        thresholds = [10, 20, 30, 40, 50]
+        solar_photo_supply = SolarPhotoSupply([], {})
+
+        # Test Case 1: Valid floor area
+        floor_area = 25
+        expected_decile = 2
+        result = solar_photo_supply.classify_floor_area(floor_area, thresholds)
+        self.assertEqual(result, expected_decile, "Decile classification did not match expected result")
+
+        # Test Case 2: Out of range floor area
+        floor_area = 60
+        expected_decile = len(thresholds)
+        result = solar_photo_supply.classify_floor_area(floor_area, thresholds)
+        self.assertEqual(result, expected_decile, "Decile classification for out of range value is incorrect")
+
+
+if __name__ == '__main__':
+    unittest.main()