Set cookie to track session requests and errors when sharing server b…

lib/capybara.rb

@@ -23,7 +23,7 @@ class << self
     attr_accessor :server_port, :exact, :match, :exact_options, :visible_text_only
     attr_accessor :default_selector, :default_max_wait_time, :ignore_hidden_elements
     attr_accessor :save_and_open_page_path, :wait_on_first_by_default, :automatic_reload
-    attr_accessor :reuse_server, :raise_server_errors, :server_errors
+    attr_accessor :reuse_server, :raise_server_errors, :server_errors, :cookie_tracking
     attr_writer :default_driver, :current_driver, :javascript_driver, :session_name, :server_host
     attr_accessor :app
 
@@ -51,6 +51,7 @@ class << self
     # [save_and_open_page_path = String]  Where to put pages saved through save_and_open_page (Default: Dir.pwd)
     # [wait_on_first_by_default = Boolean]   Whether Node#first defaults to Capybara waiting behavior for at least 1 element to match (Default: false)
     # [reuse_server = Boolean]  Reuse the server thread between multiple sessions using the same app object (Default: true)
+    # [cookie_tracking = Boolean]  Use a cookie to track the session a request is coming from (Default: false)
     # === DSL Options
     #
     # when using capybara/dsl, the following options are also available:
@@ -347,6 +348,11 @@ def reuse_server=(bool)
       @reuse_server = bool
     end
 
+    def cookie_tracking=(bool)
+      warn "Capybara.cookie_tracking == true is a BETA feature and may change in a future versions" if bool
+      @cookie_tracking = bool
+    end
+
     def deprecate(method, alternate_method, once=false)
       @deprecation_notified ||= {}
       warn "DEPRECATED: ##{method} is deprecated, please use ##{alternate_method} instead" unless once and @deprecation_notified[method]
@@ -420,6 +426,7 @@ module Selenium; end
   config.visible_text_only = false
   config.wait_on_first_by_default = false
   config.reuse_server = true
+  config.cookie_tracking = false
 end
 
 Capybara.register_driver :rack_test do |app|

lib/capybara/server.rb

@@ -5,6 +5,16 @@
 module Capybara
   class Server
     class Middleware
+      class Session
+        attr_accessor :error
+        attr_reader :counter, :id
+
+        def initialize(id)
+          @counter = Counter.new
+          @id = id
+        end
+      end
+
       class Counter
         attr_reader :value
 
@@ -22,29 +32,63 @@ def decrement
         end
       end
 
-      attr_accessor :error
-
       def initialize(app)
         @app = app
-        @counter = Counter.new
+        @sessions = Hash.new { |hash, key| hash[key] = Session.new(key) }
+      end
+
+      def session(session_id=nil)
+        @sessions[session_id && session_id.to_s]
+      end
+
+      def error(session_id)
+        session(session_id).error || session().error
+      end
+
+      def reset_error(session_id)
+        if session(session_id).error
+          session(session_id).error = nil
+        else
+          session().error = nil
+        end
+      end
+
+      def pending_requests?(session_id)
+        session(session_id).counter.value > 0 || session().counter.value > 0
       end
 
-      def pending_requests?
-        @counter.value > 0
+      def cookie_domain(hostname)
+        if (/\d{1,3}\.\d{1,3}\.\d{1,3}\.\d{1,3}$/.match(hostname))
+          {}
+        else
+          {domain: hostname.split('.')[-2..-1].join(".")}
+        end
       end
 
       def call(env)
         if env["PATH_INFO"] == "/__identify__"
-          [200, {}, [@app.object_id.to_s]]
+          [200, {}, ["#{@app.object_id.to_s}:#{self.object_id.to_s}"]]
         else
-          @counter.increment
+          request = Rack::Request.new(env)
+          session_id = request.cookies["capybara_session_id"] || request.params.delete("capybara_session_id")
+
+          current_session = session(session_id)
+          current_session.counter.increment
           begin
-            @app.call(env)
+            # @app.call(env)
+            status, headers, body = @app.call(env)
+            response = Rack::Response.new(body, status, headers)
+
+            cookie_options = { value: session_id, expires: Time.now + 24*60*60, path: '/'}
+            cookie_options.merge!(cookie_domain(request.host))
+            response.set_cookie("capybara_session_id", cookie_options ) if session_id
+
+            response.finish
           rescue *Capybara.server_errors => e
-            @error = e unless @error
+            current_session.error = e unless current_session.error
             raise e
           ensure
-            @counter.decrement
+            current_session.counter.decrement
           end
         end
       end
@@ -67,12 +111,12 @@ def initialize(app, port=Capybara.server_port, host=Capybara.server_host)
       @port ||= find_available_port
     end
 
-    def reset_error!
-      @middleware.error = nil
+    def reset_error!(session_id)
+      @middleware.reset_error(session_id)
     end
 
