From d3f95685b397cca83649052bc0014c3aeb26e152 Mon Sep 17 00:00:00 2001
From: Evgeny Seregin <zhenya.seregin@outlook.com>
Date: Wed, 28 Jun 2023 12:48:35 +0600
Subject: [PATCH] Fix TaskLockedException handling (#2206)

---
 sentry_sdk/integrations/huey.py      |  4 ++--
 tests/integrations/huey/test_huey.py | 28 ++++++++++++++++++++++++++++
 2 files changed, 30 insertions(+), 2 deletions(-)

diff --git a/sentry_sdk/integrations/huey.py b/sentry_sdk/integrations/huey.py
index 7c3fcbc70c..52b0e549a2 100644
--- a/sentry_sdk/integrations/huey.py
+++ b/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):
diff --git a/tests/integrations/huey/test_huey.py b/tests/integrations/huey/test_huey.py
index 819a4816d7..29e4d37027 100644
--- a/tests/integrations/huey/test_huey.py
+++ b/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()