Fix TaskLockedException handling

Author: Zhenay

sentry_sdk/integrations/huey.py

@@ -26,12 +26,12 @@
 
 try:
     from huey.api import Huey, Result, ResultGroup, Task
-    from huey.exceptions import CancelExecution, RetryTask
+    from huey.exceptions import CancelExecution, RetryTask, TaskLockedException
 except ImportError:
     raise DidNotEnable("Huey is not installed")
 
 
-HUEY_CONTROL_FLOW_EXCEPTIONS = (CancelExecution, RetryTask)
+HUEY_CONTROL_FLOW_EXCEPTIONS = (CancelExecution, RetryTask, TaskLockedException)
 
 
 class HueyIntegration(Integration):

tests/integrations/huey/test_huey.py

@@ -118,6 +118,34 @@ def retry_task(context):
     assert len(huey) == 0
 
 
+@pytest.mark.parametrize("lock_name", ["lock.a", "lock.b"], ids=["locked", "unlocked"])
+def test_task_lock(capture_events, init_huey, lock_name):
+    huey = init_huey()
+
+    task_lock_name = "lock.a"
+    should_be_locked = task_lock_name == lock_name
+
+    @huey.task()
+    @huey.lock_task(task_lock_name)
+    def maybe_locked_task():
+        pass
+
+    events = capture_events()
+
+    with huey.lock_task(lock_name):
+        assert huey.is_locked(task_lock_name) == should_be_locked
+        result = execute_huey_task(huey, maybe_locked_task)
+
+    (event,) = events
+
+    assert event["transaction"] == "maybe_locked_task"
+    assert event["tags"]["huey_task_id"] == result.task.id
+    assert (
+        event["contexts"]["trace"]["status"] == "aborted" if should_be_locked else "ok"
+    )
+    assert len(huey) == 0
+
+
 def test_huey_enqueue(init_huey, capture_events):
     huey = init_huey()