diff --git a/examples/function_wrapper.py b/examples/function_wrapper.py index 9d3ed32..c7e5fa1 100644 --- a/examples/function_wrapper.py +++ b/examples/function_wrapper.py @@ -1,4 +1,4 @@ -from datetime import datetime +import datetime from taskbadger import track @@ -7,7 +7,7 @@ def my_task(arg1, arg2, kwarg1=None, kwarg2="demo"): print("Hello from my_task") print(f"arg1={arg1}, arg2={arg2}, kwarg1={kwarg1}, kwarg2={kwarg2}") - return ["Hello from my_task", datetime.utcnow()] + return ["Hello from my_task", datetime.datetime.now(datetime.timezone.utc)] if __name__ == "__main__": diff --git a/integration_tests/test_basics.py b/integration_tests/test_basics.py index 8beec1e..863e9fd 100644 --- a/integration_tests/test_basics.py +++ b/integration_tests/test_basics.py @@ -1,11 +1,11 @@ -from datetime import datetime +import datetime import taskbadger as badger from taskbadger import StatusEnum def test_basics(): - data = {"now": datetime.utcnow().isoformat()} + data = {"now": datetime.datetime.now(datetime.timezone.utc).isoformat()} task = badger.create_task("test basics", data=data) task.success(100) assert task.status == StatusEnum.SUCCESS diff --git a/taskbadger/process.py b/taskbadger/process.py index 43c8115..1a97908 100644 --- a/taskbadger/process.py +++ b/taskbadger/process.py @@ -1,7 +1,7 @@ +import datetime import subprocess import threading import time -from datetime import datetime class ProcessRunner: @@ -13,7 +13,7 @@ def __init__(self, process_args, env, capture_output: bool, update_frequency: in self.returncode = None def run(self): - last_update = datetime.utcnow() + last_update = datetime.datetime.now(datetime.timezone.utc) kwargs = {} if self.capture_output: @@ -28,7 +28,7 @@ def run(self): while process.poll() is None: time.sleep(0.1) if _should_update(last_update, self.update_frequency): - last_update = datetime.utcnow() + last_update = datetime.datetime.now(datetime.timezone.utc) if self.capture_output: yield {"stdout": stdout.read(), "stderr": stderr.read()} else: @@ -75,4 +75,4 @@ def __bool__(self): def _should_update(last_update: datetime, update_frequency_seconds): - return (datetime.utcnow() - last_update).total_seconds() >= update_frequency_seconds + return (datetime.datetime.now(datetime.timezone.utc) - last_update).total_seconds() >= update_frequency_seconds diff --git a/taskbadger/sdk.py b/taskbadger/sdk.py index a610672..f455e50 100644 --- a/taskbadger/sdk.py +++ b/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,24 @@ 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 + if self._task.updated.tzinfo is None: + tz = None + else: + # Use timezone.utc for Python <3.11 compatibility + tz = datetime.timezone.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 diff --git a/tests/test_sdk.py b/tests/test_sdk.py index b96bbfe..0f5888c 100644 --- a/tests/test_sdk.py +++ b/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) diff --git a/tests/utils.py b/tests/utils.py index 2453a4a..34d9d65 100644 --- a/tests/utils.py +++ b/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.timezone.utc) + kwargs["updated"] = datetime.datetime.now(datetime.timezone.utc) return TaskInternal( task_id, "org",