-    def error
-      @middleware.error
+    def error(session_id)
+      @middleware.error(session_id)
     end
 
     def responsive?
@@ -81,14 +125,20 @@ def responsive?
       res = Net::HTTP.start(host, @port) { |http| http.get('/__identify__') }
 
       if res.is_a?(Net::HTTPSuccess) or res.is_a?(Net::HTTPRedirection)
-        return res.body == @app.object_id.to_s
+        object_ids = res.body.split(':')
+        if object_ids[0] == @app.object_id.to_s
+          @middleware = ObjectSpace._id2ref(object_ids[1].to_i)
+          return true
+        else
+          return
+        end
       end
     rescue SystemCallError
       return false
     end
 
-    def wait_for_pending_requests
-      Timeout.timeout(60) { sleep(0.01) while @middleware.pending_requests? }
+    def wait_for_pending_requests(session_id)
+      Timeout.timeout(60) { sleep(0.01) while @middleware.pending_requests?(session_id) }
     rescue Timeout::Error
       raise "Requests did not finish in 60 seconds"
     end

lib/capybara/session.rb

@@ -1,4 +1,5 @@
 require 'capybara/session/matchers'
+require 'addressable'
 
 module Capybara
 
@@ -109,7 +110,7 @@ def reset!
         assert_no_selector :xpath, "/html/body/*" if driver.browser_initialized?
         @touched = false
       end
-      @server.wait_for_pending_requests if @server
+      @server.wait_for_pending_requests(self.object_id) if server
       raise_server_error!
     end
     alias_method :cleanup!, :reset!
@@ -120,9 +121,9 @@ def reset!
     # Raise errors encountered in the server
     #
     def raise_server_error!
-      raise @server.error if Capybara.raise_server_errors and @server and @server.error
+      raise @server.error(self.object_id) if Capybara.raise_server_errors and @server and @server.error(self.object_id)
     ensure
-      @server.reset_error! if @server
+      @server.reset_error!(self.object_id) if @server
     end
 
     ##
@@ -178,7 +179,10 @@ def current_host
     # @return [String] Fully qualified URL of the current page
     #
     def current_url
-      driver.current_url
+      uri = Addressable::URI.parse(driver.current_url)
+      new_query_values = (uri.query_values || {}).delete_if { |k,v| k == "capybara_session_id"}
+      uri.query_values = (new_query_values.empty? ? nil : new_query_values)
+      uri.to_s
     end
 
     ##
@@ -230,6 +234,11 @@ def visit(visit_uri)
 
       visit_uri = uri_base.merge(visit_uri) unless uri_base.nil?
 
+      if @server && Capybara.cookie_tracking
+        visit_uri = Addressable::URI.parse(visit_uri.to_s)
+        visit_uri.query_values=({'capybara_session_id' => self.object_id}.merge(visit_uri.query_values || {}))
+      end
+
       driver.visit(visit_uri.to_s)
     end
 

lib/capybara/spec/session/reset_session_spec.rb

@@ -71,6 +71,39 @@
     end
   end
 
+  context "When cookie_tracking is true" do
+    before do
+      @cookie_tracking = Capybara.cookie_tracking
+      Capybara.cookie_tracking = true
+    end
+
+    after do
+      Capybara.cookie_tracking = @cookie_tracking
+    end
+
+    it "raises any standard errors caught inside the server during a second session", requires: [:server] do
+      Capybara.using_driver(@session.mode) do
+        Capybara.using_session(:another_session) do
+          @another_session = Capybara.current_session
+          quietly { @another_session.visit("/error") }
+
+          expect do
+            @session.visit('/foo')
+          end.not_to raise_error
+
+          expect do
+            @another_session.reset_session!
+          end.to raise_error(TestApp::TestAppError)
+          @another_session.visit("/")
+          expect(@another_session.current_path).to eq("/")
+        end
+
+      end
+    end
+
+
+  end
+
   it "raises configured errors caught inside the server", :requires => [:server] do
     prev_errors = Capybara.server_errors
 

spec/server_spec.rb

@@ -113,9 +113,9 @@
       start_request(server2, 1.0)
 
       expect {
-        server1.wait_for_pending_requests
+        server1.wait_for_pending_requests(nil)
       }.to change{done}.from(false).to(true)
-      expect(server2.instance_variable_get('@middleware').pending_requests?).to eq(false)
+      expect(server2.instance_variable_get('@middleware').pending_requests?(nil)).to eq(false)
     end
 
   end
@@ -162,9 +162,9 @@
       start_request(server2, 1.0)
 
       expect {
-        server1.wait_for_pending_requests
+        server1.wait_for_pending_requests(nil)
       }.to change{done}.from(false).to(true)
-      expect(server2.instance_variable_get('@middleware').pending_requests?).to eq(true)
+      expect(server2.instance_variable_get('@middleware').pending_requests?(nil)).to eq(true)
     end
 
   end