Made login improvements

2 files changed, 36 insertions, 9 deletions

diff --git a/src/ute_wrapper/ute.py b/src/ute_wrapper/ute.py
index 1560a64..795c598 100755
--- a/src/ute_wrapper/ute.py
+++ b/src/ute_wrapper/ute.py
@@ -113,8 +113,8 @@ class UTEClient:
         self.phone_number = phone_number
         self.device_id = device_id
         self.average_cost_per_kwh = average_cost_per_kwh
-        self.authorization = self._login()
         self.power_factor = power_factor
+        self.authorization = None
         self._validate_power_factor()
         self._initialize_device_id()
         self._initialize_average_cost_per_kwh()
@@ -176,11 +176,8 @@ class UTEClient:
             "Connection": "Keep-Alive",
         }
 
-        try:
-            if self.authorization:
-                headers["Authorization"] = f"Bearer {self.authorization}"
-        except AttributeError:
-            pass
+        if self.authorization:
+            headers["Authorization"] = f"Bearer {self.authorization}"
 
         for attempt in range(retries):
             try:
@@ -214,7 +211,9 @@ class UTEClient:
             "PhoneNumber": self.phone_number,
         }
 
-        return self._make_request("POST", url, data=data).text
+        response = self._make_request("POST", url, data=data)
+        self.authorization = response.text
+        return self.authorization
 
     def get_devices_list(self) -> List[dict]:
         """
@@ -223,6 +222,9 @@ class UTEClient:
         Returns:
             List[dict]: List of devices
         """
+        if not self.authorization:
+            self._login()
+
         accounts_url = f"{API_VERSION_1}/accounts"
         return self._make_request("GET", accounts_url).json()["data"]
 
@@ -233,6 +235,9 @@ class UTEClient:
         Returns:
             dict: UTE account information
         """
+        if not self.authorization:
+            self._login()
+
         accounts_by_id_url = f"{API_VERSION_1}/accounts/{self.device_id}"
         return self._make_request("GET", accounts_by_id_url).json()["data"]
 
@@ -243,6 +248,9 @@ class UTEClient:
         Returns:
             dict: UTE peak info
         """
+        if not self.authorization:
+            self._login()
+
         peak_by_id_url = f"{API_VERSION_1}/accounts/{self.device_id}/peak"
         return self._make_request("GET", peak_by_id_url).json()["data"]
 
@@ -253,6 +261,9 @@ class UTEClient:
         Returns:
             dict: UTE network status
         """
+        if not self.authorization:
+            self._login()
+
         network_status_url = f"{API_VERSION_1}/info/network/status"
         return self._make_request("GET", network_status_url).json()["data"]["summary"]
 
@@ -263,6 +274,9 @@ class UTEClient:
         Returns:
             str: UTE renewable sources percentage
         """
+        if not self.authorization:
+            self._login()
+
         global_demand_url = f"{API_VERSION_1}/info/demand/global"
         return self._make_request("GET", global_demand_url).json()["data"]["renewableSources"]
 
@@ -281,6 +295,9 @@ class UTEClient:
         Returns:
             dict: UTE info
         """
+        if not self.authorization:
+            self._login()
+
         if date_start is None or date_end is None:
             yesterday = datetime.now() - timedelta(days=1)
             yesterday_formatted = yesterday.strftime("%Y-%m-%d")
@@ -367,6 +384,9 @@ class UTEClient:
             ReadingRequestFailedException: If the reading request fails.
             ReadingResponseInvalidException: If the reading response is invalid.
         """
+        if not self.authorization:
+            self._login()
+
         reading_request_url = f"{API_VERSION_1}/device/readingRequest"
         reading_url = f"{API_VERSION_1}/device/{self.device_id}/lastReading/30"
 
diff --git a/tests/test_ute.py b/tests/test_ute.py
index 6dacb74..5469da2 100644
--- a/tests/test_ute.py
+++ b/tests/test_ute.py
@@ -42,7 +42,7 @@ def test_make_request_success(ute_client, mocker: MockerFixture):
 
 def test_make_request_no_authorization(ute_client, mocker: MockerFixture):
     """Test the _make_request method when there is no authorization token."""
