tests: preserve timing semantics in hardened tests

Retain not-before and negative observation assertions where elapsed time is part of the behavior. Use explicit runner barriers to prove handler removal occurs while inbox work is still pending.

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

Author: cartertinney

tests/unit/common/test_alarm.py

@@ -41,6 +41,8 @@ def test_fn_called_w_args(self, mocker, desired_function, args, kwargs):
         a = Alarm(alarm_time=alarm_time, function=desired_function, args=args, kwargs=kwargs)
         a.start()
 
+        assert desired_function.call_count == 0
+        assert not a.finished.wait(timeout=1)
         assert desired_function.call_count == 0
         assert a.finished.wait(timeout=5)
         assert time.time() >= alarm_time
@@ -76,6 +78,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
+        assert not a.finished.wait(timeout=1)
         assert desired_function.call_count == 0
         a.cancel()  # cancel the alarm
         a.join(timeout=5)

tests/unit/iothub/aio/test_async_handler_manager.py

@@ -5,12 +5,16 @@
 # --------------------------------------------------------------------------
 import logging
 import pytest
+import asyncio
 import threading
 import concurrent.futures
 from azure.iot.device.common import handle_exceptions
 from azure.iot.device.iothub import client_event
 from azure.iot.device.iothub.aio.async_handler_manager import AsyncHandlerManager
-from azure.iot.device.iothub.sync_handler_manager import HandlerManagerException
+from azure.iot.device.iothub.sync_handler_manager import (
+    HandlerManagerException,
+    HandlerRunnerKillerSentinel,
+)
 from azure.iot.device.iothub.sync_handler_manager import MESSAGE, METHOD, TWIN_DP_PATCH
 from azure.iot.device.iothub.inbox_manager import InboxManager
 from azure.iot.device.iothub.aio.async_inbox import AsyncClientInbox
@@ -451,8 +455,31 @@ async def test_handler_resolve_pending_items_before_handler_removal(
         handler,
         handler_checker,
         inbox,
-        async_wait_for,
     ):
+        original_get = inbox.get
+        original_put = inbox.put
+        runner_paused = threading.Event()
+        resume_runner = threading.Event()
+        sentinel_added = threading.Event()
+        get_count = 0
+
+        async def controlled_get():
+            nonlocal get_count
+            item = await original_get()
+            get_count += 1
+            if get_count == 2:
+                runner_paused.set()
+                assert await asyncio.to_thread(resume_runner.wait, 5)
+            return item
+
+        def tracked_put(item):
+            original_put(item)
+            if isinstance(item, HandlerRunnerKillerSentinel):
+                sentinel_added.set()
+
+        mocker.patch.object(inbox, "get", side_effect=controlled_get)
+        mocker.patch.object(inbox, "put", side_effect=tracked_put)
+
         assert inbox.empty()
         # Queue up a bunch of items in the inbox
         for _ in range(100):
