[java][bidi] Add response handler

Author: pujagani

java/src/org/openqa/selenium/remote/Network.java

@@ -19,10 +19,12 @@
 
 import java.net.URI;
 import java.util.function.Predicate;
+import java.util.function.Supplier;
 import java.util.function.UnaryOperator;
 import org.openqa.selenium.Beta;
 import org.openqa.selenium.UsernameAndPassword;
 import org.openqa.selenium.remote.http.HttpRequest;
+import org.openqa.selenium.remote.http.HttpResponse;
 
 @Beta
 public interface Network {
@@ -35,9 +37,15 @@ public interface Network {
 
   void clearAuthenticationHandlers();
 
-  long addRequestHandler(Predicate<URI> filter, UnaryOperator<HttpRequest> handler);
+  long addRequestHandler(Predicate<HttpRequest> filter, UnaryOperator<HttpRequest> handler);
 
   void removeRequestHandler(long id);
 
   void clearRequestHandlers();
+
+  long addResponseHandler(Predicate<HttpRequest> filter, Supplier<HttpResponse> handler);
+
+  void removeResponseHandler(long id);
+
+  void clearResponseHandlers();
 }

java/src/org/openqa/selenium/remote/RemoteNetwork.java

@@ -25,6 +25,7 @@
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.atomic.AtomicLong;
 import java.util.function.Predicate;
+import java.util.function.Supplier;
 import java.util.function.UnaryOperator;
 import org.openqa.selenium.Beta;
 import org.openqa.selenium.UsernameAndPassword;
@@ -36,10 +37,12 @@
 import org.openqa.selenium.bidi.network.ContinueRequestParameters;
 import org.openqa.selenium.bidi.network.Header;
 import org.openqa.selenium.bidi.network.InterceptPhase;
+import org.openqa.selenium.bidi.network.ProvideResponseParameters;
 import org.openqa.selenium.bidi.network.RequestData;
 import org.openqa.selenium.remote.http.Contents;
 import org.openqa.selenium.remote.http.HttpMethod;
 import org.openqa.selenium.remote.http.HttpRequest;
+import org.openqa.selenium.remote.http.HttpResponse;
 
 @Beta
 class RemoteNetwork implements Network {
@@ -51,6 +54,8 @@ class RemoteNetwork implements Network {
 
   private final Map<Long, RequestDetails> requestHandlers = new ConcurrentHashMap<>();
 
+  private final Map<Long, ResponseDetails> responseHandlers = new ConcurrentHashMap<>();
+
   private final AtomicLong callBackId = new AtomicLong(1);
 
   public RemoteNetwork(WebDriver driver) {
@@ -97,24 +102,60 @@ private void interceptRequest() {
           ContinueRequestParameters continueRequestParameters =
               new ContinueRequestParameters(requestId);
 
-          Optional<UnaryOperator<HttpRequest>> requestHandler = getRequestHandler(uri);
+          RequestData interceptedRequest = beforeRequestSent.getRequest();
+
+          // Build the originalRequest object from the intercepted request details.
+          HttpRequest originalRequest =
+              new HttpRequest(
+                  HttpMethod.getHttpMethod(interceptedRequest.getMethod()),
+                  interceptedRequest.getUrl());
+
+          // Populate the headers of the original request.
+          // Body is not available as part of WebDriver Spec, hence we cannot add that or use that.
+          interceptedRequest
+              .getHeaders()
+              .forEach(
+                  header ->
+                      originalRequest.addHeader(header.getName(), header.getValue().getValue()));
+
+          Optional<UnaryOperator<HttpRequest>> requestHandler = getRequestHandler(originalRequest);
+
+          Optional<Supplier<HttpResponse>> responseHandler = getResponseHandler(originalRequest);
+
+          // If request and response handler both are registered for same uri,
+          // then the response will be mocked instead of modifying the outgoing request.
+          // This can be altered in the future to let the modified request go through and have a
+          // response mock for that modified request.
+          // If in future the Browsers support intercepting at "response started" phase and allow
+          // using provide response command in that phase.
+          // Currently, only intercepting in "before request sent" phase is permitted.
+          if (responseHandler.isPresent()) {
+            ProvideResponseParameters provideResponseParameters =
+                new ProvideResponseParameters(requestId);
 
-          if (requestHandler.isPresent()) {
-            RequestData interceptedRequest = beforeRequestSent.getRequest();
+            HttpResponse modifiedResponse = responseHandler.get().get();
+
+            provideResponseParameters.statusCode(modifiedResponse.getStatus());
+
+            List<Header> headerList = new ArrayList<>();
+            modifiedResponse.forEachHeader(
+                (name, value) ->
+                    headerList.add(
+                        new Header(name, new BytesValue(BytesValue.Type.STRING, value))));
 
-            // Build the originalRequest object from the intercepted request details.
-            HttpRequest originalRequest =
-                new HttpRequest(
-                    HttpMethod.getHttpMethod(interceptedRequest.getMethod()),
-                    interceptedRequest.getUrl());
+            if (!headerList.isEmpty()) {
+              provideResponseParameters.headers(headerList);
+            }
 
-            // Populate the headers of the original request.
-            interceptedRequest
-                .getHeaders()
-                .forEach(
-                    header ->
-                        originalRequest.addHeader(header.getName(), header.getValue().getValue()));
+            Contents.Supplier content = modifiedResponse.getContent();
 
+            if (content.length() > 0) {
+              provideResponseParameters.body(
+                  new BytesValue(BytesValue.Type.STRING, Contents.utf8String(content)));
+            }
+            network.provideResponse(provideResponseParameters);
+            return;
+          } else if (requestHandler.isPresent()) {
             HttpRequest modifiedRequest = requestHandler.get().apply(originalRequest);
 
             continueRequestParameters.method(modifiedRequest.getMethod());
@@ -144,13 +185,20 @@ private void interceptRequest() {
         });
   }
 
-  private Optional<UnaryOperator<HttpRequest>> getRequestHandler(URI uri) {
+  private Optional<UnaryOperator<HttpRequest>> getRequestHandler(HttpRequest request) {
     return requestHandlers.values().stream()
-        .filter(requestDetails -> requestDetails.getFilter().test(uri))
+        .filter(requestDetails -> requestDetails.getFilter().test(request))
         .map(RequestDetails::getHandler)
         .findFirst();
   }
 
+  private Optional<Supplier<HttpResponse>> getResponseHandler(HttpRequest request) {
+    return responseHandlers.values().stream()
+        .filter(responseDetails -> responseDetails.getFilter().test(request))
+        .map(ResponseDetails::getHandler)
+        .findFirst();
+  }
+
   @Override
   public long addAuthenticationHandler(UsernameAndPassword usernameAndPassword) {
     return addAuthenticationHandler(url -> true, usernameAndPassword);
@@ -176,7 +224,7 @@ public void clearAuthenticationHandlers() {
   }
 
   @Override
-  public long addRequestHandler(Predicate<URI> filter, UnaryOperator<HttpRequest> handler) {
+  public long addRequestHandler(Predicate<HttpRequest> filter, UnaryOperator<HttpRequest> handler) {
     long id = this.callBackId.incrementAndGet();
 
     requestHandlers.put(id, new RequestDetails(filter, handler));
@@ -193,6 +241,25 @@ public void clearRequestHandlers() {
     requestHandlers.clear();
   }
 
+  // Allows mocking the response body
+  @Override
+  public long addResponseHandler(Predicate<HttpRequest> filter, Supplier<HttpResponse> handler) {
+    long id = this.callBackId.incrementAndGet();
+
+    responseHandlers.put(id, new ResponseDetails(filter, handler));
+    return id;
+  }
+
+  @Override
+  public void removeResponseHandler(long id) {
+    responseHandlers.remove(id);
+  }
+
+  @Override
+  public void clearResponseHandlers() {
+    responseHandlers.clear();
+  }
+
   private class AuthDetails {
     private final Predicate<URI> filter;
     private final UsernameAndPassword usernameAndPassword;
@@ -212,20 +279,38 @@ public UsernameAndPassword getUsernameAndPassword() {
   }
 
   private class RequestDetails {
-    private final Predicate<URI> filter;
+    private final Predicate<HttpRequest> filter;
     private final UnaryOperator<HttpRequest> handler;
 
-    public RequestDetails(Predicate<URI> filter, UnaryOperator<HttpRequest> handler) {
+    public RequestDetails(Predicate<HttpRequest> filter, UnaryOperator<HttpRequest> handler) {
       this.filter = filter;
       this.handler = handler;
     }
 
-    public Predicate<URI> getFilter() {
+    public Predicate<HttpRequest> getFilter() {
       return this.filter;
     }
 
     public UnaryOperator<HttpRequest> getHandler() {
       return this.handler;
     }
   }
+
+  private class ResponseDetails {
+    private final Predicate<HttpRequest> filter;
+    private final Supplier<HttpResponse> handler;
+
+    public ResponseDetails(Predicate<HttpRequest> filter, Supplier<HttpResponse> handler) {
+      this.filter = filter;
+      this.handler = handler;
+    }
+
+    public Predicate<HttpRequest> getFilter() {
+      return this.filter;
+    }
+
+    public Supplier<HttpResponse> getHandler() {
+      return this.handler;
+    }
+  }
 }