-    del ute_client.authorization
+    ute_client.authorization = None
     mocked_response = mocker.Mock(status_code=requests.codes.ok, json=lambda: {"success": True})
     mocker.patch("requests.get", return_value=mocked_response)
     response = ute_client._make_request("GET", "http://example.com")
@@ -72,7 +72,7 @@ def test_method_not_supported(ute_client):
 def test_login(ute_client, mocker: MockerFixture):
     """Test the _login method."""
     mocked_response = mocker.Mock(text="mocked_token")
-    mocker.patch("requests.post", return_value=mocked_response)
+    mocker.patch.object(ute_client, "_make_request", return_value=mocked_response)
     token = ute_client._login()
     assert token == "mocked_token"
 
@@ -143,6 +143,7 @@ def test_get_historic_consumption(ute_client, mocker: MockerFixture, date_start,
             {"magnitudeVO": "IMPORT_ACTIVE_ENERGY", "date": "2023-05-02T00:00:00+00:00", "value": "20.0"},
         ],
     }
+    mocker.patch.object(ute_client, "_login")
     mocker.patch.object(ute_client, "_make_request", return_value=mocker.Mock(json=lambda: mocked_response))
     result = ute_client.get_historic_consumption(date_start, date_end)
     assert result["total"]["sum_in_kwh"] == 30.0
@@ -158,6 +159,7 @@ def test_get_current_usage_info_success(ute_client, mocker: MockerFixture):
     ]
     mocked_reading_request = mocker.Mock(status_code=requests.codes.ok)
     mocked_reading_response = mocker.Mock(json=lambda: {"success": True, "data": {"readings": readings}})
+    mocker.patch.object(ute_client, "_login")
     mocker.patch.object(ute_client, "_make_request", side_effect=[mocked_reading_request, mocked_reading_response])
 
     result = ute_client.get_current_usage_info()
@@ -170,6 +172,7 @@ def test_get_current_usage_info_success(ute_client, mocker: MockerFixture):
 
 def test_get_current_usage_info_failed_request(ute_client, mocker: MockerFixture):
     """Test the get_current_usage_info method with a failed request."""
+    mocker.patch.object(ute_client, "_login")
     mocker.patch.object(
         ute_client,
         "_make_request",
@@ -181,6 +184,7 @@ def test_get_current_usage_info_failed_request(ute_client, mocker: MockerFixture
 
 def test_get_current_usage_info_invalid_response(ute_client, mocker: MockerFixture):
     """Test the get_current_usage_info method with an invalid response."""
+    mocker.patch.object(ute_client, "_login")
     mocker.patch.object(
         ute_client,
         "_make_request",
@@ -223,6 +227,7 @@ def test_get_peak(ute_client, mocker: MockerFixture):
     }
 
     mocked_response = mocker.Mock(json=lambda: {"data": expected_peak_data})
+    mocker.patch.object(ute_client, "_login")
     mocker.patch.object(ute_client, "_make_request", return_value=mocked_response)
 
     result = ute_client.get_peak()
@@ -250,6 +255,7 @@ def test_get_network_status(ute_client, mocker: MockerFixture):
     ]
 
     mocked_response = mocker.Mock(json=lambda: {"data": {"summary": expected_network_status}})
+    mocker.patch.object(ute_client, "_login")
     mocker.patch.object(ute_client, "_make_request", return_value=mocked_response)
 
     result = ute_client.get_network_status()
@@ -266,6 +272,7 @@ def test_get_renewable_sources(ute_client, mocker: MockerFixture):
     expected_renewable_sources = "50.0%"
 
     mocked_response = mocker.Mock(json=lambda: {"data": {"renewableSources": expected_renewable_sources}})
+    mocker.patch.object(ute_client, "_login")
     mocker.patch.object(ute_client, "_make_request", return_value=mocked_response)
 
     result = ute_client.get_renewable_sources()