diff --git a/tests/test_taskqueue.py b/tests/test_taskqueue.py index d7e1a72..925442c 100644 --- a/tests/test_taskqueue.py +++ b/tests/test_taskqueue.py @@ -216,6 +216,17 @@ def test_dataset_get_status(): collection_id="1", ) + task_three = Task( + task_id="3", + job_id="abc", + delivery_tag="", + operation="index", + context={}, + payload={}, + priority=5, + collection_id="1", + ) + # Adding a task updates `start_time` and `last_update` with time_machine.travel("2024-01-01T00:00:00"): dataset.add_task(task_one.task_id, task_one.operation) @@ -284,7 +295,7 @@ def test_dataset_get_status(): # Adding a new task to an inactive dataset sets `start_time` with time_machine.travel("2024-01-07T00:00:00"): - dataset.add_task("3", "analyze") + dataset.add_task(task_three.task_id, task_three.operation) status = dataset.get_status() assert status["pending"] == 1 @@ -293,6 +304,29 @@ def test_dataset_get_status(): assert status["start_time"].startswith("2024-01-07T00:00:00") assert status["last_update"].startswith("2024-01-07T00:00:00") + # Cancelling a dataset flushes status data + with time_machine.travel("2024-01-08T00:00:00"): + dataset.checkout_task(task_three.task_id, task_three.operation) + dataset.cancel() + + status = dataset.get_status() + assert status["pending"] == 0 + assert status["running"] == 0 + assert status["finished"] == 0 + assert status["start_time"] is None + assert status["last_update"] is None + + # Tasks that were already running when the dataset was cancelled + # have no effect + with time_machine.travel("2024-01-09T00:00:00"): + dataset.mark_done(task_three) + + assert status["pending"] == 0 + assert status["running"] == 0 + assert status["finished"] == 0 + assert status["start_time"] is None + assert status["last_update"] is None + def test_get_priority_bucket(): redis = get_fakeredis()