From a01905b5cdbb248f8fe28e3283083ff7d0f66361 Mon Sep 17 00:00:00 2001
From: Vitaly Kh <82949155+bgklika@users.noreply.github.com>
Date: Tue, 29 Aug 2023 19:40:20 +0200
Subject: [PATCH] fix(uat): ease race between MQTT connect and agent reconnect
 (#399)

---
 .../client/control/implementation/EngineControlImpl.java    | 5 +++--
 .../java/com/aws/greengrass/steps/MqttControlSteps.java     | 6 +++---
 2 files changed, 6 insertions(+), 5 deletions(-)

diff --git a/uat/mqtt-client-control/src/main/java/com/aws/greengrass/testing/mqtt/client/control/implementation/EngineControlImpl.java b/uat/mqtt-client-control/src/main/java/com/aws/greengrass/testing/mqtt/client/control/implementation/EngineControlImpl.java
index 67e40e5fe..2bcefbfa3 100644
--- a/uat/mqtt-client-control/src/main/java/com/aws/greengrass/testing/mqtt/client/control/implementation/EngineControlImpl.java
+++ b/uat/mqtt-client-control/src/main/java/com/aws/greengrass/testing/mqtt/client/control/implementation/EngineControlImpl.java
@@ -178,9 +178,10 @@ public void onDiscoveryAgent(@NonNull String agentId, @NonNull String address, i
             } else {
                 logger.atInfo().log("Agent {} relocated to {}:{}", agentId, address, port);
                 agents.remove(agentId);
-                stopAgent(agentId, agentControl, false);
-                // recursion
+                // recursion just to add new relocated agent
                 onDiscoveryAgent(agentId, address, port);
+                // stop can take long time (about 2 seconds)
+                stopAgent(agentId, agentControl, false);
             }
         }
     }
diff --git a/uat/testing-features/src/main/java/com/aws/greengrass/steps/MqttControlSteps.java b/uat/testing-features/src/main/java/com/aws/greengrass/steps/MqttControlSteps.java
index 485327776..aca931c59 100644
--- a/uat/testing-features/src/main/java/com/aws/greengrass/steps/MqttControlSteps.java
+++ b/uat/testing-features/src/main/java/com/aws/greengrass/steps/MqttControlSteps.java
@@ -759,9 +759,6 @@ public void connect(String clientDeviceId, String componentId, String brokerId,
         }
 
 
-        // get agent control by componentId
-        AgentControl agentControl = getAgentControl(componentId);
-
         // request for new MQTT connection
         final String clientDeviceThingName = getClientDeviceThingName(clientDeviceId);
 
@@ -777,6 +774,9 @@ public void connect(String clientDeviceId, String componentId, String brokerId,
                      brokerId, host, port, clientDeviceThingName, componentId, mqttVersion);
             MqttConnectRequest request = buildMqttConnectRequest(
                     clientDeviceThingName, caList, host, port, version);
+            // get agent control by componentId as late as possible to reduce race probability
+            AgentControl agentControl = getAgentControl(componentId);
+
             try {
                 ConnectionControl connectionControl
                         = agentControl.createMqttConnection(request, connectionEvents);