From a8be4837b35219c9abf6d1633fec1e0c5c4f1793 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Sven=20L=C3=BCtkemeier?= <sven.luetkemeier@outlook.com>
Date: Sat, 5 Apr 2025 13:22:19 +0200
Subject: [PATCH] ostrom tariff: reuse access token until expired

---
 .../electricity_tariffs/ostrom/config.py      | 14 ++++-
 .../electricity_tariffs/ostrom/tariff.py      | 57 +++++++++++++++----
 2 files changed, 58 insertions(+), 13 deletions(-)

diff --git a/packages/modules/electricity_tariffs/ostrom/config.py b/packages/modules/electricity_tariffs/ostrom/config.py
index 26fd3ef1d8..a35fa0cf43 100644
--- a/packages/modules/electricity_tariffs/ostrom/config.py
+++ b/packages/modules/electricity_tariffs/ostrom/config.py
@@ -1,14 +1,26 @@
 from typing import Optional
 
 
+class OstromToken():
+    def __init__(self,
+                 access_token: Optional[str] = None,
+                 expires_in: Optional[str] = None,
+                 created_at: Optional[str] = None) -> None:
+        self.access_token = access_token
+        self.expires_in = expires_in
+        self.created_at = created_at
+
+
 class OstromTariffConfiguration:
     def __init__(self,
                  client_id: Optional[str] = None,
                  client_secret: Optional[str] = None,
-                 zip: Optional[str] = None):
+                 zip: Optional[str] = None,
+                 token: OstromToken = None):
         self.client_id = client_id
         self.client_secret = client_secret
         self.zip = zip
+        self.token = token or OstromToken()
 
 
 class OstromTariff:
diff --git a/packages/modules/electricity_tariffs/ostrom/tariff.py b/packages/modules/electricity_tariffs/ostrom/tariff.py
index d3c9629586..aa4d6a8546 100644
--- a/packages/modules/electricity_tariffs/ostrom/tariff.py
+++ b/packages/modules/electricity_tariffs/ostrom/tariff.py
@@ -1,17 +1,33 @@
 #!/usr/bin/env python3
+import logging
 from base64 import b64encode
 from datetime import datetime, timezone, timedelta
 from urllib.parse import quote
 from typing import Dict
+from requests.exceptions import HTTPError
+from helpermodules import timecheck
 from modules.common import req
 from modules.common.abstract_device import DeviceDescriptor
 from modules.common.component_state import TariffState
 from modules.electricity_tariffs.ostrom.config import OstromTariffConfiguration
-from modules.electricity_tariffs.ostrom.config import OstromTariff
+from modules.electricity_tariffs.ostrom.config import OstromTariff, OstromToken
 
+log = logging.getLogger(__name__)
 
-def fetch_prices(config: OstromTariffConfiguration) -> Dict[int, float]:
-    access_token = req.get_http_session().post(
+
+def validate_token(config: OstromTariffConfiguration) -> None:
+    if config.token.expires_in:
+        expiration = config.token.created_at + config.token.expires_in
+        log.debug("No need to authenticate. Valid token already present.")
+        if timecheck.create_timestamp() > expiration:
+            log.debug("Access token expired. Refreshing token.")
+            _refresh_token(config)
+    else:
+        _refresh_token(config)
+
+
+def _refresh_token(config: OstromTariffConfiguration):
+    response = req.get_http_session().post(
         url="https://auth.production.ostrom-api.io/oauth2/token",
         data={"grant_type": "client_credentials"},
         headers={
@@ -19,7 +35,24 @@ def fetch_prices(config: OstromTariffConfiguration) -> Dict[int, float]:
             "content-type": "application/x-www-form-urlencoded",
             "authorization": "Basic " + b64encode((config.client_id + ":" + config.client_secret).encode()).decode()
         }
-    ).json()["access_token"]
+    ).json()
+    config.token = OstromToken(access_token=response["access_token"],
+                               expires_in=response["expires_in"],
+                               created_at=timecheck.create_timestamp())
+
+
+def fetch_prices(config: OstromTariffConfiguration) -> Dict[int, float]:
+    def get_raw_prices():
+        return req.get_http_session().get(
+            url="https://production.ostrom-api.io/spot-prices?" +
+                f"startDate={startDate}&endDate={endDate}&resolution=HOUR{zip}",
+            headers={
+                "accept": "application/json",
+                "authorization": "Bearer " + config.token.access_token
+            }
+        ).json()["data"]
+
+    validate_token(config)
     utcnow = datetime.now(timezone.utc)
     startDate = quote(utcnow.strftime("%Y-%m-%dT%H:00:00.000Z"))
     endDate = quote((utcnow + timedelta(days=1)).strftime("%Y-%m-%dT%H:00:00.000Z"))
@@ -27,14 +60,14 @@ def fetch_prices(config: OstromTariffConfiguration) -> Dict[int, float]:
         zip = f"&zip={config.zip}"
     else:
         zip = ""
-    raw_prices = req.get_http_session().get(
-        url="https://production.ostrom-api.io/spot-prices?" +
-            f"startDate={startDate}&endDate={endDate}&resolution=HOUR{zip}",
-        headers={
-            "accept": "application/json",
-            "authorization": "Bearer " + access_token
-        }
-    ).json()["data"]
+    try:
+        raw_prices = get_raw_prices()
+    except HTTPError as error:
+        if error.response.status_code == 401:
+            _refresh_token(config)
+            raw_prices = get_raw_prices()
+        else:
+            raise error
     prices: Dict[int, float] = {}
     for raw_price in raw_prices:
         # Note: with Python >= 3.11, we can use: timestamp = datetime.fromisoformat(raw_price["date"]).timestamp()