Add rate limiting for task updates

taskbadger/sdk.py

@@ -1,3 +1,4 @@
+import datetime
 import logging
 import os
 from typing import Any
@@ -326,18 +327,33 @@ def update_status(self, status: StatusEnum):
         """Update the task status"""
         self.update(status=status)
 
-    def increment_progress(self, amount: int):
-        """Increment the task progress.
+    def increment_progress(self, amount: int, min_value_interval: int = None, min_time_interval: int = None):
+        """Increment the task progress by adding the specified amount to the current value.
         If the task value is not set it will be set to `amount`.
+
+        Arguments:
+            amount: The amount to increment the task value by.
+            min_value_interval: The minimum change in value required to trigger an update.
+            min_time_interval: The minimum interval between updates in seconds.
         """
         value = self._task.value
         value_norm = value if value is not UNSET and value is not None else 0
         new_amount = value_norm + amount
-        self.update(value=new_amount)
+        self.update_progress(new_amount, min_value_interval, min_time_interval)
+
+    def update_progress(self, value: int, min_value_interval: int = None, min_time_interval: int = None):
+        """Update task progress.
 
-    def update_progress(self, value: int):
-        """Update task progress."""
-        self.update(value=value)
+        Arguments:
+            value: The new value to set.
+            min_value_interval: The minimum change in value required to trigger an update.
+            min_time_interval: The minimum interval between updates in seconds.
+        """
+        skip_check = not (min_value_interval or min_time_interval)
+        time_check = min_time_interval and self._check_update_time_interval(min_time_interval)
+        value_check = min_value_interval and self._check_update_value_interval(value, min_value_interval)
+        if skip_check or time_check or value_check:
+            self.update(value=value)
 
     def set_value_max(self, value_max: int):
         """Set the `value_max`."""
@@ -392,10 +408,16 @@ def tag(self, tags: dict[str, str]):
         """Add tags to the task."""
         self.update(tags=tags)
 
-    def ping(self):
+    def ping(self, min_time_interval=None):
         """Update the task without changing any values. This can be used in conjunction
-        with 'stale_timeout' to indicate that the task is still running."""
-        self.update()
+        with 'stale_timeout' to indicate that the task is still running.
+
+        Arguments:
+            min_time_interval: The minimum interval between pings in seconds. If set this will only
+                update the task if the last update was more than `min_time_interval` seconds ago.
+        """
+        if self._check_update_time_interval(min_time_interval):
+            self.update()
 
     @property
     def tags(self):
@@ -410,6 +432,20 @@ def safe_update(self, **kwargs):
         except Exception as e:
             log.warning("Error updating task '%s': %s", self._task.id, e)
 
+    def _check_update_time_interval(self, min_time_interval: int = None):
+        if min_time_interval and self._task.updated:
+            # tzinfo should always be set but for the sake of safety we check
+            tz = None if self._task.updated.tzinfo is None else datetime.UTC
+            now = datetime.datetime.now(tz)
+            time_since = now - self._task.updated
+            return time_since.total_seconds() >= min_time_interval
+        return True
+
+    def _check_update_value_interval(self, new_value, min_value_interval: int = None):
+        if min_value_interval and self._task.value:
+            return new_value - self._task.value >= min_value_interval
+        return True
+
 
 def _none_to_unset(value):
     return UNSET if value is None else value

tests/test_sdk.py

@@ -1,3 +1,4 @@
+import datetime
 from http import HTTPStatus
 from unittest import mock
 
@@ -153,6 +154,85 @@ def test_update_data(settings, patched_update):
     _verify_update(settings, patched_update, data={"a": 1})
 
 
+def test_ping(settings, patched_update):
+    task = Task(task_for_test())
+
+    updated_at = task.updated
+    patched_update.return_value = Response(HTTPStatus.OK, b"", {}, task_for_test())
+    task.ping(min_time_interval=1)
+    assert len(patched_update.call_args_list) == 0
+
+    task.ping()
+    _verify_update(settings, patched_update)
+    assert task.updated > updated_at
+
+    task.ping(min_time_interval=1)
+    assert len(patched_update.call_args_list) == 1
+
+    task._task.updated = task._task.updated - datetime.timedelta(seconds=1)
+    task.ping(min_time_interval=1)
+    assert len(patched_update.call_args_list) == 2
+
+
+def test_update_progress_min_time_interval(settings, patched_update):
+    task = Task(task_for_test(value=1))
+
+    updated_at = task.updated
+    patched_update.return_value = Response(HTTPStatus.OK, b"", {}, task_for_test())
+    task.update_progress(2, min_time_interval=1)
+    assert len(patched_update.call_args_list) == 0
+
+    task.update_progress(2)
+    _verify_update(settings, patched_update, value=2)
+    assert task.updated > updated_at
+
+    task.update_progress(3, min_time_interval=1)
+    assert len(patched_update.call_args_list) == 1
+
+    task._task.updated = task._task.updated - datetime.timedelta(seconds=1)
+    task.update_progress(3, min_time_interval=1)
+    assert len(patched_update.call_args_list) == 2
+
+
+def test_update_progress_min_value_interval(settings, patched_update):
+    task = Task(task_for_test(value=1))
+
+    patched_update.return_value = Response(HTTPStatus.OK, b"", {}, task_for_test(value=4))
+    task.update_progress(4, min_value_interval=5)
+    assert len(patched_update.call_args_list) == 0
+
+    task.update_progress(4)
+    _verify_update(settings, patched_update, value=4)
+
+    task.update_progress(8, min_value_interval=5)
+    assert len(patched_update.call_args_list) == 1
+
+    task.update_progress(9, min_value_interval=5)
+    assert len(patched_update.call_args_list) == 2
+
+
+def test_update_progress_min_interval_both(settings, patched_update):
+    task = Task(task_for_test(value=1))
+
+    patched_update.return_value = Response(HTTPStatus.OK, b"", {}, task_for_test(value=4))
+    # neither checks pass
+    task.update_progress(4, min_time_interval=1, min_value_interval=5)
+    assert len(patched_update.call_args_list) == 0
+
+    # value check passes
+    task.update_progress(6, min_time_interval=1, min_value_interval=5)
+    _verify_update(settings, patched_update, value=6)
+
+    # neither checks pass
+    task.update_progress(8, min_time_interval=1, min_value_interval=5)
+    assert len(patched_update.call_args_list) == 1
+
+    # time check passes
+    task._task.updated = task._task.updated - datetime.timedelta(seconds=1)
+    task.update_progress(6, min_time_interval=1, min_value_interval=5)
+    assert len(patched_update.call_args_list) == 2
+
+
 def test_increment_progress(settings, patched_update):
     api_task = task_for_test()
     task = Task(api_task)

tests/utils.py

@@ -1,4 +1,4 @@
-from datetime import datetime
+import datetime
 from uuid import uuid4
 
 from taskbadger.internal.models import Task as TaskInternal
@@ -12,8 +12,8 @@ def task_for_test(**kwargs):
     kwargs["url"] = None
     kwargs["public_url"] = None
     kwargs["value_percent"] = None
-    kwargs["created"] = datetime.utcnow()
-    kwargs["updated"] = datetime.utcnow()
+    kwargs["created"] = datetime.datetime.now(datetime.UTC)
+    kwargs["updated"] = datetime.datetime.now(datetime.UTC)
     return TaskInternal(
         task_id,
         "org",