feature: stf support (without manual reserve)

zebrunner · Mar 9, 2024 · 979e826 · 979e826
1 parent f92f503
commit 979e826
Show file tree

Hide file tree

Showing 31 changed files with 304 additions and 229 deletions.
diff --git a/agent/pom.xml b/agent/pom.xml
@@ -12,6 +12,11 @@
         <project.build.sourceEncoding>UTF-8</project.build.sourceEncoding>
         <selenium.version>4.18.1</selenium.version>
         <commons-lang3.version>3.14.0</commons-lang3.version>
+        <jackson-databind.version>2.16.0</jackson-databind.version>
+        <httpclient.version>5.2.2</httpclient.version>
+        <genson.version>1.6</genson.version>
+        <jersey-bundle.version>1.19.4</jersey-bundle.version>
+        <lombok.version>1.18.30</lombok.version>
         <maven-shade-plugin.version>3.5.0</maven-shade-plugin.version>
         <maven-compiler-plugin.version>3.11.0</maven-compiler-plugin.version>
     </properties>
@@ -38,6 +43,42 @@
             <artifactId>commons-lang3</artifactId>
             <version>${commons-lang3.version}</version>
         </dependency>
+        <dependency>
+            <groupId>org.projectlombok</groupId>
+            <artifactId>lombok</artifactId>
+            <version>${lombok.version}</version>
+            <scope>provided</scope>
+        </dependency>
+        <dependency>
+            <groupId>org.apache.commons</groupId>
+            <artifactId>commons-lang3</artifactId>
+            <version>${commons-lang3.version}</version>
+        </dependency>
+        <dependency>
+            <groupId>com.fasterxml.jackson.core</groupId>
+            <artifactId>jackson-databind</artifactId>
+            <version>${jackson-databind.version}</version>
+        </dependency>
+        <dependency>
+            <groupId>org.apache.httpcomponents.client5</groupId>
+            <artifactId>httpclient5</artifactId>
+            <version>${httpclient.version}</version>
+        </dependency>
+        <dependency>
+            <groupId>com.owlike</groupId>
+            <artifactId>genson</artifactId>
+            <version>${genson.version}</version>
+        </dependency>
+        <dependency>
+            <groupId>com.sun.jersey</groupId>
+            <artifactId>jersey-bundle</artifactId>
+            <version>${jersey-bundle.version}</version>
+        </dependency>
+        <dependency>
+            <groupId>jakarta.annotation</groupId>
+            <artifactId>jakarta.annotation-api</artifactId>
+            <version>3.0.0-M1</version>
+        </dependency>
     </dependencies>
     <build>
         <plugins>

diff --git a/agent/src/main/java/com/zebrunner/mcloud/grid/agent/MobileCapabilityMatcher.java b/agent/src/main/java/com/zebrunner/mcloud/grid/agent/MobileCapabilityMatcher.java
@@ -28,7 +28,7 @@
 import java.util.List;
 import java.util.logging.Logger;
 
