diff --git a/src/main/java/com/zebrunner/agent/core/webdriver/StartSessionInterceptor.java b/src/main/java/com/zebrunner/agent/core/webdriver/StartSessionInterceptor.java index 89c4a66..8af3094 100644 --- a/src/main/java/com/zebrunner/agent/core/webdriver/StartSessionInterceptor.java +++ b/src/main/java/com/zebrunner/agent/core/webdriver/StartSessionInterceptor.java @@ -3,6 +3,9 @@ import com.zebrunner.agent.core.config.ConfigurationHolder; import com.zebrunner.agent.core.registrar.TestSessionRegistrar; import com.zebrunner.agent.core.registrar.descriptor.SessionStartDescriptor; +import lombok.Data; +import lombok.NoArgsConstructor; +import lombok.RequiredArgsConstructor; import lombok.extern.slf4j.Slf4j; import net.bytebuddy.implementation.bind.annotation.Argument; import net.bytebuddy.implementation.bind.annotation.RuntimeType; @@ -21,12 +24,18 @@ import java.net.URISyntaxException; import java.net.URL; import java.util.Arrays; +import java.util.function.Consumer; @Slf4j public class StartSessionInterceptor { private static final TestSessionRegistrar REGISTRAR = TestSessionRegistrar.getInstance(); private static final CapabilitiesCustomizerChain CAPABILITIES_CUSTOMIZER_CHAIN = CapabilitiesCustomizerChain.getInstance(); + private static final ThreadLocal<Consumer<SessionRegisterDescriptor>> SESSION_REGISTER_DESCRIPTOR_CONSUMER = new ThreadLocal<>(); + + public static void setSessionRegisterConsumer(Consumer<SessionRegisterDescriptor> consumer) { + SESSION_REGISTER_DESCRIPTOR_CONSUMER.set(consumer); + } @RuntimeType public static void onSessionStart(@This RemoteWebDriver driver, @@ -37,7 +46,9 @@ public static void onSessionStart(@This RemoteWebDriver driver, capabilities = customizeCapabilities(methodInvocationProxy, capabilities); } - SessionStartDescriptor startDescriptor = SessionStartDescriptor.initiatedWith(capabilities.asMap()); + SessionRegisterDescriptor sessionRegisterDescriptor = new SessionRegisterDescriptor(); + sessionRegisterDescriptor.setSessionStartDescriptor(SessionStartDescriptor.initiatedWith(capabilities.asMap())); + try { methodInvocationProxy.run(); @@ -61,14 +72,24 @@ public static void onSessionStart(@This RemoteWebDriver driver, driverCapabilities = (Capabilities) capabilitiesField.get(driver); } - startDescriptor.successfullyStartedWith(sessionId, driverCapabilities.asMap()); + sessionRegisterDescriptor.getSessionStartDescriptor() + .successfullyStartedWith(sessionId, driverCapabilities.asMap()); } catch (Exception e) { + sessionRegisterDescriptor.setException(e); StringWriter errorMessageStringWriter = new StringWriter(); e.printStackTrace(new PrintWriter(errorMessageStringWriter)); - startDescriptor.failedToStart(errorMessageStringWriter.toString()); + sessionRegisterDescriptor.getSessionStartDescriptor() + .failedToStart(errorMessageStringWriter.toString()); throw e; } finally { - REGISTRAR.registerStart(startDescriptor); + if (SESSION_REGISTER_DESCRIPTOR_CONSUMER.get() == null) { + REGISTRAR.registerStart(sessionRegisterDescriptor.getSessionStartDescriptor()); + } else { + Consumer<SessionRegisterDescriptor> consumer = SESSION_REGISTER_DESCRIPTOR_CONSUMER.get(); + SESSION_REGISTER_DESCRIPTOR_CONSUMER.remove(); + sessionRegisterDescriptor.setTestSessionRegistrar(REGISTRAR); + consumer.accept(sessionRegisterDescriptor); + } } } @@ -157,4 +178,12 @@ private static Capabilities customizeCapabilities(Runnable methodInvocationProxy return capabilities; } + @Data + @NoArgsConstructor + @RequiredArgsConstructor + public static class SessionRegisterDescriptor { + private TestSessionRegistrar testSessionRegistrar; + private SessionStartDescriptor sessionStartDescriptor; + private Exception exception; + } }