simplifying cookies plugin by having less cookie response introspection (now there's a client cookies store)

lib/httpx/plugins/cookies.rb

@@ -3,31 +3,81 @@
 module HTTPX
   module Plugins
     module Cookies
+      using URIExtensions
+
+      class Store
+        def initialize(cookies = nil)
+          @store = Hash.new { |hash, origin| hash[origin] = HTTP::CookieJar.new }
+          return unless cookies
+
+          cookies = cookies.split(/ *; */) if cookies.is_a?(String)
+          @default_cookies = cookies.map do |cookie, v|
+            if cookie.is_a?(HTTP::Cookie)
+              cookie
+            else
+              HTTP::Cookie.new(cookie.to_s, v.to_s)
+            end
+          end
+        end
+
+        def set(origin, cookies)
+          return unless cookies
+
+          @store[origin].parse(cookies, origin)
+        end
+
+        def [](uri)
+          store = @store[uri.origin]
+          @default_cookies.each do |cookie|
+            c = cookie.dup
+            c.domain ||= uri.authority
+            c.path ||= uri.path
+            store.add(c)
+          end if @default_cookies
+          store
+        end
+      end
+
       def self.load_dependencies(*)
         require "http/cookie"
       end
 
       module InstanceMethods
+        attr_reader :cookies_store
+
         def initialize(*)
-          @cookies_store = {}
           super
+          @cookies_store = @options.cookies || Store.new
         end
 
-        def cookies(cookies)
+        def with_cookies(cookies)
           branch(default_options.with_cookies(cookies))
         end
 
+        def wrap
+          return unless block_given?
+
+          super do |client|
+            old_cookies_store = @cookies_store
+            @cookies_store = old_cookies_store.dup
+            begin
+              yield client
+            ensure
+              @cookies_store = old_cookies_store
+            end
+          end
+        end
+
         private
 
         def on_response(request, response)
-          @cookies_store[request.origin] = response.cookies
+          @cookies_store.set(request.origin, response.headers["set-cookie"])
           super
         end
 
         def __build_req(*)
           request = super
-          request.headers.cookies(@cookies_store[request.origin], request)
-          request.headers.cookies(@options.cookies, request)
+          request.headers.cookies(@cookies_store[request.uri], request)
           request
         end
       end
@@ -36,40 +86,20 @@ module HTTPX
         def cookies(jar, request)
           return unless jar
 
-          unless jar.is_a?(HTTP::CookieJar)
-            jar = jar.each_with_object(HTTP::CookieJar.new) do |(cookie, v), j|
-              unless cookie.is_a?(HTTP::Cookie)
-                cookie = HTTP::Cookie.new(cookie.to_s, v.to_s)
-                cookie.domain = request.authority
-                cookie.path = request.path
-              end
-              j.add(cookie)
-            end
-          end
-          add("cookie", HTTP::Cookie.cookie_value(jar.cookies(request.uri)))
-        end
-      end
+          cookie_value = HTTP::Cookie.cookie_value(jar.cookies(request.uri))
+          return if cookie_value.empty?
 
-      module ResponseMethods
-        def cookie_jar
-          return @cookie_jar if defined?(@cookie_jar)
-          return nil unless headers.key?("set-cookie")
-
-          @cookie_jar = begin
-            jar = HTTP::CookieJar.new
-            jar.parse(headers["set-cookie"], @request.uri)
-            jar
-          end
+          add("cookie", cookie_value)
         end
-        alias_method :cookies, :cookie_jar
       end
 
       module OptionsMethods
         def self.included(klass)
           super
           klass.def_option(:cookies) do |cookies|
-            cookies.split(/ *; */) if cookies.is_a?(String)
-            cookies
+            return cookies if cookies.is_a?(Store)
+
+            Store.new(cookies)
           end
         end
       end

test/support/requests/plugins/cookies.rb

@@ -5,14 +5,12 @@ module Requests
     module Cookies
       def test_plugin_cookies_get
         client = HTTPX.plugin(:cookies)
-        assert client.respond_to?(:cookies), "client should be cookie-enabled"
         response = client.get(cookies_uri)
-        assert response.respond_to?(:cookies), "response should have cookies"
         body = json_body(response)
         assert body.key?("cookies")
         assert body["cookies"].empty?
 
-        session_response = client.cookies("abc" => "def").get(cookies_uri)
+        session_response = client.with_cookies("abc" => "def").get(cookies_uri)
         body = json_body(session_response)
         assert body.key?("cookies")
         assert body["cookies"]["abc"] == "def", "abc wasn't properly set"
@@ -24,18 +22,24 @@ module Requests
         session_uri = cookies_set_uri(session_cookies)
         session_response = client.get(session_uri)
         verify_status(session_response, 302)
-        verify_cookies(session_response.cookies, session_cookies)
+        verify_cookies(client.cookies_store[URI(session_uri)], session_cookies)
 
-        # first iteration sets the session
+        # first request sets the session
         response = client.get(cookies_uri)
         body = json_body(response)
         assert body.key?("cookies")
         verify_cookies(body["cookies"], session_cookies)
 
-        extra_cookie_response = client.cookies("c" => "d").get(cookies_uri)
+        # second request reuses the session
+        extra_cookie_response = client.with_cookies("e" => "f").get(cookies_uri)
         body = json_body(extra_cookie_response)
         assert body.key?("cookies")
-        verify_cookies(body["cookies"], session_cookies.merge("c" => "d"))
+        verify_cookies(body["cookies"], session_cookies.merge("e" => "f"))
+
+        # redirect to a different origin only uses the option cookies
+        other_origin_response = client.with_cookies("e" => "f").get(redirect_uri(origin("google.com")))
+        verify_status(other_origin_response, 302)
+        assert !other_origin_response.headers.key?("set-cookie"), "cookies should not transition to next origin"
       end
 
       def test_plugin_cookies_follow

test/support/requests/plugins/follow_redirects.rb

@@ -52,8 +52,8 @@ module Requests
 
       private
 
-      def redirect_uri
-        build_uri("/redirect-to?url=" + redirect_location)
+      def redirect_uri(redirect_uri = redirect_location)
+        build_uri("/redirect-to?url=" + redirect_uri)
       end
 
       def max_redirect_uri(n)