-import static com.zebrunner.mcloud.grid.agent.utils.CapabilityUtils.getAppiumCapability;
+import static com.zebrunner.mcloud.grid.agent.util.CapabilityUtils.getAppiumCapability;
 
 @SuppressWarnings("unused")
 public final class MobileCapabilityMatcher extends DefaultSlotMatcher {

diff --git a/agent/src/main/java/com/zebrunner/mcloud/grid/agent/NodeAgent.java b/agent/src/main/java/com/zebrunner/mcloud/grid/agent/NodeAgent.java
@@ -20,14 +20,16 @@
 public class NodeAgent {
     private static final Logger LOGGER = Logger.getLogger(NodeAgent.class.getName());
     private static final String RELAY_SESSION_FACTORY_CLASS = "org.openqa.selenium.grid.node.relay.RelaySessionFactory";
-    private static final String TEST_METHOD_NAME = "test";
+    private static final String SESSION_SLOT_CLASS = "org.openqa.selenium.grid.node.local.SessionSlot";
 
     public static void premain(String args, Instrumentation instrumentation) {
         try {
             new AgentBuilder.Default()
                     .with(new AgentBuilder.InitializationStrategy.SelfInjection.Eager())
                     .type(named(RELAY_SESSION_FACTORY_CLASS))
                     .transform((builder, type, classloader, module, protectionDomain) -> addTestMethodInterceptor(builder))
+                    .type(named(SESSION_SLOT_CLASS))
+                    .transform((builder, type, classloader, module, protectionDomain) -> addSessionSlotMethodInterceptor(builder))
                     .installOn(instrumentation);
         } catch (Exception e) {
             LOGGER.warning(() -> "Could not init instrumentation.");
@@ -39,15 +41,46 @@ private static DynamicType.Builder<?> addTestMethodInterceptor(DynamicType.Build
                 .intercept(to(testMethodInterceptor()));
     }
 
+    private static DynamicType.Builder<?> addSessionSlotMethodInterceptor(DynamicType.Builder<?> builder) {
+        return builder.method(isReleaseMethod())
+                .intercept(to(releaseMethodInterceptor()))
+                .method(isReserveMethod())
+                .intercept(to(reserveMethodInterceptor()));
+    }
+
     public static ElementMatcher<? super MethodDescription> isTestMethod() {
         return isPublic()
                 .and(not(isStatic()))
-                .and(new NameMatcher<>(TEST_METHOD_NAME::equals));
+                .and(new NameMatcher<>("test"::equals));
     }
 
     private static TypeDescription testMethodInterceptor() {
         return TypePool.Default.ofSystemLoader()
                 .describe(RelaySessionFactoryInterceptor.class.getName())
                 .resolve();
     }
+
+    public static ElementMatcher<? super MethodDescription> isReleaseMethod() {
+        return isPublic()
+                .and(not(isStatic()))
+                .and(new NameMatcher<>("release"::equals));
+    }
+
+    private static TypeDescription releaseMethodInterceptor() {
+        return TypePool.Default.ofSystemLoader()
+                .describe(SessionSlotReleaseInterceptor.class.getName())
+                .resolve();
+    }
+
+    public static ElementMatcher<? super MethodDescription> isReserveMethod() {
+        return isPublic()
+                .and(not(isStatic()))
+                .and(new NameMatcher<>("reserve"::equals));
+    }
+
+    private static TypeDescription reserveMethodInterceptor() {
+        return TypePool.Default.ofSystemLoader()
+                .describe(SessionSlotReserveInterceptor.class.getName())
+                .resolve();
+    }
 }
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
@@ -17,4 +17,5 @@ public static Object onTestMethodInvocation(@This final RelaySessionFactory fact
             @SuperCall final Callable<Object> proxy, @Argument(0) Capabilities capabilities) throws Exception {
         return CAPABILITY_MATCHER.matches(factory.getStereotype(), capabilities);
     }
+
 }
diff --git a/agent/src/main/java/com/zebrunner/mcloud/grid/agent/SessionSlotReleaseInterceptor.java b/agent/src/main/java/com/zebrunner/mcloud/grid/agent/SessionSlotReleaseInterceptor.java
@@ -0,0 +1,44 @@
+package com.zebrunner.mcloud.grid.agent;
+
+import com.zebrunner.mcloud.grid.agent.stf.entity.Path;
+import com.zebrunner.mcloud.grid.agent.util.HttpClient;
+import net.bytebuddy.implementation.bind.annotation.RuntimeType;
+import net.bytebuddy.implementation.bind.annotation.SuperCall;
+import net.bytebuddy.implementation.bind.annotation.This;
+import org.apache.commons.lang3.StringUtils;
+import org.openqa.selenium.grid.node.local.SessionSlot;
+
+import java.util.concurrent.atomic.AtomicReference;
+import java.util.logging.Logger;
+
+public class SessionSlotReleaseInterceptor {
+    private static final Logger LOGGER = Logger.getLogger(SessionSlotReleaseInterceptor.class.getName());
+    private static final String STF_URL = System.getenv("STF_URL");
+    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 {
+                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))
+                            .delete(Void.class).getStatus() != 200) {
+                        LOGGER.warning(() -> "[STF] Could not return device to the STF.");
+                    }
+                }
+            } catch (Exception e) {
+                LOGGER.warning(() -> String.format("[STF] Could not return device to the STF. Error: %s", e.getMessage()));
+            }
+        }
+        proxy.run();
+    }
+
+    private static String buildAuthToken(String authToken) {
+        return "Bearer " + authToken;
+    }
+}
diff --git a/agent/src/main/java/com/zebrunner/mcloud/grid/agent/SessionSlotReserveInterceptor.java b/agent/src/main/java/com/zebrunner/mcloud/grid/agent/SessionSlotReserveInterceptor.java
@@ -0,0 +1,113 @@
+package com.zebrunner.mcloud.grid.agent;
+
+import com.zebrunner.mcloud.grid.agent.stf.entity.Devices;
+import com.zebrunner.mcloud.grid.agent.stf.entity.Path;
+import com.zebrunner.mcloud.grid.agent.stf.entity.STFDevice;
+import com.zebrunner.mcloud.grid.agent.stf.entity.User;
+import com.zebrunner.mcloud.grid.agent.util.HttpClient;
+import net.bytebuddy.implementation.bind.annotation.RuntimeType;
+import net.bytebuddy.implementation.bind.annotation.SuperCall;
+import net.bytebuddy.implementation.bind.annotation.This;
+import org.apache.commons.lang3.StringUtils;
+import org.openqa.selenium.grid.node.local.SessionSlot;
+
+import java.util.HashMap;
+import java.util.Map;
+import java.util.Optional;
+import java.util.concurrent.TimeUnit;
+import java.util.logging.Logger;
+
+import static com.zebrunner.mcloud.grid.agent.SessionSlotReleaseInterceptor.DISCONNECT;
+
+public class SessionSlotReserveInterceptor {
+
+    private static final Logger LOGGER = Logger.getLogger(SessionSlotReserveInterceptor.class.getName());
+    private static final String NOT_AUTHENTICATED_ERROR = "[STF] Not authenticated at STF successfully! URL: '%s'; Token: '%s'";
+    private static final String UNABLE_GET_DEVICES_STATUS_ERROR = "[STF] Unable to get devices status. HTTP status: %s";
+    private static final String COULD_NOT_FIND_DEVICE_ERROR = "[STF] Could not find STF device with udid: %s";
+    private static final String STF_URL = System.getenv("STF_URL");
+    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));
+    // Max time is seconds for reserving devices in STF
+    private static final String DEFAULT_STF_TIMEOUT = Optional.ofNullable(System.getenv("STF_TIMEOUT"))
+            .filter(StringUtils::isNotBlank)
+            .orElse("3600");
+    private static final String UDID = System.getenv("DEVICE_UDID");
+
+    @RuntimeType
+    public static void onTestMethodInvocation(@This final SessionSlot slot, @SuperCall final Runnable proxy) throws Exception {
+        try {
+            if (STF_ENABLED) {
+                //                String stfToken = CapabilityUtils.getZebrunnerCapability(slot, "STF_TOKEN")
+                //                        .map(String::valueOf)
+                //                        .orElse(DEFAULT_STF_TOKEN);
+                String stfToken = DEFAULT_STF_TOKEN;
+
+                HttpClient.Response<User> user = HttpClient.uri(Path.STF_USER_PATH, STF_URL)
+                        .withAuthorization(buildAuthToken(stfToken))
+                        .get(User.class);
+                if (user.getStatus() != 200) {
+                    LOGGER.warning(() -> String.format(NOT_AUTHENTICATED_ERROR, STF_URL, stfToken));
+                    return;
+                }
+                HttpClient.Response<Devices> devices = HttpClient.uri(Path.STF_DEVICES_PATH, STF_URL)
+                        .withAuthorization(buildAuthToken(stfToken))
+                        .get(Devices.class);
+
+                if (devices.getStatus() != 200) {
+                    LOGGER.warning(() -> String.format(UNABLE_GET_DEVICES_STATUS_ERROR, devices.getStatus()));
+                    return;
+                }
+
+                Optional<STFDevice> optionalSTFDevice = devices.getObject()
+                        .getDevices()
+                        .stream()
+                        .filter(device -> StringUtils.equals(device.getSerial(), UDID))
+                        .findFirst();
+                if (optionalSTFDevice.isEmpty()) {
+                    LOGGER.warning(() -> String.format(COULD_NOT_FIND_DEVICE_ERROR, UDID));
+                    return;
+                }
+
+                STFDevice device = optionalSTFDevice.get();
+                LOGGER.info(() -> String.format("[STF] STF device info: %s", device));
+
+                boolean reserve = true;
+
+                if (device.getOwner() != null) {
+                    if (!(StringUtils.equals(device.getOwner().getName(), user.getObject().getUser().getName()) &&
+                            device.getPresent() &&
+                            device.getReady())) {
+                        LOGGER.warning(() -> String.format("[STF] STF device busy by %s or not present/ready.", device.getOwner().getName()));
+                        return;
+                    } 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;
+                    }
+                }
+                if (reserve) {
+                    Map<String, Object> entity = new HashMap<>();
+                    entity.put("serial", UDID);
+                    entity.put("timeout",
+                            TimeUnit.SECONDS.toMillis(Integer.parseInt(DEFAULT_STF_TIMEOUT)));
+                    if (HttpClient.uri(Path.STF_USER_DEVICES_PATH, STF_URL)
+                            .withAuthorization(buildAuthToken(stfToken))
+                            .post(Void.class, entity).getStatus() != 200) {
+                        LOGGER.warning(() -> String.format("[STF] Could not reserve STF device with udid: %s.", UDID));
+                    } else {
+                        LOGGER.info(() -> "[STF] Device successfully reserved.");
+                    }
+                }
+            }
+        } catch (Exception e) {
+            LOGGER.warning(() -> String.format("[STF] Could not reserve STF device. Error: %s", e.getMessage()));
+        } finally {
+            proxy.run();
+        }
+    }
+
+    private static String buildAuthToken(String authToken) {
+        return "Bearer " + authToken;
+    }
+}
diff --git a/...zebrunner/mcloud/grid/models/stf/App.java → ...ner/mcloud/grid/agent/stf/entity/App.java b/...zebrunner/mcloud/grid/models/stf/App.java → ...ner/mcloud/grid/agent/stf/entity/App.java
@@ -14,7 +14,7 @@
  * limitations under the License.
  *******************************************************************************/
 
