From 3cdbcb2fca79a43124dd4edcb06fa33486d74670 Mon Sep 17 00:00:00 2001
From: Andrei Kamarouski <akamarouski@solvd.com>
Date: Sat, 9 Mar 2024 00:42:23 +0100
Subject: [PATCH] fix

---
 .../agent/RelaySessionFactoryInterceptor.java  |  3 +++
 .../grid/agent/SessionSlotInterceptor.java     | 18 +++---------------
 2 files changed, 6 insertions(+), 15 deletions(-)

diff --git a/agent/src/main/java/com/zebrunner/mcloud/grid/agent/RelaySessionFactoryInterceptor.java b/agent/src/main/java/com/zebrunner/mcloud/grid/agent/RelaySessionFactoryInterceptor.java
index b59b917..7426818 100644
--- a/agent/src/main/java/com/zebrunner/mcloud/grid/agent/RelaySessionFactoryInterceptor.java
+++ b/agent/src/main/java/com/zebrunner/mcloud/grid/agent/RelaySessionFactoryInterceptor.java
@@ -21,6 +21,8 @@
 import java.util.concurrent.TimeUnit;
 import java.util.logging.Logger;
 
+import static com.zebrunner.mcloud.grid.agent.SessionSlotInterceptor.DISCONNECT;
+
 public class RelaySessionFactoryInterceptor {
 
     private static final Logger LOGGER = Logger.getLogger(RelaySessionFactoryInterceptor.class.getName());
@@ -83,6 +85,7 @@ public static Object onTestMethodInvocation(@This final RelaySessionFactory fact
                         LOGGER.warning(() -> String.format("[STF] STF device busy by %s or not present/ready.", device.getOwner().getName()));
                         return false;
                     } else if(!StringUtils.equals(stfToken, DEFAULT_STF_TOKEN)) {
+                        DISCONNECT.set(false);
                         LOGGER.info(() -> String.format("[STF] STF device manually reserved by the same user: %s.", device.getOwner().getName()));
                         reserve = false;
                     }
diff --git a/agent/src/main/java/com/zebrunner/mcloud/grid/agent/SessionSlotInterceptor.java b/agent/src/main/java/com/zebrunner/mcloud/grid/agent/SessionSlotInterceptor.java
index a7780ed..c14ff16 100644
--- a/agent/src/main/java/com/zebrunner/mcloud/grid/agent/SessionSlotInterceptor.java
+++ b/agent/src/main/java/com/zebrunner/mcloud/grid/agent/SessionSlotInterceptor.java
@@ -10,6 +10,7 @@
 import org.openqa.selenium.NoSuchSessionException;
 import org.openqa.selenium.grid.node.local.SessionSlot;
 
+import java.util.concurrent.atomic.AtomicReference;
 import java.util.logging.Logger;
 
 public class SessionSlotInterceptor {
@@ -18,26 +19,13 @@ public class SessionSlotInterceptor {
     private static final String DEFAULT_STF_TOKEN = System.getenv("STF_TOKEN");
     private static final boolean STF_ENABLED = (!StringUtils.isEmpty(STF_URL) && !StringUtils.isEmpty(DEFAULT_STF_TOKEN));
     private static final String UDID = System.getenv("DEVICE_UDID");
+    static final AtomicReference<Boolean> DISCONNECT = new AtomicReference<>(true);
 
     @RuntimeType
     public static void onTestMethodInvocation(@This final SessionSlot slot, @SuperCall final Runnable proxy) throws Exception {
         if (STF_ENABLED) {
             try {
-                boolean disconnect = true;
-
-                try {
-                    String stfToken = CapabilityUtils.getZebrunnerCapability(slot.getSession().getCapabilities(), "STF_TOKEN")
-                            .map(String::valueOf)
-                            .orElse(DEFAULT_STF_TOKEN);
-                    if (!StringUtils.equalsIgnoreCase(stfToken, DEFAULT_STF_TOKEN)) {
-                        LOGGER.info(() -> String.format("STF device manually reserved by the same user, so it will not be disconnected."));
-                        disconnect = false;
-                    }
-                } catch (NoSuchSessionException e) {
-                    //do nothing
-                }
-
-                if (disconnect) {
+                if (DISCONNECT.getAndSet(true)) {
                     LOGGER.info(() -> "[STF] Return STF Device.");
                     if (HttpClient.uri(Path.STF_USER_DEVICES_BY_ID_PATH, STF_URL, UDID)
                             .withAuthorization(buildAuthToken(DEFAULT_STF_TOKEN))