From be5b8c2e54822de5be795bbee467477d160115d7 Mon Sep 17 00:00:00 2001
From: "Daniel A. Wozniak" <dwozniak@vmware.com>
Date: Fri, 2 Feb 2024 14:57:43 -0700
Subject: [PATCH] Connect master tests need to be coroutines

---
 salt/transport/zeromq.py          |  7 +++-
 tests/pytests/unit/test_minion.py | 55 +++++++++++++++++++++++++++----
 2 files changed, 54 insertions(+), 8 deletions(-)

diff --git a/salt/transport/zeromq.py b/salt/transport/zeromq.py
index 07f2bc1bf896..703b6ab2cf59 100644
--- a/salt/transport/zeromq.py
+++ b/salt/transport/zeromq.py
@@ -281,7 +281,12 @@ def on_recv(self, callback):
 
         :param func callback: A function which should be called when data is received
         """
-        return self.stream.on_recv(callback)
+        try:
+            return self.stream.on_recv(callback)
+        except OSError as exc:
+            if callback is None and str(exc) == "Stream is closed":
+                return
+            raise
 
     @salt.ext.tornado.gen.coroutine
     def send(self, msg):
diff --git a/tests/pytests/unit/test_minion.py b/tests/pytests/unit/test_minion.py
index 6839275ffa3c..e5ca73a3f9a7 100644
--- a/tests/pytests/unit/test_minion.py
+++ b/tests/pytests/unit/test_minion.py
@@ -22,6 +22,31 @@
 log = logging.getLogger(__name__)
 
 
+@pytest.fixture
+def connect_master_mock():
+    class ConnectMasterMock:
+        """
+        Mock connect master call.
+
+        The first call will raise an exception stored on the exc attribute.
+        Subsequent calls will return True.
+        """
+
+        def __init__(self):
+            self.calls = 0
+            self.exc = Exception
+
+        @salt.ext.tornado.gen.coroutine
+        def __call__(self, *args, **kwargs):
+            self.calls += 1
+            if self.calls == 1:
+                raise self.exc()
+            else:
+                return True
+
+    return ConnectMasterMock()
+
+
 def test_minion_load_grains_false(minion_opts):
     """
     Minion does not generate grains when load_grains is False
@@ -1121,37 +1146,53 @@ def test_load_args_and_kwargs(minion_opts):
         ret = salt.minion.load_args_and_kwargs(test_mod.rand_sleep, _args)
 
 
-def test_connect_master_salt_client_error(minion_opts):
+async def test_connect_master_salt_client_error(minion_opts, connect_master_mock):
     """
     Ensure minion's destory method is called on an salt client error while connecting to master.
     """
+    minion_opts["acceptance_wait_time"] = 0
     mm = salt.minion.MinionManager(minion_opts)
     minion = salt.minion.Minion(minion_opts)
-    minion.connect_master = MagicMock(side_effect=SaltClientError)
+
+    connect_master_mock.exc = SaltClientError
+    minion.connect_master = connect_master_mock
     minion.destroy = MagicMock()
-    mm._connect_minion(minion)
+    await mm._connect_minion(minion)
     minion.destroy.assert_called_once()
 
+    # The first call raised an error which caused minion.destroy to get called,
+    # the second call is a success.
+    assert minion.connect_master.calls == 2
+
 
-def test_connect_master_unresolveable_error(minion_opts):
+async def test_connect_master_unresolveable_error(minion_opts, connect_master_mock):
     """
     Ensure minion's destory method is called on an unresolvable while connecting to master.
     """
     mm = salt.minion.MinionManager(minion_opts)
     minion = salt.minion.Minion(minion_opts)
-    minion.connect_master = MagicMock(side_effect=SaltMasterUnresolvableError)
+    connect_master_mock.exc = SaltMasterUnresolvableError
+    minion.connect_master = connect_master_mock
     minion.destroy = MagicMock()
     mm._connect_minion(minion)
     minion.destroy.assert_called_once()
 
+    # Unresolvable errors break out of the loop.
+    assert minion.connect_master.calls == 1
 
-def test_connect_master_general_exception_error(minion_opts):
+
+async def test_connect_master_general_exception_error(minion_opts, connect_master_mock):
     """
     Ensure minion's destory method is called on an un-handled exception while connecting to master.
     """
     mm = salt.minion.MinionManager(minion_opts)
     minion = salt.minion.Minion(minion_opts)
-    minion.connect_master = MagicMock(side_effect=Exception)
+    connect_master_mock.exc = Exception
+    minion.connect_master = connect_master_mock
     minion.destroy = MagicMock()
     mm._connect_minion(minion)
     minion.destroy.assert_called_once()
+
+    # The first call raised an error which caused minion.destroy to get called,
+    # the second call is a success.
+    assert minion.connect_master.calls == 2