-package com.zebrunner.mcloud.grid.models.stf;
+package com.zebrunner.mcloud.grid.agent.stf.entity;
 
 import com.fasterxml.jackson.annotation.JsonAnyGetter;
 import com.fasterxml.jackson.annotation.JsonAnySetter;

diff --git a/...unner/mcloud/grid/models/stf/Battery.java → ...mcloud/grid/agent/stf/entity/Battery.java b/...unner/mcloud/grid/models/stf/Battery.java → ...mcloud/grid/agent/stf/entity/Battery.java
@@ -14,20 +14,19 @@
  * limitations under the License.
  *******************************************************************************/
 
-package com.zebrunner.mcloud.grid.models.stf;
-
-import java.util.HashMap;
-import java.util.Map;
+package com.zebrunner.mcloud.grid.agent.stf.entity;
 
 import com.fasterxml.jackson.annotation.JsonAnyGetter;
 import com.fasterxml.jackson.annotation.JsonAnySetter;
 import com.fasterxml.jackson.annotation.JsonIgnore;
 import com.fasterxml.jackson.annotation.JsonInclude;
-
 import lombok.Getter;
 import lombok.NoArgsConstructor;
 import lombok.Setter;
 
+import java.util.HashMap;
+import java.util.Map;
+
 @Getter
 @Setter
 @NoArgsConstructor

