Merge pull request #894 from TikhomirovSergey/master

 #875 FIX

Author: TikhomirovSergey

build.gradle

@@ -50,7 +50,7 @@ compileJava {
     ]
 }
 
-ext.seleniumVersion = '3.11.0'
+ext.seleniumVersion = '3.12.0'
 
 dependencies {
     compile ("org.seleniumhq.selenium:selenium-java:${seleniumVersion}") {

src/main/java/io/appium/java_client/android/nativekey/AndroidKey.java

@@ -1411,6 +1411,8 @@ public boolean isSystemKey() {
     }
 
     /**
+     * Is it wake key or not.
+     *
      * @return true if this a wakeup key.
      */
     public boolean isWakeKey() {

src/main/java/io/appium/java_client/events/DefaultAspect.java

@@ -24,6 +24,7 @@
 import org.aspectj.lang.JoinPoint;
 import org.aspectj.lang.ProceedingJoinPoint;
 import org.aspectj.lang.annotation.After;
+import org.aspectj.lang.annotation.AfterReturning;
 import org.aspectj.lang.annotation.Around;
 import org.aspectj.lang.annotation.Aspect;
 import org.aspectj.lang.annotation.Before;
@@ -38,6 +39,7 @@
 import java.util.Collection;
 import java.util.List;
 
+@SuppressWarnings("unused")
 @Aspect
 class DefaultAspect {
 
@@ -85,6 +87,8 @@ class DefaultAspect {
         + "context(..))";
     private static final String EXECUTION_SWITCH_TO_WINDOW = "execution(* org.openqa.selenium.WebDriver.TargetLocator"
             + ".window(..))";
+    private static final String EXECUTION_TAKE_SCREENSHOT_AS = "execution(* org.openqa.selenium.TakesScreenshot"
+            + ".getScreenshotAs(..))";
     private static final String AROUND = "execution(* org.openqa.selenium.WebDriver.*(..)) || "
         + "execution(* org.openqa.selenium.WebElement.*(..)) || "
         + "execution(* org.openqa.selenium.WebDriver.Navigation.*(..)) || "
@@ -114,7 +118,8 @@ class DefaultAspect {
         + "execution(* io.appium.java_client.MobileElement.*(..)) || "
         + "execution(* org.openqa.selenium.remote.RemoteWebDriver.*(..)) || "
         + "execution(* org.openqa.selenium.remote.RemoteWebElement.*(..)) || "
-        + "execution(* org.openqa.selenium.Alert.*(..))";
+        + "execution(* org.openqa.selenium.Alert.*(..)) || "
+        + "execution(* org.openqa.selenium.TakesScreenshot.*(..))";
 
     private final AbstractApplicationContext context;
     private final WebDriver driver;
@@ -483,6 +488,24 @@ public void afterSwitchToWindow(JoinPoint joinPoint) throws Throwable {
         }
     }
 
+    @Before(EXECUTION_TAKE_SCREENSHOT_AS)
+    public void beforeTakeScreenShot(JoinPoint joinPoint) throws Throwable {
+        try {
+            listener.beforeGetScreenshotAs(castArgument(joinPoint, 0));
+        } catch (Throwable t) {
+            throw getRootCause(t);
+        }
+    }
+
+    @AfterReturning(value = EXECUTION_TAKE_SCREENSHOT_AS, returning = "result")
+    public void afterTakeScreenShot(JoinPoint joinPoint, Object result) throws Throwable {
+        try {
+            listener.afterGetScreenshotAs(castArgument(joinPoint, 0), result);
+        } catch (Throwable t) {
+            throw getRootCause(t);
+        }
+    }
+
     @Before(EXECUTION_ROTATE)
     public void beforeRotation(JoinPoint joinPoint) throws Throwable {
         try {

src/main/java/io/appium/java_client/events/DefaultListener.java

@@ -30,6 +30,7 @@
 import org.openqa.selenium.Alert;
 import org.openqa.selenium.By;
 import org.openqa.selenium.Dimension;
+import org.openqa.selenium.OutputType;
 import org.openqa.selenium.Point;
 import org.openqa.selenium.ScreenOrientation;
 import org.openqa.selenium.WebDriver;
@@ -133,6 +134,16 @@ class DefaultListener
         ((ListensToException) dispatcher).onException(throwable, driver);
     }
 
+    @Override
+    public <X> void beforeGetScreenshotAs(OutputType<X> target) {
+        ((WebDriverEventListener) dispatcher).beforeGetScreenshotAs(target);
+    }
+
+    @Override
+    public <X> void afterGetScreenshotAs(OutputType<X> target, X screenshot) {
+        ((WebDriverEventListener) dispatcher).afterGetScreenshotAs(target, screenshot);
+    }
+
     public void add(Collection<Listener> listeners) {
         this.listeners.addAll(listeners);
     }

src/main/java/io/appium/java_client/remote/AppiumCommandExecutor.java

@@ -18,25 +18,45 @@
 
 import static com.google.common.base.Preconditions.checkNotNull;
 import static com.google.common.base.Throwables.throwIfUnchecked;
+import static java.lang.String.format;
+import static java.nio.charset.StandardCharsets.UTF_8;
 import static java.util.Optional.ofNullable;
+import static java.util.logging.Logger.getLogger;
+import static org.openqa.selenium.remote.DriverCommand.NEW_SESSION;
 
 import com.google.common.base.Supplier;
 import com.google.common.base.Throwables;
 
+import com.google.common.io.CountingOutputStream;
+import com.google.common.io.FileBackedOutputStream;
+
+import org.openqa.selenium.Capabilities;
+import org.openqa.selenium.ImmutableCapabilities;
+import org.openqa.selenium.SessionNotCreatedException;
 import org.openqa.selenium.WebDriverException;
 import org.openqa.selenium.remote.Command;
 import org.openqa.selenium.remote.CommandCodec;
 import org.openqa.selenium.remote.CommandInfo;
+import org.openqa.selenium.remote.Dialect;
 import org.openqa.selenium.remote.DriverCommand;
 import org.openqa.selenium.remote.HttpCommandExecutor;
+import org.openqa.selenium.remote.ProtocolHandshake;
 import org.openqa.selenium.remote.Response;
+import org.openqa.selenium.remote.ResponseCodec;
 import org.openqa.selenium.remote.http.HttpClient;
 import org.openqa.selenium.remote.http.HttpRequest;
+import org.openqa.selenium.remote.http.HttpResponse;
 import org.openqa.selenium.remote.http.W3CHttpCommandCodec;
 import org.openqa.selenium.remote.service.DriverService;
 
+import java.io.BufferedInputStream;
 import java.io.IOException;
+import java.io.InputStream;
+import java.io.OutputStreamWriter;
+import java.io.Writer;
 import java.lang.reflect.Field;
+import java.lang.reflect.InvocationTargetException;
+import java.lang.reflect.Method;
 import java.net.ConnectException;
 import java.net.URL;
 import java.util.Map;
@@ -111,6 +131,76 @@ private void setCommandCodec(CommandCodec<HttpRequest> newCodec) {
         setPrivateFieldValue("commandCodec", newCodec);
     }
 
+    private void setResponseCodec(ResponseCodec<HttpResponse> codec) {
+        setPrivateFieldValue("responseCodec", codec);
+    }
+
+    private HttpClient getClient() {
+        //noinspection unchecked
+        return getPrivateFieldValue("client", HttpClient.class);
+    }
+
+    private Response createSession(Command command) throws IOException {
+        if (getCommandCodec() != null) {
+            throw new SessionNotCreatedException("Session already exists");
+        }
+        ProtocolHandshake handshake = new ProtocolHandshake() {
+            @SuppressWarnings("unchecked")
+            public Result createSession(HttpClient client, Command command)
+                    throws IOException {
+                Capabilities desiredCapabilities = (Capabilities) command.getParameters().get("desiredCapabilities");
+                Capabilities desired = desiredCapabilities == null ? new ImmutableCapabilities() : desiredCapabilities;
+
+                //the number of bytes before the stream should switch to buffering to a file
+                int threshold = (int) Math.min(Runtime.getRuntime().freeMemory() / 10, Integer.MAX_VALUE);
+                FileBackedOutputStream os = new FileBackedOutputStream(threshold);
+                try {
+
+                    CountingOutputStream counter = new CountingOutputStream(os);
+                    Writer writer = new OutputStreamWriter(counter, UTF_8);
+                    NewAppiumSessionPayload payload = NewAppiumSessionPayload.create(desired);
+                    payload.writeTo(writer);
+
+                    try (InputStream rawIn = os.asByteSource().openBufferedStream();
+                         BufferedInputStream contentStream = new BufferedInputStream(rawIn)) {
+
+                        Method createSessionMethod = this.getClass().getSuperclass()
+                                .getDeclaredMethod("createSession", HttpClient.class, InputStream.class, long.class);
+                        createSessionMethod.setAccessible(true);
+
+                        Optional<Result> result = (Optional<Result>) createSessionMethod
+                                .invoke(this, client, contentStream, counter.getCount());
+
+                        return result.map(result1 -> {
+                            Result toReturn = result.get();
+                            getLogger(ProtocolHandshake.class.getName())
+                                    .info(format("Detected dialect: %s", toReturn.getDialect()));
+                            return toReturn;
+                        }).orElseThrow(() -> new SessionNotCreatedException(
+                                format("Unable to create new remote session. desired capabilities = %s", desired)));
+                    } catch (NoSuchMethodException | IllegalAccessException | InvocationTargetException e) {
+                        throw new WebDriverException(format("It is impossible to create a new session "
+                                        + "because 'createSession' which takes %s, %s and %s was not found "
+                                        + "or it is not accessible",
+                                HttpClient.class.getSimpleName(),
+                                InputStream.class.getSimpleName(),
+                                long.class.getSimpleName()), e);
+                    }
+                } finally {
+                    os.reset();
+                }
+            }
+        };
+
+        ProtocolHandshake.Result result = handshake
+                .createSession(getClient(), command);
+        Dialect dialect = result.getDialect();
+        setCommandCodec(dialect.getCommandCodec());
+        getAdditionalCommands().forEach(this::defineCommand);
+        setResponseCodec(dialect.getResponseCodec());
+        return result.createResponse();
+    }
+
     @Override
     public Response execute(Command command) throws WebDriverException {
         if (DriverCommand.NEW_SESSION.equals(command.getName())) {
@@ -125,7 +215,7 @@ public Response execute(Command command) throws WebDriverException {
 
         Response response;
         try {
-            response = super.execute(command);
+            response = NEW_SESSION.equals(command.getName()) ? createSession(command) : super.execute(command);
         } catch (Throwable t) {
             Throwable rootCause = Throwables.getRootCause(t);
             if (rootCause instanceof ConnectException