tests: replace timing sleeps with condition-based waits

Avoid flaky assertions that assume cross-thread and cross-event-loop work completes within a fixed delay. Wait for observable outcomes instead, directly await callback futures, and explicitly synchronize blocking inbox tests.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

Author: cartertinney

tests/conftest.py

@@ -4,8 +4,55 @@
 # license information.
 # --------------------------------------------------------------------------
 
+import asyncio
+import threading
+import time
+
 import pytest
 
+
+@pytest.fixture
+def wait_for():
+    def wait_for_condition(condition, timeout=5, interval=0.01):
+        deadline = time.monotonic() + timeout
+        while not condition():
+            remaining = deadline - time.monotonic()
+            if remaining <= 0:
+                pytest.fail("Timed out waiting for condition")
+            time.sleep(min(interval, remaining))
+
+    return wait_for_condition
+
+
+@pytest.fixture
+def async_wait_for():
+    async def wait_for_condition(condition, timeout=5, interval=0.01):
+        deadline = time.monotonic() + timeout
+        while not condition():
+            remaining = deadline - time.monotonic()
+            if remaining <= 0:
+                pytest.fail("Timed out waiting for condition")
+            await asyncio.sleep(min(interval, remaining))
+
+    return wait_for_condition
+
+
+@pytest.fixture
+def track_call_started(mocker):
+    def track_call(obj, method_name):
+        call_started = threading.Event()
+        original_method = getattr(obj, method_name)
+
+        def tracked_call(*args, **kwargs):
+            call_started.set()
+            return original_method(*args, **kwargs)
+
+        mocker.patch.object(obj, method_name, side_effect=tracked_call)
+        return call_started
+
+    return track_call
+
+
 """
 NOTE: ALL (yes, ALL) tests need some kind of non-specific, arbitrary exception should use
 one of the following fixtures. This is to ensure the tests operate correctly - many tests used to

tests/unit/common/pipeline/test_pipeline_stages_base.py

@@ -28,7 +28,6 @@
 from .fixtures import ArbitraryOperation
 from tests.unit.common.pipeline import pipeline_stage_test
 
-
 this_module = sys.modules[__name__]
 logging.basicConfig(level=logging.DEBUG)
 pytestmark = pytest.mark.usefixtures("fake_pipeline_thread")
@@ -184,11 +183,11 @@ def event(self):
         return pipeline_events_base.ConnectedEvent()
 
     @pytest.mark.it("Invokes the 'on_connected_handler' handler function, if set")
-    def test_invoke_handler(self, mocker, stage, event):
+    def test_invoke_handler(self, mocker, stage, event, wait_for):
         mock_handler = mocker.MagicMock()
         stage.on_connected_handler = mock_handler
         stage.handle_pipeline_event(event)
-        time.sleep(0.1)  # Needs a brief sleep so thread can switch
+        wait_for(lambda: mock_handler.call_count >= 1)
         assert mock_handler.call_count == 1
         assert mock_handler.call_args == mocker.call()
 
@@ -204,11 +203,11 @@ def event(self):
         return pipeline_events_base.DisconnectedEvent()
 
     @pytest.mark.it("Invokes the 'on_disconnected_handler' handler function, if set")
-    def test_invoke_handler(self, mocker, stage, event):
+    def test_invoke_handler(self, mocker, stage, event, wait_for):
         mock_handler = mocker.MagicMock()
         stage.on_disconnected_handler = mock_handler
         stage.handle_pipeline_event(event)
-        time.sleep(0.1)  # Needs a brief sleep so thread can switch
+        wait_for(lambda: mock_handler.call_count >= 1)
         assert mock_handler.call_count == 1
         assert mock_handler.call_args == mocker.call()
 
@@ -224,11 +223,11 @@ def event(self):
         return pipeline_events_base.NewSasTokenRequiredEvent()
 
     @pytest.mark.it("Invokes the 'on_new_sastoken_required_handler' handler function, if set")
-    def test_invoke_handler(self, mocker, stage, event):
+    def test_invoke_handler(self, mocker, stage, event, wait_for):
         mock_handler = mocker.MagicMock()
         stage.on_new_sastoken_required_handler = mock_handler
         stage.handle_pipeline_event(event)
-        time.sleep(0.1)  # Needs a brief sleep so thread can switch
+        wait_for(lambda: mock_handler.call_count >= 1)
         assert mock_handler.call_count == 1
         assert mock_handler.call_args == mocker.call()
 
@@ -246,11 +245,11 @@ def event(self, arbitrary_exception):
     @pytest.mark.it(
         "Invokes the 'on_background_exception_handler' handler function, passing the exception object, if set"
     )
-    def test_invoke_handler(self, mocker, stage, event):
+    def test_invoke_handler(self, mocker, stage, event, wait_for):
         mock_handler = mocker.MagicMock()
         stage.on_background_exception_handler = mock_handler
         stage.handle_pipeline_event(event)
-        time.sleep(0.1)  # Needs a brief sleep so thread can switch
+        wait_for(lambda: mock_handler.call_count >= 1)
         assert mock_handler.call_count == 1
         assert mock_handler.call_args == mocker.call(event.e)
 
@@ -266,11 +265,11 @@ def event(self, arbitrary_event):
         return arbitrary_event
 
     @pytest.mark.it("Invokes the 'on_pipeline_event_handler' handler function, if set")
-    def test_invoke_handler(self, mocker, stage, event):
+    def test_invoke_handler(self, mocker, stage, event, wait_for):
         mock_handler = mocker.MagicMock()
         stage.on_pipeline_event_handler = mock_handler
         stage.handle_pipeline_event(event)
-        time.sleep(0.1)  # Needs a brief sleep so thread can switch
+        wait_for(lambda: mock_handler.call_count >= 1)
         assert mock_handler.call_count == 1
         assert mock_handler.call_args == mocker.call(event)
 

tests/unit/common/test_alarm.py

@@ -42,9 +42,8 @@ def test_fn_called_w_args(self, mocker, desired_function, args, kwargs):
         a.start()
 
         assert desired_function.call_count == 0
-        time.sleep(1)  # hasn't been 2 seconds yet
-        assert desired_function.call_count == 0
-        time.sleep(1.1)  # it has now been just over 2 seconds, so the call HAS been made
+        assert a.finished.wait(timeout=5)
+        assert time.time() >= alarm_time
         assert desired_function.call_count == 1
         assert desired_function.call_args == mocker.call(*args, **kwargs)
 
@@ -55,16 +54,18 @@ def test_fn_called_no_args(self, mocker, desired_function):
         a.start()
 
         assert desired_function.call_count == 0
-        time.sleep(1.1)  # it has now been just over 1 second, so the call HAS been made
+        assert a.finished.wait(timeout=5)
+        assert time.time() >= alarm_time
         assert desired_function.call_count == 1
-        desired_function.call_args == mocker.call()
+        assert desired_function.call_args == mocker.call()
 
     @pytest.mark.it("Invokes the function immediately if the given alarm time is in the past")
     def test_alarm_already_expired(self, mocker, desired_function):
         alarm_time = time.time() - 1
         a = Alarm(alarm_time=alarm_time, function=desired_function)
         a.start()
 
+        assert a.finished.wait(timeout=5)
         assert desired_function.call_count == 1
 
     @pytest.mark.it(
@@ -75,9 +76,8 @@ def test_cancel_alarm(self, mocker, desired_function):
         a = Alarm(alarm_time=alarm_time, function=desired_function)
         a.start()
 
-        assert desired_function.call_count == 0
-        time.sleep(1)  # hasn't been 2 seconds yet
         assert desired_function.call_count == 0
         a.cancel()  # cancel the alarm
-        time.sleep(1.5)  # it has now been more than 2 seconds
-        assert desired_function.call_count == 0  # still not called
+        a.join(timeout=5)
+        assert not a.is_alive()
+        assert desired_function.call_count == 0

tests/unit/common/test_async_adapter.py

@@ -6,7 +6,6 @@
 
 import pytest
 import inspect
-import asyncio
 import logging
 import azure.iot.device.common.async_adapter as async_adapter
 
@@ -86,10 +85,9 @@ async def test_calling_object_completes_future(self):
         callback = async_adapter.AwaitableCallback()
         assert not callback.future.done()
         callback()
-        await asyncio.sleep(0.1)  # wait to give time to complete the callback
+        await callback.completion()
         assert callback.future.done()
         assert not callback.future.exception()
-        await callback.completion()
 
     @pytest.mark.it(
         "Completes the instance Future when a call is invoked on the instance (with return_arg_name)"
@@ -100,10 +98,10 @@ async def test_calling_object_completes_future_with_return_arg_name(
         callback = async_adapter.AwaitableCallback(return_arg_name="arg_name")
         assert not callback.future.done()
         callback(arg_name=fake_return_arg_value)
-        await asyncio.sleep(0.1)  # wait to give time to complete the callback
+        result = await callback.completion()
         assert callback.future.done()
         assert not callback.future.exception()
-        assert await callback.completion() == fake_return_arg_value
+        assert result == fake_return_arg_value
 
     @pytest.mark.it(
         "Raises a TypeError when a call is invoked on the instance without the correct return argument (with return_arg_name)"
@@ -122,11 +120,10 @@ async def test_raises_error_without_return_arg_name(self, arbitrary_exception):
         callback = async_adapter.AwaitableCallback()
         assert not callback.future.done()
         callback(error=arbitrary_exception)
-        await asyncio.sleep(0.1)  # wait to give time to complete the callback
-        assert callback.future.done()
-        assert callback.future.exception() == arbitrary_exception
         with pytest.raises(arbitrary_exception.__class__) as e_info:
             await callback.completion()
+        assert callback.future.done()
+        assert callback.future.exception() == arbitrary_exception
         assert e_info.value is arbitrary_exception
 
     @pytest.mark.it(
@@ -136,9 +133,8 @@ async def test_raises_error_with_return_arg_name(self, arbitrary_exception):
         callback = async_adapter.AwaitableCallback(return_arg_name="arg_name")
         assert not callback.future.done()
         callback(error=arbitrary_exception)
-        await asyncio.sleep(0.1)  # wait to give time to complete the callback
-        assert callback.future.done()
-        assert callback.future.exception() == arbitrary_exception
         with pytest.raises(arbitrary_exception.__class__) as e_info:
             await callback.completion()
+        assert callback.future.done()
+        assert callback.future.exception() == arbitrary_exception
         assert e_info.value is arbitrary_exception