@@ -463,9 +490,22 @@ async def test_handler_resolve_pending_items_before_handler_removal(
         assert not inbox.empty()
         # Set the handler
         setattr(handler_manager, handler_name, handler)
-        # Immediately remove the handler
-        setattr(handler_manager, handler_name, None)
-        await async_wait_for(lambda: handler_checker.handler_call_count >= 100)
+        try:
+            assert runner_paused.wait(timeout=5)
+            assert handler_checker.handler_call_count < 100
+            assert not inbox.empty()
+
+            # Remove the handler while items are still pending
+            loop = asyncio.get_running_loop()
+            removal_future = loop.run_in_executor(
+                None, setattr, handler_manager, handler_name, None
+            )
+            assert sentinel_added.wait(timeout=5)
+            assert not removal_future.done()
+        finally:
+            resume_runner.set()
+        await asyncio.wait_for(removal_future, timeout=5)
+
         # Despite removal, handler has been called for everything that was in the inbox at the
         # time of the removal
         assert handler_checker.handler_call_count == 100
@@ -474,6 +514,8 @@ async def test_handler_resolve_pending_items_before_handler_removal(
         # Add some more items
         for _ in range(100):
             inbox.put(mocker.MagicMock())
+        # Give any incorrectly retained runner an opportunity to invoke the handler
+        await asyncio.sleep(0.1)
         # Despite more items added to inbox, no further handler calls have been made beyond the
         # initial calls that were made when the original items were added
         assert handler_checker.handler_call_count == 100

tests/unit/iothub/test_sync_handler_manager.py

@@ -6,9 +6,15 @@
 import logging
 import pytest
 import threading
+import time
+import concurrent.futures
 from azure.iot.device.common import handle_exceptions
 from azure.iot.device.iothub import client_event
-from azure.iot.device.iothub.sync_handler_manager import SyncHandlerManager, HandlerManagerException
+from azure.iot.device.iothub.sync_handler_manager import (
+    SyncHandlerManager,
+    HandlerManagerException,
+    HandlerRunnerKillerSentinel,
+)
 from azure.iot.device.iothub.sync_handler_manager import MESSAGE, METHOD, TWIN_DP_PATCH
 from azure.iot.device.iothub.inbox_manager import InboxManager
 from azure.iot.device.iothub.sync_inbox import SyncClientInbox
@@ -387,6 +393,30 @@ def test_handler_invoked_multiple(self, mocker, handler_name, handler_manager, i
     def test_handler_resolve_pending_items_before_handler_removal(
         self, mocker, handler_name, handler_manager, inbox
     ):
+        original_get = inbox.get
+        original_put = inbox.put
+        runner_paused = threading.Event()
+        resume_runner = threading.Event()
+        sentinel_added = threading.Event()
+        get_count = 0
+
+        def controlled_get(*args, **kwargs):
+            nonlocal get_count
+            item = original_get(*args, **kwargs)
+            get_count += 1
+            if get_count == 2:
+                runner_paused.set()
+                assert resume_runner.wait(timeout=5)
+            return item
+
+        def tracked_put(item):
+            original_put(item)
+            if isinstance(item, HandlerRunnerKillerSentinel):
+                sentinel_added.set()
+
+        mocker.patch.object(inbox, "get", side_effect=controlled_get)
+        mocker.patch.object(inbox, "put", side_effect=tracked_put)
+
         # Use a threadsafe mock to ensure accurate counts
         mock_handler = ThreadsafeMock()
         assert inbox.empty()
@@ -399,13 +429,23 @@ def test_handler_resolve_pending_items_before_handler_removal(
         assert not inbox.empty()
         # Set the handler
         setattr(handler_manager, handler_name, mock_handler)
-        # The handler has not yet been called for everything that was in the inbox
-        # NOTE: I'd really like to show that the handler call count is also > 0 here, but
-        # it's pretty difficult to make the timing work
-        assert mock_handler.call_count < 100
+        try:
+            assert runner_paused.wait(timeout=5)
+            assert mock_handler.call_count < 100
+            assert not inbox.empty()
+
+            # Remove the handler while items are still pending
+            with concurrent.futures.ThreadPoolExecutor(max_workers=1) as executor:
+                removal_future = executor.submit(setattr, handler_manager, handler_name, None)
+                try:
+                    assert sentinel_added.wait(timeout=5)
+                    assert not removal_future.done()
+                finally:
+                    resume_runner.set()
+                removal_future.result(timeout=5)
+        finally:
+            resume_runner.set()
 
-        # Immediately remove the handler
-        setattr(handler_manager, handler_name, None)
         # Despite removal, handler has been called for everything that was in the inbox at the
         # time of the removal
         assert mock_handler.call_count == 100
@@ -414,6 +454,8 @@ def test_handler_resolve_pending_items_before_handler_removal(
         # Add some more items
         for _ in range(100):
             inbox.put(mocker.MagicMock())
+        # Give any incorrectly retained runner an opportunity to invoke the handler
+        time.sleep(0.2)
         # Despite more items added to inbox, no further handler calls have been made beyond the
         # initial calls that were made when the original items were added
         assert mock_handler.call_count == 100