diff --git a/...unner/mcloud/grid/models/stf/Browser.java → ...mcloud/grid/agent/stf/entity/Browser.java b/...unner/mcloud/grid/models/stf/Browser.java → ...mcloud/grid/agent/stf/entity/Browser.java
@@ -14,7 +14,7 @@
  * limitations under the License.
  *******************************************************************************/
 
-package com.zebrunner.mcloud.grid.models.stf;
+package com.zebrunner.mcloud.grid.agent.stf.entity;
 
 import com.fasterxml.jackson.annotation.JsonAnyGetter;
 import com.fasterxml.jackson.annotation.JsonAnySetter;

diff --git a/...runner/mcloud/grid/models/stf/Device.java → .../mcloud/grid/agent/stf/entity/Device.java b/...runner/mcloud/grid/models/stf/Device.java → .../mcloud/grid/agent/stf/entity/Device.java
@@ -14,20 +14,19 @@
  * limitations under the License.
  *******************************************************************************/
 
-package com.zebrunner.mcloud.grid.models.stf;
-
-import java.util.HashMap;
-import java.util.Map;
+package com.zebrunner.mcloud.grid.agent.stf.entity;
 
 import com.fasterxml.jackson.annotation.JsonAnyGetter;
 import com.fasterxml.jackson.annotation.JsonAnySetter;
 import com.fasterxml.jackson.annotation.JsonIgnore;
 import com.fasterxml.jackson.annotation.JsonInclude;
