From d21c7b3ecc625e80a491f52fc2d14cb21a71599d Mon Sep 17 00:00:00 2001
From: Alexey Belkevich <belkevich.alexey@gmail.com>
Date: Thu, 10 Oct 2013 11:29:34 +0300
Subject: [PATCH] Added reverting changed methods on 'funtime' object dealloc

---
 Classes/OCFMethod.m                 |  6 ++-
 Classes/OCFuntime.m                 |  5 +++
 Spec/OCFuntimeSpec/OCFuntimeSpec.mm | 62 +++++++++++++++++++++++++----
 3 files changed, 64 insertions(+), 9 deletions(-)

diff --git a/Classes/OCFMethod.m b/Classes/OCFMethod.m
index 44aaa4a..4b3eefc 100644
--- a/Classes/OCFMethod.m
+++ b/Classes/OCFMethod.m
@@ -69,7 +69,11 @@ - (void)changeImplementationWithBlock:(id)block
 
 - (void)revertImplementation
 {
-    method_setImplementation(method, defaultImplementation);
+    if (defaultImplementation)
+    {
+        method_setImplementation(method, defaultImplementation);
+        defaultImplementation = NULL;
+    }
 }
 
 @end
\ No newline at end of file
diff --git a/Classes/OCFuntime.m b/Classes/OCFuntime.m
index c91fa39..05b9096 100644
--- a/Classes/OCFuntime.m
+++ b/Classes/OCFuntime.m
@@ -24,6 +24,11 @@ - (id)init
     return self;
 }
 
+- (void)dealloc
+{
+    [self revertAll];
+}
+
 #pragma mark - public
 
 - (void)changeClass:(Class)theClass instanceMethod:(SEL)method implementation:(id)block
diff --git a/Spec/OCFuntimeSpec/OCFuntimeSpec.mm b/Spec/OCFuntimeSpec/OCFuntimeSpec.mm
index 7511b66..9260d3c 100644
--- a/Spec/OCFuntimeSpec/OCFuntimeSpec.mm
+++ b/Spec/OCFuntimeSpec/OCFuntimeSpec.mm
@@ -14,9 +14,10 @@
 
 SPEC_BEGIN(OCFuntimeSpec)
 
-describe(@"OCFuntime", ^{
-    __block OCFuntime *funtime = nil;
-    __block OCFMock *mock = nil;
+__block OCFuntime *funtime;
+__block OCFMock *mock;
+
+describe(@"OCFuntime with changed method", ^{
 
     beforeEach(^
                {
@@ -82,8 +83,8 @@
     it(@"should call newly changed instance method if it was changed again", ^
     {
         [funtime changeClass:[OCFMock class] instanceMethod:@selector(funInstanceMethod)
-              implementation:^{
-                  // still returns NO
+              implementation:^
+              {
                   NSLog(@"One more time changed FUN instance method");
                   return 2;
               }];
@@ -94,7 +95,6 @@
     {
         [funtime changeClass:[OCFMock class] classMethod:@selector(funClassMethod)
               implementation:^{
-                  // still returns NO
                   NSLog(@"One more time changed FUN class method");
                   return 2;
               }];
@@ -104,7 +104,8 @@
     it(@"should revert to default instance method, not previous setted method", ^
     {
         [funtime changeClass:[OCFMock class] instanceMethod:@selector(funInstanceMethod)
-              implementation:^{
+              implementation:^
+              {
                   // still returns NO
                   NSLog(@"One more time changed FUN instance method");
                   return 2;
@@ -116,7 +117,8 @@
     it(@"should revert to default class method, not previous setted method", ^
     {
         [funtime changeClass:[OCFMock class] classMethod:@selector(funClassMethod)
-              implementation:^{
+              implementation:^
+              {
                   // still returns NO
                   NSLog(@"One more time changed FUN class method");
                   return 2;
@@ -124,6 +126,50 @@
         [funtime revertClass:[OCFMock class]];
         [OCFMock funClassMethod] should equal(0);
     });
+});
+
+describe(@"OCFuntime memory management", ^
+{
+    it(@"should call default instance method if 'funtime' instance deallocated", ^
+    {
+        funtime = [[OCFuntime alloc] init];
+        [funtime changeClass:[OCFMock class] instanceMethod:@selector(funInstanceMethod)
+              implementation:^
+              {
+                  return 1;
+              }];
+        mock = [[OCFMock alloc] init];
+        [mock funInstanceMethod] should equal(1);
+        [funtime release];
+        [mock funInstanceMethod] should equal(0);
+        [mock release];
+    });
+
+    it(@"should call default class method if 'funtime' instance deallocated", ^
+    {
+        funtime = [[OCFuntime alloc] init];
+        [funtime changeClass:[OCFMock class] classMethod:@selector(funClassMethod)
+              implementation:^
+              {
+                  return 1;
+              }];
+        [OCFMock funClassMethod] should equal(1);
+        [funtime release];
+        [OCFMock funClassMethod] should equal(0);
+    });
+});
+
+describe(@"OCFuntime protection", ^{
+
+    beforeEach(^
+               {
+                   funtime = [[OCFuntime alloc] init];
+               });
+
+    afterEach(^
+              {
+                  [funtime release];
+              });
 
     it(@"should throw exception if instance method doesn't exist", ^
     {