chore: fix streaming pull close test flakiness

Author: plamut

tests/unit/pubsub_v1/subscriber/test_streaming_pull_manager.py

@@ -550,6 +550,26 @@ def make_running_manager(**kwargs):
     )
 
 
+def await_manager_shutdown(timeout=None):  # pragma: NO COVER
+    # NOTE: This method should be called after manager.close(), i.e. after the shutdown
+    # thread has been started.
+    shutdown_thread = next(
+        (
+            thread
+            for thread in threading.enumerate()
+            if thread.name == streaming_pull_manager._REGULAR_SHUTDOWN_THREAD_NAME
+        ),
+        None,
+    )
+
+    if shutdown_thread is None:
+        return  # Shutdown already finished.
+
+    shutdown_thread.join(timeout=timeout)
+    if shutdown_thread.is_alive():
+        pytest.fail("Shutdown not completed in time.")
+
+
 def test_close():
     (
         manager,
@@ -561,6 +581,7 @@ def test_close():
     ) = make_running_manager()
 
     manager.close()
+    await_manager_shutdown(timeout=3)
 
     consumer.stop.assert_called_once()
     leaser.stop.assert_called_once()
@@ -583,6 +604,7 @@ def test_close_inactive_consumer():
     consumer.is_active = False
 
     manager.close()
+    await_manager_shutdown(timeout=3)
 
     consumer.stop.assert_not_called()
     leaser.stop.assert_called_once()
@@ -596,6 +618,7 @@ def test_close_idempotent():
 
     manager.close()
     manager.close()
+    await_manager_shutdown(timeout=3)
 
     assert scheduler.shutdown.call_count == 1
 
@@ -640,6 +663,7 @@ def test_close_no_dispatcher_error():
     dispatcher.start()
 
     manager.close()
+    await_manager_shutdown(timeout=3)
 
     error_callback.assert_not_called()
 
@@ -651,6 +675,7 @@ def test_close_callbacks():
 
     manager.add_close_callback(callback)
     manager.close(reason="meep")
+    await_manager_shutdown(timeout=3)
 
     callback.assert_called_once_with(manager, "meep")
 
@@ -660,6 +685,7 @@ def test_close_blocking_scheduler_shutdown():
     scheduler = manager._scheduler
 
     manager.close()
+    await_manager_shutdown(timeout=3)
 
     scheduler.shutdown.assert_called_once_with(await_msg_callbacks=True)
 
@@ -669,6 +695,7 @@ def test_close_nonblocking_scheduler_shutdown():
     scheduler = manager._scheduler
 
     manager.close()
+    await_manager_shutdown(timeout=3)
 
     scheduler.shutdown.assert_called_once_with(await_msg_callbacks=False)
 
@@ -690,6 +717,7 @@ def fake_nack(self):
     manager._messages_on_hold._messages_on_hold.append(messages[2])
 
     manager.close()
+    await_manager_shutdown(timeout=3)
 
     assert sorted(nacked_messages) == [b"msg1", b"msg2", b"msg3"]