Fix race condition in celery system integration tests

Celery stores the task result before firing the task_success signal.
This means result.get() can return before the success handler runs,
causing the mock context to exit and the second update_task_safe call
to hit the real function instead of the mock.

Add _wait_for_mock_calls() to poll until the expected call count is
reached before exiting the mock context.

Author: snopoke

tests/test_celery_system_integration.py

@@ -10,6 +10,7 @@
 
 import logging
 import sys
+import time
 import weakref
 from http import HTTPStatus
 from unittest import mock
@@ -26,6 +27,18 @@
 from tests.utils import task_for_test
 
 
+def _wait_for_mock_calls(mock_obj, expected_count, timeout=5):
+    """Wait for a mock to reach the expected call count.
+
+    Celery stores the task result before firing task_success, so
+    ``result.get()`` can return before the success signal handler runs.
+    Without this wait the mock context may exit before the handler fires.
+    """
+    deadline = time.monotonic() + timeout
+    while mock_obj.call_count < expected_count and time.monotonic() < deadline:
+        time.sleep(0.05)
+
+
 @pytest.fixture()
 def _bind_settings_with_system():
     systems = [CelerySystemIntegration()]
@@ -71,6 +84,7 @@ def add_normal(self, a, b):
         result = add_normal.delay(2, 2)
         assert result.info.get("taskbadger_task_id") == tb_task.id
         assert result.get(timeout=10, propagate=True) == 4
+        _wait_for_mock_calls(update, 2)
 
     create.assert_called_once()
     assert get_task.call_count == 1
@@ -102,6 +116,7 @@ def add_normal(self, a, b):
         result = add_normal.delay(2, 2)
         assert result.info.get("taskbadger_task_id") == tb_task.id
         assert result.get(timeout=10, propagate=True) == 4
+        _wait_for_mock_calls(update, 2)
 
     create.assert_called_once_with(
         "tests.test_celery_system_integration.add_normal",