Added a function to get the average price for a given plan

2 files changed, 23 insertions, 1 deletions

diff --git a/script.py b/script.py
index 1e3cbf1..6d5cce2 100755
--- a/script.py
+++ b/script.py
@@ -6,6 +6,8 @@ from typing import Optional
 import requests
 from dotenv import load_dotenv
 
+from utils import get_average_price
+
 load_dotenv()
 
 DEVICE_ID = os.environ.get("DEVICE_ID")
@@ -189,7 +191,7 @@ def get_current_usage_info(device_id: str, authorization: str) -> dict:
 if __name__ == "__main__":
     date_start = "2023-08-01"
     date_end = "2023-08-31"
-    average_price = (10.680 * 0.1667) + (2.223 * 0.2917) + (4.875 * 0.5416)  # Each price * percentage of time on the day
+    average_price = get_average_price("triple")
     authorization = login(EMAIL, PHONE_NUMBER)
 
     ute_historic_usage = get_ute_historic_info(DEVICE_ID, authorization, average_price, date_start, date_end)
diff --git a/utils.py b/utils.py
new file mode 100644
index 0000000..83775d5
--- /dev/null
+++ b/utils.py
@@ -0,0 +1,20 @@
+def get_average_price(plan: str) -> float:
+    """
+    Get the average price for a plan
+
+    Args:
+        plan (str): Plan name. Can be "triple" or "doble"
+
+    Returns:
+        float: Average price
+
+    Raises:
+        Exception: If the plan is invalid
+    """
+
+    if plan == "triple":
+        return (10.680 * 0.1667) + (2.223 * 0.2917) + (4.875 * 0.5416)
+    if plan == "doble":
+        return (10.680 * 0.1667) + (4.280 * 0.8333)
+
+    raise Exception("Invalid plan")