From 6a5f4d2d0e8e9525870545741e11eaf3a22aa92b Mon Sep 17 00:00:00 2001
From: Ryan Schlesinger <ryan@instanceinc.com>
Date: Fri, 26 Nov 2010 15:47:01 -0800
Subject: [PATCH] Adding request inspection support.

---
 lib/fake_web.rb              | 14 +++++++----
 lib/fake_web/ext/net_http.rb |  2 +-
 lib/fake_web/registry.rb     | 48 +++++++++++++++++++++++++++---------
 test/test_fake_web.rb        |  2 +-
 4 files changed, 48 insertions(+), 18 deletions(-)

diff --git a/lib/fake_web.rb b/lib/fake_web.rb
index 77fbdae..40d54fc 100644
--- a/lib/fake_web.rb
+++ b/lib/fake_web.rb
@@ -147,13 +147,17 @@ class MultipleMatchingURIsError < StandardError; end;
   # response more than once:
   #
   #   FakeWeb.register_uri(:get, "http://example.com", :set_cookie => ["name=value", "example=1"])
-  def self.register_uri(*args)
+  def self.register_uri(*args, &block)
     case args.length
     when 3
       Registry.instance.register_uri(*args)
     when 2
-      print_missing_http_method_deprecation_warning(*args)
-      Registry.instance.register_uri(:any, *args)
+      if block_given?
+        Registry.instance.register_uri(*args, &block)
+      else
+        print_missing_http_method_deprecation_warning(*args)
+        Registry.instance.register_uri(:any, *args)
+      end
     else
       raise ArgumentError.new("wrong number of arguments (#{args.length} for 3)")
     end
@@ -165,13 +169,13 @@ def self.register_uri(*args)
   # Returns the faked Net::HTTPResponse object associated with +method+ and +uri+.
   def self.response_for(*args, &block) #:nodoc: :yields: response
     case args.length
-    when 2
+    when 3, 2
       Registry.instance.response_for(*args, &block)
     when 1
       print_missing_http_method_deprecation_warning(*args)
       Registry.instance.response_for(:any, *args, &block)
     else
-      raise ArgumentError.new("wrong number of arguments (#{args.length} for 2)")
+      raise ArgumentError.new("wrong number of arguments (#{args.length} for 3)")
     end
   end
 
diff --git a/lib/fake_web/ext/net_http.rb b/lib/fake_web/ext/net_http.rb
index 4ff3e9a..64372ec 100644
--- a/lib/fake_web/ext/net_http.rb
+++ b/lib/fake_web/ext/net_http.rb
@@ -44,7 +44,7 @@ def request_with_fakeweb(request, body = nil, &block)
       if FakeWeb.registered_uri?(method, uri)
         @socket = Net::HTTP.socket_type.new
         FakeWeb::Utility.produce_side_effects_of_net_http_request(request, body)
-        FakeWeb.response_for(method, uri, &block)
+        FakeWeb.response_for(method, uri, request, &block)
       elsif FakeWeb.allow_net_connect?(uri)
         connect_without_fakeweb
         request_without_fakeweb(request, body, &block)
diff --git a/lib/fake_web/registry.rb b/lib/fake_web/registry.rb
index 9a4a34e..8a623e5 100644
--- a/lib/fake_web/registry.rb
+++ b/lib/fake_web/registry.rb
@@ -2,7 +2,7 @@ module FakeWeb
   class Registry #:nodoc:
     include Singleton
 
-    attr_accessor :uri_map, :passthrough_uri_map
+    attr_accessor :uri_map, :passthrough_uri_map, :uri_block_map
 
     def initialize
       clean_registry
@@ -10,19 +10,37 @@ def initialize
 
     def clean_registry
       self.uri_map = Hash.new { |hash, key| hash[key] = {} }
+      self.uri_block_map = Hash.new { |hash, key| hash[key] = {} }
     end
 
-    def register_uri(method, uri, options)
-      uri_map[normalize_uri(uri)][method] = [*[options]].flatten.collect do |option|
-        FakeWeb::Responder.new(method, uri, option, option[:times])
+    def register_uri(method, uri, options = nil, &block)
+      if block_given?
+        uri_block_map[normalize_uri(uri)][method] = block
+      else
+        if options.nil?
+          raise ArgumentError.new("wrong number of arguments (2 for 3)")
+        end
+
+        uri_map[normalize_uri(uri)][method] = [*[options]].flatten.collect do |option|
+          FakeWeb::Responder.new(method, uri, option, option[:times])
+        end
       end
     end
 
     def registered_uri?(method, uri)
-      !responders_for(method, uri).empty?
+      !block_responder_for(method, uri).nil? || !responders_for(method, uri).empty?
     end
 
-    def response_for(method, uri, &block)
+    def response_for(method, uri, request = nil, &block)
+      block_responder = block_responder_for(method, uri)
+      if block_responder
+        block_args = {:method => method, :uri => uri, :request => request}
+        option = block_responder.call(block_args)
+
+        responder = FakeWeb::Responder.new(method, uri, option, option[:times])
+        return responder.response(&block)
+      end
+
       responders = responders_for(method, uri)
       return nil if responders.empty?
 
@@ -55,13 +73,21 @@ def passthrough_uri_matches?(uri)
     private
 
     def responders_for(method, uri)
+      responders_for_helper(uri_map, method, uri, [])
+    end
+
+    def block_responder_for(method, uri)
+      responders_for_helper(uri_block_map, method, uri, nil)
+    end
+
+    def responders_for_helper(map, method, uri, default)
       uri = normalize_uri(uri)
 
-      uri_map_matches(uri_map, method, uri, URI) ||
-      uri_map_matches(uri_map, :any,   uri, URI) ||
-      uri_map_matches(uri_map, method, uri, Regexp) ||
-      uri_map_matches(uri_map, :any,   uri, Regexp) ||
-      []
+      uri_map_matches(map, method, uri, URI) ||
+      uri_map_matches(map, :any,   uri, URI) ||
+      uri_map_matches(map, method, uri, Regexp) ||
+      uri_map_matches(map, :any,   uri, Regexp) ||
+      default
     end
 
     def uri_map_matches(map, method, uri, type_to_check = URI)
diff --git a/test/test_fake_web.rb b/test/test_fake_web.rb
index c6e6b59..faba6eb 100644
--- a/test/test_fake_web.rb
+++ b/test/test_fake_web.rb
@@ -30,7 +30,7 @@ def test_response_for_with_wrong_number_of_arguments
       FakeWeb.response_for
     end
     assert_raises ArgumentError do
-      FakeWeb.response_for(:get, "http://example.com", "/example")
+      FakeWeb.response_for(:get, "http://example.com", "/example", "/foo")
     end
   end