github.com/bazelbuild/rules_webtesting@v0.2.0/java/com/google/testing/web/screenshotter/Screenshotter.java (about)

     1  // Copyright 2017 Google Inc.
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  //      http://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    12  // See the License for the specific language governing permissions and
    13  // limitations under the License.
    14  //
    15  // //////////////////////////////////////////////////////////////////////////////
    16  //
    17  package com.google.testing.web.screenshotter;
    18  
    19  import com.google.common.collect.ImmutableSet;
    20  import java.io.IOException;
    21  import java.util.Collection;
    22  import java.util.Optional;
    23  import okhttp3.MediaType;
    24  import okhttp3.OkHttpClient;
    25  import okhttp3.Request;
    26  import okhttp3.RequestBody;
    27  import okhttp3.Response;
    28  import org.json.JSONArray;
    29  import org.json.JSONException;
    30  import org.json.JSONObject;
    31  import org.openqa.selenium.WebDriver;
    32  import org.openqa.selenium.WebElement;
    33  import org.openqa.selenium.remote.Dialect;
    34  import org.openqa.selenium.remote.RemoteWebDriver;
    35  import org.openqa.selenium.remote.RemoteWebElement;
    36  
    37  /**
    38   * API for interacting with advanced screenshot handler installed into Web Test Launcher.
    39   *
    40   * <p>Screenshotter objects are immutable; methods that return Screenshotter return an new
    41   * Screenshotter object that is a copy of the original with appropriate modifications.
    42   */
    43  public final class Screenshotter {
    44  
    45    private static final MediaType JSON = MediaType.parse("application/json; charset=utf-8");
    46  
    47    private static final OkHttpClient client = new OkHttpClient();
    48  
    49    private final String url;
    50    private final Optional<WebElement> element;
    51    private final ImmutableSet<WebElement> excluding;
    52  
    53    /**
    54     * Create a new Screenshotter that interacts with the WebDriver session to which driver is
    55     * connected.
    56     */
    57    public Screenshotter(WebDriver driver) {
    58      this(
    59          String.format(
    60              "%s/session/%s/google/screenshot",
    61              System.getenv("WEB_TEST_WEBDRIVER_SERVER"),
    62              String.valueOf(((RemoteWebDriver) driver).getSessionId())),
    63          Optional.empty(),
    64          ImmutableSet.of());
    65    }
    66  
    67    private Screenshotter(
    68        String url, Optional<WebElement> element, Collection<WebElement> excluding) {
    69      this.url = url;
    70      this.element = element;
    71      this.excluding = ImmutableSet.copyOf(excluding);
    72    }
    73  
    74    /** Returns a copy of this configured to take a screenshot of element. */
    75    public Screenshotter of(WebElement element) {
    76      return new Screenshotter(this.url, Optional.of(element), this.excluding);
    77    }
    78  
    79    /** Returns a copy of this configured to blackout elements. */
    80    public Screenshotter excluding(WebElement... elements) {
    81      ImmutableSet<WebElement> ex =
    82          new ImmutableSet.Builder<WebElement>().addAll(this.excluding).add(elements).build();
    83      return new Screenshotter(this.url, this.element, ex);
    84    }
    85  
    86    public Screenshot takeScreenshot() throws IOException, JSONException {
    87      Response response =
    88          client
    89              .newCall(
    90                  new Request.Builder()
    91                      .url(url)
    92                      .post(RequestBody.create(JSON, getRequestBodyJson().toString()))
    93                      .build())
    94              .execute();
    95      JSONObject responseBody = new JSONObject(response.body().string());
    96  
    97      if (!response.isSuccessful()) {
    98        throw new IOException(responseBody.getString("message"));
    99      }
   100      return new Screenshot(responseBody);
   101    }
   102  
   103    private JSONObject getRequestBodyJson() throws JSONException {
   104      JSONObject params = new JSONObject();
   105  
   106      if (element.isPresent()) {
   107        params.put("Element", webElementToJSON(element.get()));
   108      }
   109  
   110      if (!excluding.isEmpty()) {
   111        JSONArray ex = new JSONArray();
   112        for (WebElement el : excluding) {
   113          ex.put(webElementToJSON(el));
   114        }
   115        params.put("Exclude", ex);
   116      }
   117  
   118      return params;
   119    }
   120  
   121    private static JSONObject webElementToJSON(WebElement element) throws JSONException {
   122      String id = ((RemoteWebElement) element).getId();
   123      JSONObject object = new JSONObject();
   124      object.put(Dialect.OSS.getEncodedElementKey(), id);
   125      object.put(Dialect.W3C.getEncodedElementKey(), id);
   126      return object;
   127    }
   128  }