-
 import lombok.Getter;
 import lombok.NoArgsConstructor;
 import lombok.Setter;
 
+import java.util.HashMap;
+import java.util.Map;
+
 @Getter
 @Setter
 @NoArgsConstructor

diff --git a/...unner/mcloud/grid/models/stf/Devices.java → ...mcloud/grid/agent/stf/entity/Devices.java b/...unner/mcloud/grid/models/stf/Devices.java → ...mcloud/grid/agent/stf/entity/Devices.java
@@ -14,7 +14,7 @@
  * limitations under the License.
  *******************************************************************************/
 
-package com.zebrunner.mcloud.grid.models.stf;
+package com.zebrunner.mcloud.grid.agent.stf.entity;
 
 import com.fasterxml.jackson.annotation.JsonAnyGetter;
 import com.fasterxml.jackson.annotation.JsonAnySetter;

diff --git a/...unner/mcloud/grid/models/stf/Display.java → ...mcloud/grid/agent/stf/entity/Display.java b/...unner/mcloud/grid/models/stf/Display.java → ...mcloud/grid/agent/stf/entity/Display.java
@@ -14,7 +14,7 @@
  * limitations under the License.
  *******************************************************************************/
 
-package com.zebrunner.mcloud.grid.models.stf;
+package com.zebrunner.mcloud.grid.agent.stf.entity;
 
 import com.fasterxml.jackson.annotation.JsonAnyGetter;
 import com.fasterxml.jackson.annotation.JsonAnySetter;

diff --git a/...unner/mcloud/grid/models/stf/Network.java → ...mcloud/grid/agent/stf/entity/Network.java b/...unner/mcloud/grid/models/stf/Network.java → ...mcloud/grid/agent/stf/entity/Network.java
@@ -13,7 +13,7 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  *******************************************************************************/
-package com.zebrunner.mcloud.grid.models.stf;
+package com.zebrunner.mcloud.grid.agent.stf.entity;
 
 import com.fasterxml.jackson.annotation.JsonAnyGetter;
 import com.fasterxml.jackson.annotation.JsonAnySetter;

diff --git a/.../mcloud/grid/integration/client/Path.java → ...er/mcloud/grid/agent/stf/entity/Path.java b/.../mcloud/grid/integration/client/Path.java → ...er/mcloud/grid/agent/stf/entity/Path.java
@@ -13,7 +13,7 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  *******************************************************************************/
-package com.zebrunner.mcloud.grid.integration.client;
+package com.zebrunner.mcloud.grid.agent.stf.entity;
 
 public enum Path {