From 2e6d3aea6589163d3034acc1e287ffbb77a514e4 Mon Sep 17 00:00:00 2001
From: lasotap <Patryk.Lasota@ig.com>
Date: Mon, 9 May 2022 13:53:24 +0200
Subject: [PATCH] fix multiple resetting + add and adjust tests

---
 locust/runners.py           |  2 +-
 locust/test/test_runners.py | 27 +++++++++++++++++++++++----
 2 files changed, 24 insertions(+), 5 deletions(-)

diff --git a/locust/runners.py b/locust/runners.py
index ac01dbbc59..fc61c80f26 100644
--- a/locust/runners.py
+++ b/locust/runners.py
@@ -887,6 +887,7 @@ def reset_connection(self):
         try:
             self.server.close()
             self.server = rpc.Server(self.master_bind_host, self.master_bind_port)
+            self.connection_broken = False
         except RPCError as e:
             logger.error(f"Temporary failure when resetting connection: {e}, will retry later.")
 
@@ -905,7 +906,6 @@ def client_listener(self):
                 self.connection_broken = True
                 gevent.sleep(FALLBACK_INTERVAL)
                 continue
-            self.connection_broken = False
             msg.node_id = client_id
             if msg.type == "client_ready":
                 if not msg.data:
diff --git a/locust/test/test_runners.py b/locust/test/test_runners.py
index 031f1bec6c..4e27e74307 100644
--- a/locust/test/test_runners.py
+++ b/locust/test/test_runners.py
@@ -51,10 +51,11 @@
 NETWORK_BROKEN = "network broken"
 
 
-def mocked_rpc():
+def mocked_rpc(raise_on_close=True):
     class MockedRpcServerClient:
         queue = Queue()
         outbox = []
+        raise_error_on_close = raise_on_close
 
         def __init__(self, *args, **kwargs):
             pass
@@ -85,7 +86,10 @@ def recv_from_client(self):
             return msg.node_id, msg
 
         def close(self):
-            raise RPCError()
+            if self.raise_error_on_close:
+                raise RPCError()
+            else:
+                pass
 
     return MockedRpcServerClient
 
@@ -2678,17 +2682,32 @@ class MyUser(User):
     def test_master_reset_connection(self):
         """Test that connection will be reset when network issues found"""
         with mock.patch("locust.runners.FALLBACK_INTERVAL", new=0.1):
-            with mock.patch("locust.rpc.rpc.Server", mocked_rpc()) as server:
+            with mock.patch("locust.rpc.rpc.Server", mocked_rpc(raise_on_close=False)) as server:
                 master = self.get_runner()
                 self.assertEqual(0, len(master.clients))
                 server.mocked_send(Message("client_ready", NETWORK_BROKEN, "fake_client"))
                 self.assertTrue(master.connection_broken)
                 server.mocked_send(Message("client_ready", __version__, "fake_client"))
-                sleep(0.2)
+                sleep(1)
                 self.assertFalse(master.connection_broken)
                 self.assertEqual(1, len(master.clients))
                 master.quit()
 
+    def test_reset_connection_after_RPCError(self):
+        with mock.patch("locust.rpc.rpc.Server", mocked_rpc(raise_on_close=False)) as server:
+            master = self.get_runner()
+            server.mocked_send(Message("client_ready", __version__, "fake_client"))
+            sleep(0.2)
+            self.assertFalse(master.connection_broken)
+            self.assertEqual(1, len(master.clients))
+
+            # Trigger RPCError
+            server.mocked_send(Message("lets_trigger_RPCError", NETWORK_BROKEN, "fake_client"))
+            self.assertTrue(master.connection_broken)
+            sleep(1)
+            self.assertFalse(master.connection_broken)
+            master.quit()
+
     def test_attributes_populated_when_calling_start(self):
         class MyUser1(User):
             @task