From 0b9266b941a0c7bbeb3d4e17e961088b24e00e64 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Caina=CC=83=20Costa?= <cainan.costa@gmail.com>
Date: Mon, 10 Oct 2011 14:56:50 -0300
Subject: [PATCH] Add check of the arity of the mocked method.

---
 README.md                  | 3 +--
 lib/minitest/fire_mock.rb  | 4 ++++
 test/unit/firemock_test.rb | 7 +++++++
 3 files changed, 12 insertions(+), 2 deletions(-)

diff --git a/README.md b/README.md
index 0840152..3292694 100644
--- a/README.md
+++ b/README.md
@@ -28,11 +28,10 @@ class MyOtherClassTest < MiniTest::Unit::TestCase
 end
 ```
 
-The only real difference of using `MiniTest::FireMock` instead of `MiniTest::Mock` is that if `MyClass` is defined, and `my_method` isn't there, it'll raise a `MockExpectationError`.
+The only real difference of using `MiniTest::FireMock` instead of `MiniTest::Mock` is that if `MyClass` is defined, and `my_method` isn't there, it'll raise a `MockExpectationError`. It checks also for the arity of the method, so it'll raise a `MockExpectationError` if the real method have a different arity than the expectation.
 
 TODO
 ----
 
-- Check for the arity of the methods if they are defined.
 - Mock class/module methods too.
 - Make it work with method_missing (as of now it doesn't, even if the #responds_to? is correct)
diff --git a/lib/minitest/fire_mock.rb b/lib/minitest/fire_mock.rb
index e618018..8a260de 100644
--- a/lib/minitest/fire_mock.rb
+++ b/lib/minitest/fire_mock.rb
@@ -14,6 +14,10 @@ def expect(name, retval, args = [])
       raise MockExpectationError, "expected #{@constant_name} to define `#{name}`, but it doesn't"
     end
 
+    if method and method.arity != args.size
+      raise MockExpectationError, "`#{name}` expects #{method.arity} arguments, given #{args.size}"
+    end
+
     super(name, retval, args)
   end
 
diff --git a/test/unit/firemock_test.rb b/test/unit/firemock_test.rb
index 67678b5..dbc0db4 100644
--- a/test/unit/firemock_test.rb
+++ b/test/unit/firemock_test.rb
@@ -43,4 +43,11 @@ def test_mock_with_namespace_is_invalid_when_defined_but_dont_responds_to_method
       mock.expect(:not_defined_method, 42)
     end
   end
+
+  def test_valid_mock_with_different_arity
+    mock = MiniTest::FireMock.new('DefinedConstant')
+    assert_raises MockExpectationError, "`defined_method` expects 0 arguments, given 3" do
+      mock.expect(:defined_method, 42, [1,2,3])
+    end
+  end
 end