feat: Listen to 202 status

examples/HEMS/assets_setup.py

@@ -70,6 +70,7 @@ async def create_weather_station(client: FlexMeasuresClient):
     print(f"Account ID: {account_id}")
     # Create top-level weather station asset (not public, but still under the toy account)
     # Generic asset type 7 (process) used for weather stations since no dedicated type exists
+    # TODO: remove hard-coded ID, we should actually create a weather station type somehow
     all_top_level_assets = await client.get_assets(
         include_public=True,
         depth=0,

src/flexmeasures_client/client.py

@@ -235,6 +235,7 @@ async def request(
                     headers = await self.get_headers(include_auth=include_auth)
                     try:
                         async with async_timeout.timeout(self.request_timeout):
+                            previous_polling_step = polling_step
                             (
                                 response,
                                 polling_step,
@@ -249,7 +250,10 @@ async def request(
                                 polling_step=polling_step,
                                 reauth_once=reauth_once,
                             )
-                            if response.status < 300:
+                            if (
+                                response.status < 300
+                                and polling_step == previous_polling_step
+                            ):
                                 break
                     except asyncio.TimeoutError:
                         sleep_interval = self.polling_interval * (2**polling_step)
@@ -271,6 +275,10 @@ async def request(
                         raise ConnectionError(
                             f"Error occurred while communicating with the API: {exception}"
                         ) from exception
+                else:
+                    raise ConnectionError(
+                        "Max polling steps reached while waiting for the API response."
+                    )
 
         except asyncio.TimeoutError as exception:
             raise ConnectionError(
@@ -339,7 +347,7 @@ async def request_once(
             self.server_version = header_version
 
         polling_step, reauth_once, url = await check_response(
-            self, response, polling_step, reauth_once, url
+            self, response, polling_step, reauth_once, url, method=method
         )
         return response, polling_step, reauth_once, url
 
@@ -1528,40 +1536,17 @@ async def get_forecast(
 
         This function raises a ValueError when an unhandled status code is returned.
         """
-        polling_step = 0
-        try:
-            async with async_timeout.timeout(self.polling_timeout):
-                while polling_step < self.max_polling_steps:
-                    forecast, status = await self.request(
-                        uri=f"sensors/{sensor_id}/forecasts/{forecast_id}",
-                        method="GET",
-                        minimum_server_version="0.31.0",
-                    )
-                    if status == 200:
-                        if not isinstance(forecast, dict):
-                            raise ContentTypeError(
-                                f"Expected a forecast dictionary, but got {type(forecast)}",
-                            )
-                        return forecast
-                    elif status == 202:
-                        job_status = (
-                            forecast.get("status", "unknown")
-                            if isinstance(forecast, dict)
-                            else "unknown"
-                        )
-                        message = f"Forecast job status: {job_status}. Polling step: {polling_step}. Retrying in {self.polling_interval} seconds..."
-                        self.logger.debug(message)
-                        polling_step += 1
-                        await asyncio.sleep(self.polling_interval)
-                    else:
-                        check_for_status(status, 200)
-        except asyncio.TimeoutError as exception:
-            raise ConnectionError(
-                "Client polling timeout while waiting for forecast job to complete."
-            ) from exception
-        raise ConnectionError(
-            "Max polling steps reached while waiting for forecast job to complete."
+        forecast, status = await self.request(
+            uri=f"sensors/{sensor_id}/forecasts/{forecast_id}",
+            method="GET",
+            minimum_server_version="0.31.0",
         )
+        check_for_status(status, 200)
+        if not isinstance(forecast, dict):
+            raise ContentTypeError(
+                f"Expected a forecast dictionary, but got {type(forecast)}",
+            )
+        return forecast
 
     async def trigger_and_get_forecast(
         self,

src/flexmeasures_client/response_handling.py

@@ -16,22 +16,43 @@
 
 
 async def check_response(
-    self: FlexMeasuresClient, response, polling_step: int, reauth_once: bool, url: URL
+    self: FlexMeasuresClient,
+    response,
+    polling_step: int,
+    reauth_once: bool,
+    url: URL,
+    method: str = "GET",
 ) -> tuple[int, bool, URL]:
     """
     <300: passes
+    202 on GET: job not ready yet
     303: redirect to new url
     400 + custom message: schedule not ready yet
     401: reauthenticate
     503 + Retry-After header: poll again
     otherwise: call error_handler
+
+    Returns: tuple of (polling_step, reauth_once, url)
+     - polling_step: incremented if we need to poll again, otherwise unchanged
+     - reauth_once: set to False if we re-authenticated (on 401), otherwise unchanged
+     - url: updated if we get a redirect (303), otherwise unchanged
     """
     status = response.status
     payload = await response.json()
     if payload is None:
         payload = {}
     headers = response.headers
-    if status < 300:
+    if status == 202 and method.upper() == "GET":
+        sleep_interval = self.polling_interval * (2**polling_step)
+        job_status = payload.get("status")
+        message = "Server accepted the request but the result is not ready yet."
+        if job_status:
+            message += f" Job status: {job_status}."
+        message += f" Retrying in {sleep_interval} seconds..."
+        self.logger.debug(message)
+        polling_step += 1
+        await asyncio.sleep(sleep_interval)
+    elif status < 300:
         pass
     elif response.status == 303:
         message = f"Redirect to fallback schedule: {response.headers['location']}"  # noqa: E501

tests/client/test_forecast.py

@@ -179,6 +179,41 @@ async def test_get_forecast_failed_job() -> None:
         await flexmeasures_client.close()
 
 
+@pytest.mark.asyncio
+async def test_get_forecast_polling_max_steps() -> None:
+    """Test getting a forecast fails if the job never completes."""
+    sensor_id = 1
+    forecast_id = "pending-uuid"
+    url = f"http://localhost:5000/api/v3_0/sensors/{sensor_id}/forecasts/{forecast_id}"
+
+    with aioresponses() as m:
+        m.get(
+            url=url,
+            status=202,
+            payload={"status": "QUEUED"},
+            repeat=2,
+        )
+
+        flexmeasures_client = FlexMeasuresClient(
+            email="test@test.test",
+            password="test",
+            request_timeout=2,
+            polling_interval=0,
+            max_polling_steps=2,
+            access_token="skip-auth",
+        )
+
+        with pytest.raises(
+            ConnectionError,
+            match="Max polling steps reached while waiting for the API response.",
+        ):
+            await flexmeasures_client.get_forecast(
+                sensor_id=sensor_id, forecast_id=forecast_id
+            )
+
+        await flexmeasures_client.close()
+
+
 @pytest.mark.asyncio
 async def test_trigger_and_get_forecast() -> None:
     """Test triggering and getting a forecast in one call."""

tests/client/test_schedule.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import asyncio
+import logging
 from unittest.mock import patch
 
 import pytest
@@ -116,6 +117,49 @@ async def test_get_schedule_polling() -> None:
     await flexmeasures_client.close()
 
 
+@pytest.mark.asyncio
+async def test_get_schedule_polling_accepted(caplog) -> None:
+    url = "http://localhost:5000/api/v3_0/sensors/1/schedules/some-uuid?duration=P0DT0H45M0S"  # noqa 501
+    with aioresponses() as m:
+        m.get(
+            url=url,
+            status=202,
+            payload={"status": "QUEUED"},
+        )
+        m.get(
+            url=url,
+            status=202,
+            payload={"status": "STARTED"},
+        )
+        m.get(
+            url=url,
+            status=200,
+            payload={
+                "values": [2.15, 3, 2],
+                "start": "2015-06-02T10:00:00+00:00",
+                "duration": "PT45M",
+                "unit": "MW",
+            },
+        )
+        flexmeasures_client = FlexMeasuresClient(
+            email="test@test.test",
+            password="test",
+            request_timeout=2,
+            polling_interval=0.2,
+            access_token="skip-auth",
+        )
+
+        with caplog.at_level(logging.DEBUG, logger="flexmeasures_client.client"):
+            schedule = await flexmeasures_client.get_schedule(
+                sensor_id=1, schedule_id="some-uuid", duration="PT45M"
+            )
+    assert schedule["values"] == [2.15, 3, 2]
+    assert "result is not ready yet" in caplog.text
+    assert "Job status: QUEUED" in caplog.text
+    assert "Job status: STARTED" in caplog.text
+    await flexmeasures_client.close()
+
+
 @pytest.mark.asyncio
 async def test_get_schedule_polling_exponential_backoff() -> None:
     """Test that polling uses exponential backoff (doubling the sleep interval each retry)."""