fix

zebrunner · Mar 8, 2024 · 3cdbcb2 · 3cdbcb2
1 parent 63f755f
commit 3cdbcb2
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 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
@@ -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
@@ -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))