feat: Update scheduled retry mechanism for EHR (M2-9394) (#1884)

* feat: Update scheduled retry mechanism for EHR (M2-9394)

Adds a scheduled retry mechanism to the task to improve reliability.

This change introduces a function to check if a task should be retried based on a retry time stored in the task context. This prevents tasks from being retried immediately after a failure and allows for delayed retries.
The retry time is calculated using an exponential backoff strategy, and the task is requeued if the current time is before the retry time.

* fix: uncomment & simplify test_task_retries_on_connection_error test

* chore: fix formatting

* chore: add back user_id

---------

Co-authored-by: Marty <2614025+mbanting@users.noreply.github.com>

Author: rcmerlo

src/apps/integrations/oneup_health/service/task.py

@@ -1,9 +1,11 @@
 import random
 import uuid
 from contextlib import asynccontextmanager
-from datetime import datetime
+from datetime import datetime, timedelta, timezone
+from typing import Annotated
 
 import httpx
+from taskiq import Context, TaskiqDepends
 
 from apps.answers.crud.answers import AnswersEHRCRUD
 from apps.answers.deps.preprocess_arbitrary import get_answer_session, preprocess_arbitrary_url
@@ -136,10 +138,11 @@ async def _schedule_retry(
     delay = _exponential_backoff(retry_count)
     if delay > 0:
         retry_count += 1
-        logger.info(f"Scheduling retry #{retry_count} in {delay} seconds")
+        retry_time = datetime.now(tz=timezone.utc) + timedelta(seconds=delay)
+        logger.info(f"Scheduling retry #{retry_count} at {retry_time}")
         await (
             task_ingest_user_data.kicker()
-            .with_labels(delay=delay)
+            .with_labels(delay=60, retry_time=retry_time.isoformat())
             .kiq(
                 user_id=user_id,
                 target_subject_id=target_subject_id,
@@ -155,6 +158,24 @@ async def _schedule_retry(
     return delay > 0
 
 
+async def _check_retry_time(context: Annotated[Context, TaskiqDepends()]) -> bool:
+    """
+    Determine if the task should be retried based on the context.
+
+    Args:
+        context (Annotated[Context, TaskiqDepends()]): The context provided by Taskiq
+
+    Returns:
+        bool: True if the task should be retried, False otherwise
+    """
+    retry_time = context.message.labels.get("retry_time")
+    if retry_time and datetime.fromisoformat(retry_time) > datetime.now(tz=timezone.utc):
+        await context.requeue()
+        return False
+
+    return True
+
+
 @broker.task
 async def task_ingest_user_data(
     user_id: uuid.UUID,
@@ -165,6 +186,7 @@ async def task_ingest_user_data(
     start_date: datetime | None = None,
     retry_count: int = 0,
     failed_attempts: int = 0,
+    check_retry_time: Annotated[bool, TaskiqDepends(_check_retry_time)] = True,  # ignore:unused argument
 ) -> str | None:
     """
     Asynchronous task to ingest user health data from OneUp Health.
@@ -181,6 +203,7 @@ async def task_ingest_user_data(
         start_date (datetime, optional): The start date of the transfer process
         retry_count (int): The current retry attempt count
         failed_attempts (int): The current error retry attempt count
+        check_retry_time (bool): Flag to check if the task should be retried
 
     Returns:
         list | None: List of retrieved resources if successful, None otherwise

src/apps/integrations/oneup_health/tests/test_task.py

@@ -1,5 +1,6 @@
 import re
 import uuid
+from datetime import datetime, timedelta, timezone
 from unittest.mock import AsyncMock, MagicMock, patch
 
 import pytest
@@ -9,7 +10,6 @@
 from apps.answers.domain import EHRIngestionStatus
 from apps.applets.domain.applet_full import AppletFull
 from apps.integrations.oneup_health.service.task import task_ingest_user_data
-from broker import broker
 
 
 class TestTaskIngestUserData:
@@ -394,43 +394,50 @@ async def test_schedule_retry(self, applet_one: AppletFull):
         target_subject_id = uuid.uuid4()
         user_id = uuid.uuid4()
 
-        with patch("apps.integrations.oneup_health.service.task.task_ingest_user_data") as mock_task:
-            kicker = MagicMock()
-            with_labels = MagicMock()
-            kiq = AsyncMock()
-
-            mock_task.kicker.return_value = kicker
-            kicker.with_labels.return_value = with_labels
-            with_labels.kiq = kiq
-
-            with patch(
-                "apps.integrations.oneup_health.service.task._exponential_backoff",
-                return_value=10,
-            ) as mock_backoff:
-                await _schedule_retry(
-                    user_id=user_id,
-                    target_subject_id=target_subject_id,
-                    applet_id=applet_one.id,
-                    submit_id=submit_id,
-                    activity_id=applet_one.activities[0].id,
-                    start_date=start_date,
-                    retry_count=retry_count,
-                    failed_attempts=failed_attempts,
-                )
-
-                mock_backoff.assert_called_once_with(retry_count)
-                mock_task.kicker.assert_called_once()
-                kicker.with_labels.assert_called_once_with(delay=10)
-                kiq.assert_awaited_once_with(
-                    user_id=user_id,
-                    target_subject_id=target_subject_id,
-                    applet_id=applet_one.id,
-                    submit_id=submit_id,
-                    activity_id=applet_one.activities[0].id,
-                    start_date=start_date,
-                    retry_count=retry_count + 1,
-                    failed_attempts=failed_attempts,
-                )
+        with patch("apps.integrations.oneup_health.service.task.datetime") as mock_datetime:
+            mock_retry_time = datetime(2025, 6, 16, 12, 0, 10)  # Fixed time for testing
+            mock_datetime.now.return_value = mock_retry_time
+
+            retry_time = mock_retry_time + timedelta(seconds=10)
+
+            with patch("apps.integrations.oneup_health.service.task.task_ingest_user_data") as mock_task:
+                kicker = MagicMock()
+                with_labels = MagicMock()
+                kiq = AsyncMock()
+
+                mock_task.kicker.return_value = kicker
+                kicker.with_labels.return_value = with_labels
+                with_labels.kiq = kiq
+
+                with patch(
+                    "apps.integrations.oneup_health.service.task._exponential_backoff",
+                    return_value=10,
+                ) as mock_backoff:
+                    await _schedule_retry(
+                        user_id=user_id,
+                        target_subject_id=target_subject_id,
+                        applet_id=applet_one.id,
+                        submit_id=submit_id,
+                        activity_id=applet_one.activities[0].id,
+                        start_date=start_date,
+                        retry_count=retry_count,
+                        failed_attempts=failed_attempts,
+                    )
+
+                    mock_backoff.assert_called_once_with(retry_count)
+                    mock_task.kicker.assert_called_once()
+
+                    kicker.with_labels.assert_called_once_with(delay=60, retry_time=retry_time.isoformat())
+                    kiq.assert_awaited_once_with(
+                        user_id=user_id,
+                        target_subject_id=target_subject_id,
+                        applet_id=applet_one.id,
+                        submit_id=submit_id,
+                        activity_id=applet_one.activities[0].id,
+                        start_date=start_date,
+                        retry_count=retry_count + 1,
+                        failed_attempts=failed_attempts,
+                    )
 
     @pytest.mark.asyncio
     async def test_task_retries_on_connection_error(self, applet_one: AppletFull):
@@ -447,26 +454,51 @@ async def test_task_retries_on_connection_error(self, applet_one: AppletFull):
         target_subject_id = uuid.uuid4()
         user_id = uuid.uuid4()
 
-        with patch(
-            "apps.answers.crud.answers.AnswersEHRCRUD.upsert",
-            new=AsyncMock(side_effect=httpx.RequestError("Connection error")),
-        ):
-            submit_id = uuid.uuid4()
-            with patch.object(task_module, "_schedule_retry", wraps=task_module._schedule_retry) as mock_retry:
-                task = await task_ingest_user_data.kicker().kiq(
-                    user_id=user_id,
-                    target_subject_id=target_subject_id,
-                    applet_id=applet_one.id,
-                    submit_id=submit_id,
-                    activity_id=applet_one.activities[0].id,
-                )
-                result = await task.wait_result()
-                # The result should be None due to the connection error
-                assert result.return_value is None
-
-                # Wait for all scheduled retries to be invoked
-                # This is necessary because the retry function is asynchronous, and we need to give it time to execute.
-                await broker.wait_all()  # type: ignore
-
-                # The function should have been retried a total of 4 times, so the retry function is called 5 times.
-                assert mock_retry.call_count == 5
+        # Mock the exponential backoff to simplify and return a constant value
+        # so that retries are scheduled consistently
+        with patch("apps.integrations.oneup_health.service.task._exponential_backoff", return_value=1):
+            with patch("apps.integrations.oneup_health.service.task.datetime") as mock_datetime:
+                mock_now = datetime(2025, 6, 16, 12, 0, 0, tzinfo=timezone.utc)
+                mock_datetime.now.return_value = mock_now
+
+                with patch("apps.integrations.oneup_health.service.task.task_ingest_user_data.kicker") as mock_kicker:
+                    kicker = MagicMock()
+                    with_labels = MagicMock()
+                    kiq = AsyncMock()
+
+                    mock_kicker.return_value = kicker
+                    kicker.with_labels.return_value = with_labels
+                    with_labels.kiq = kiq
+
+                    with patch.object(task_module, "_check_retry_time", return_value=True):
+                        with patch(
+                            "apps.answers.crud.answers.AnswersEHRCRUD.upsert",
+                            new=AsyncMock(side_effect=httpx.RequestError("Connection error")),
+                        ):
+                            submit_id = uuid.uuid4()
+                            with patch.object(
+                                task_module, "_schedule_retry", wraps=task_module._schedule_retry
+                            ) as mock_retry:
+                                # Trigger the initial task
+                                task = await task_ingest_user_data(
+                                    user_id=user_id,
+                                    target_subject_id=target_subject_id,
+                                    applet_id=applet_one.id,
+                                    submit_id=submit_id,
+                                    activity_id=applet_one.activities[0].id,
+                                )
+
+                                # The result should be None due to the connection error
+                                assert task is None
+
+                                # Verify that _schedule_retry was called once for the initial error
+                                assert mock_retry.call_count == 1
+
+                                # Verify that the task was scheduled with the correct parameters
+                                kicker.with_labels.assert_called_with(
+                                    delay=60, retry_time=(mock_now + timedelta(seconds=1)).isoformat()
+                                )
+
+                                # For each retry, we would need to manually simulate the retry process
+                                # In a real scenario, we would have 5 calls total (1 initial + 4 retries)
+                                # But in this test, we're just verifying the initial scheduling works correctly