Support mutating responses in network interception

Author: p0deje

rb/lib/selenium/webdriver/common/driver_extensions/has_network_interception.rb

@@ -33,7 +33,7 @@ module HasNetworkInterception
         #     request.continue
         #   end
         #
-        # @example Stub response for image requests
+        # @example Stub requests for images
         #   driver.intercept do |request|
         #     if request.url.match?(/\.png$/)
         #       request.respond(body: File.read('myfile.png'))
@@ -42,25 +42,72 @@ module HasNetworkInterception
         #     end
         #   end
         #
+        # @example Log responses and pass through
+        #   driver.intercept do |request|
+        #     request.continue do |response|
+        #       puts "#{response.code} #{response.body}"
+        #       response.continue
+        #     end
+        #   end
+        #
+        # @example Mutate specific response
+        #   driver.intercept do |request|
+        #     request.continue do |response|
+        #       if request.url.include?('/myurl')
+        #         request.respond(body: "#{response.body}, Added by Selenium!")
+        #       else
+        #         response.continue
+        #       end
+        #     end
+        #   end
+        #
         # @param [#call] block which is called when request is interecepted
         # @yieldparam [DevTools::Request]
         #
 
-        def intercept
+        def intercept(&block)
           devtools.network.set_cache_disabled(cache_disabled: true)
           devtools.fetch.on(:request_paused) do |params|
-            request = DevTools::Request.new(
-              devtools: devtools,
-              id: params['requestId'],
-              url: params.dig('request', 'url'),
-              method: params.dig('request', 'method'),
-              headers: params.dig('request', 'headers')
-            )
-            yield request
+            id = params['requestId']
+            if params.key?('responseStatusCode') || params.key?('responseErrorReason')
+              intercept_response(id, params, &intercepted_requests[id].on_response)
+            else
+              intercept_request(id, params, &block)
+            end
           end
-          devtools.fetch.enable
+          devtools.fetch.enable(patterns: [{requestStage: 'Request'}, {requestStage: 'Response'}])
         end
 
+        private
+
+        def intercepted_requests
+          @intercepted_requests ||= {}
+        end
+
+        def intercept_request(id, params)
+          request = DevTools::Request.new(
+            devtools: devtools,
+            id: id,
+            url: params.dig('request', 'url'),
+            method: params.dig('request', 'method'),
+            headers: params.dig('request', 'headers')
+          )
+          intercepted_requests[id] = request
+
+          yield request
+        end
+
+        def intercept_response(id, params)
+          response = DevTools::Response.new(
+            devtools: devtools,
+            id: id,
+            code: params['responseStatusCode'],
+            headers: params['responseHeaders']
+          )
+          intercepted_requests.delete(id)
+
+          yield response
+        end
       end # HasNetworkInterception
     end # DriverExtensions
   end # WebDriver

rb/lib/selenium/webdriver/devtools.rb

@@ -25,6 +25,7 @@ class DevTools
       autoload :MutationEvent, 'selenium/webdriver/devtools/mutation_event'
       autoload :PinnedScript, 'selenium/webdriver/devtools/pinned_script'
       autoload :Request, 'selenium/webdriver/devtools/request'
+      autoload :Response, 'selenium/webdriver/devtools/response'
 
       def initialize(url:)
         @messages = []

rb/lib/selenium/webdriver/devtools/request.rb

@@ -24,18 +24,39 @@ class Request
 
         attr_reader :url, :method, :headers
 
+        # @api private
+        attr_reader :on_response
+
         def initialize(devtools:, id:, url:, method:, headers:)
           @devtools = devtools
           @id = id
           @url = url
           @method = method
           @headers = headers
+          @on_response = Proc.new(&:continue)
         end
 
-        def continue
+        #
+        # Continues the request, optionally yielding
+        # the response before it reaches the browser.
+        #
+        # @param [#call] block which is called when response is intercepted
+        # @yieldparam [DevTools::Response]
+        #
+
+        def continue(&block)
+          @on_response = block if block_given?
           @devtools.fetch.continue_request(request_id: @id)
         end
 
+        #
+        # Fulfills the request providing the stubbed response.
+        #
+        # @param [Integer] code
+        # @param [Hash] headers
+        # @param [String] body
+        #
+
         def respond(code: 200, headers: {}, body: '')
           @devtools.fetch.fulfill_request(
             request_id: @id,

rb/lib/selenium/webdriver/devtools/response.rb

@@ -0,0 +1,63 @@
+# frozen_string_literal: true
+
+# Licensed to the Software Freedom Conservancy (SFC) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The SFC licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+module Selenium
+  module WebDriver
+    class DevTools
+      class Response
+
+        attr_reader :code, :headers
+
+        def initialize(devtools:, id:, code:, headers:)
+          @devtools = devtools
+          @id = id
+          @code = code
+          @headers = headers
+        end
+
+        #
+        # Returns the response body.
+        # @return [String]
+        #
+
+        def body
+          @body ||= begin
+            result = @devtools.fetch.get_response_body(request_id: @id)
+            encoded_body = result.dig('result', 'body')
+
+            Base64.strict_decode64(encoded_body)
+          end
+        end
+
+        #
+        # Continues the response unmodified.
+        #
+
+        def continue
+          @devtools.fetch.continue_request(request_id: @id)
+        end
+
+        def inspect
+          %(#<#{self.class.name} @code="#{code}")
+        end
+
+      end # Response
+    end # DevTools
+  end # WebDriver
+end # Selenium

rb/spec/integration/selenium/webdriver/devtools_spec.rb

@@ -160,15 +160,12 @@ module WebDriver
           expect(requests).not_to be_empty
         end
 
-        it 'allows to stub responses' do
-          requests = []
+        it 'allows to stub requests' do
           driver.intercept do |request|
-            requests << request
             request.respond(body: '<title>Intercepted!</title>')
           end
           driver.navigate.to url_for('html5Page.html')
           expect(driver.title).to eq('Intercepted!')
-          expect(requests).not_to be_empty
         end
 
         it 'intercepts specific requests' do
@@ -193,6 +190,43 @@ module WebDriver
           expect(driver.title).to eq('Intercepted!')
           expect(stubbed).not_to be_empty
         end
+
+        it 'allows to continue responses' do
+          responses = []
+          driver.intercept do |request|
+            request.continue do |response|
+              responses << response
+              response.continue
+            end
+          end
+          driver.navigate.to url_for('html5Page.html')
+          expect(driver.title).to eq('HTML5')
+          expect(responses).not_to be_empty
+        end
+
+        it 'allows to stub responses' do
+          driver.intercept do |request|
+            request.continue do |response|
+              request.respond(body: "#{response.body}, '<h4 id=\"appended\">Appended!</h4>'")
+            end
+          end
+          driver.navigate.to url_for('html5Page.html')
+          expect(driver.find_elements(id: "appended")).not_to be_empty
+        end
+
+        it 'intercepts specific responses' do
+          driver.intercept do |request|
+            request.continue do |response|
+              if request.url.include?('html5Page.html')
+                request.respond(body: "#{response.body}, '<h4 id=\"appended\">Appended!</h4>'")
+              else
+                response.continue
+              end
+            end
+          end
+          driver.navigate.to url_for('html5Page.html')
+          expect(driver.find_elements(id: "appended")).not_to be_empty
+        end
       end
 
       context 'script pinning' do