From e8a392577d65203d21fb5074224b35c85d43f7f3 Mon Sep 17 00:00:00 2001
From: Pedro Algarvio <palgarvio@vmware.com>
Date: Thu, 17 Aug 2023 09:54:54 +0100
Subject: [PATCH] Fix merged forwarded code and test

Signed-off-by: Pedro Algarvio <palgarvio@vmware.com>
---
 salt/transport/zeromq.py                    |  7 ++++---
 tests/pytests/unit/transport/test_zeromq.py | 10 ++--------
 tests/support/mock.py                       |  5 +++--
 3 files changed, 9 insertions(+), 13 deletions(-)

diff --git a/salt/transport/zeromq.py b/salt/transport/zeromq.py
index d2105a61b88..373399c5230 100644
--- a/salt/transport/zeromq.py
+++ b/salt/transport/zeromq.py
@@ -543,14 +543,15 @@ class RequestServer(salt.transport.base.DaemonizedRequestServer):
                 await self._socket.send(self.encode_payload(reply))
             except asyncio.exceptions.TimeoutError:
                 continue
-            except salt.exceptions.SaltDeserializationError:
-                await self._socket.send(self.encode_payload({"msg": "bad load"}))
             except Exception as exc:  # pylint: disable=broad-except
                 log.error("Exception in request handler", exc_info=True)
                 break
 
     async def handle_message(self, stream, payload):
-        payload = self.decode_payload(payload)
+        try:
+            payload = self.decode_payload(payload)
+        except salt.exceptions.SaltDeserializationError:
+            return {"msg": "bad load"}
         return await self.message_handler(payload)
 
     def encode_payload(self, payload):
diff --git a/tests/pytests/unit/transport/test_zeromq.py b/tests/pytests/unit/transport/test_zeromq.py
index 1de55463baa..e8304eede00 100644
--- a/tests/pytests/unit/transport/test_zeromq.py
+++ b/tests/pytests/unit/transport/test_zeromq.py
@@ -3,7 +3,6 @@ import pytest
 
 import salt.config
 import salt.transport.zeromq
-from tests.support.mock import MagicMock
 
 
 async def test_req_server_garbage_request(io_loop):
@@ -22,14 +21,9 @@ async def test_req_server_garbage_request(io_loop):
     byts = msgpack.dumps({"foo": "bar"})
     badbyts = byts[:3] + b"^M" + byts[3:]
 
-    valid_response = msgpack.dumps("Invalid payload")
-
-    stream = MagicMock()
-    request_server.stream = stream
-
     try:
-        await request_server.handle_message(stream, badbyts)
+        ret = await request_server.handle_message(None, badbyts)
     except Exception as exc:  # pylint: disable=broad-except
         pytest.fail(f"Exception was raised {exc}")
 
-    request_server.stream.send.assert_called_once_with(valid_response)
+    assert ret == {"msg": "bad load"}
diff --git a/tests/support/mock.py b/tests/support/mock.py
index 2256ad8f5da..2f9970d4b04 100644
--- a/tests/support/mock.py
+++ b/tests/support/mock.py
@@ -27,6 +27,7 @@ from mock import (
     ANY,
     DEFAULT,
     FILTER_DIR,
+    AsyncMock,
     MagicMock,
     Mock,
     NonCallableMagicMock,
@@ -202,7 +203,7 @@ class MockFH:
                 )
             elif not self.binary_mode and content_type is not str:
                 raise TypeError(
-                    "write() argument must be str, not {}".format(content_type.__name__)
+                    f"write() argument must be str, not {content_type.__name__}"
                 )
 
     def _writelines(self, lines):
@@ -233,7 +234,7 @@ class MockCall:
                 ret = ret[:-2]
         else:
             for key, val in self.kwargs.items():
-                ret += "{}={}".format(salt.utils.stringutils.to_str(key), repr(val))
+                ret += f"{salt.utils.stringutils.to_str(key)}={repr(val)}"
         ret += ")"
         return ret