base/mac/scoped_objc_class_swizzler_unittest.mm

   1 // Copyright 2014 The Chromium Authors. All rights reserved.
   2 // Use of this source code is governed by a BSD-style license that can be
   3 // found in the LICENSE file.
   4
   5 #import "base/mac/scoped_objc_class_swizzler.h"
   6
   7 #import "base/mac/scoped_nsobject.h"
   8 #include "testing/gtest/include/gtest/gtest.h"
   9
  10 @interface ObjCClassSwizzlerTestOne : NSObject
  11 + (NSInteger)function;
  12 - (NSInteger)method;
  13 - (NSInteger)modifier;
  14 @end
  15
  16 @interface ObjCClassSwizzlerTestTwo : NSObject
  17 + (NSInteger)function;
  18 - (NSInteger)method;
  19 - (NSInteger)modifier;
  20 @end
  21
  22 @implementation ObjCClassSwizzlerTestOne : NSObject
  23
  24 + (NSInteger)function {
  25   return 10;
  26 }
  27
  28 - (NSInteger)method {
  29   // Multiply by a modifier to ensure |self| in a swizzled implementation
  30   // refers to the original object.
  31   return 1 * [self modifier];
  32 }
  33
  34 - (NSInteger)modifier {
  35   return 3;
  36 }
  37
  38 @end
  39
  40 @implementation ObjCClassSwizzlerTestTwo : NSObject
  41
  42 + (NSInteger)function {
  43   return 20;
  44 }
  45
  46 - (NSInteger)method {
  47   return 2 * [self modifier];
  48 }
  49
  50 - (NSInteger)modifier {
  51   return 7;
  52 }
  53
  54 @end
  55
  56 @interface ObjCClassSwizzlerTestOne (AlternateCategory)
  57 - (NSInteger)alternate;
  58 @end
  59
  60 @implementation ObjCClassSwizzlerTestOne (AlternateCategory)
  61 - (NSInteger)alternate {
  62   return 3 * [self modifier];
  63 }
  64 @end
  65
  66 @interface ObjCClassSwizzlerTestOneChild : ObjCClassSwizzlerTestOne
  67 - (NSInteger)childAlternate;
  68 @end
  69
  70 @implementation ObjCClassSwizzlerTestOneChild
  71 - (NSInteger)childAlternate {
  72   return 5 * [self modifier];
  73 }
  74 @end
  75
  76 namespace base {
  77 namespace mac {
  78
  79 TEST(ObjCClassSwizzlerTest, SwizzleInstanceMethods) {
  80   base::scoped_nsobject<ObjCClassSwizzlerTestOne> object_one(
  81       [[ObjCClassSwizzlerTestOne alloc] init]);
  82   base::scoped_nsobject<ObjCClassSwizzlerTestTwo> object_two(
  83       [[ObjCClassSwizzlerTestTwo alloc] init]);
  84   EXPECT_EQ(3, [object_one method]);
  85   EXPECT_EQ(14, [object_two method]);
  86
  87   {
  88     base::mac::ScopedObjCClassSwizzler swizzler(
  89         [ObjCClassSwizzlerTestOne class],
  90         [ObjCClassSwizzlerTestTwo class],
  91         @selector(method));
  92     EXPECT_EQ(6, [object_one method]);
  93     EXPECT_EQ(7, [object_two method]);
  94
  95     IMP original = swizzler.GetOriginalImplementation();
  96     id expected_result = reinterpret_cast<id>(3);
  97     EXPECT_EQ(expected_result, original(object_one, @selector(method)));
  98   }
  99
 100   EXPECT_EQ(3, [object_one method]);
 101   EXPECT_EQ(14, [object_two method]);
 102 }
 103
 104 TEST(ObjCClassSwizzlerTest, SwizzleClassMethods) {
 105   EXPECT_EQ(10, [ObjCClassSwizzlerTestOne function]);
 106   EXPECT_EQ(20, [ObjCClassSwizzlerTestTwo function]);
 107
 108   {
 109     base::mac::ScopedObjCClassSwizzler swizzler(
 110         [ObjCClassSwizzlerTestOne class],
 111         [ObjCClassSwizzlerTestTwo class],
 112         @selector(function));
 113     EXPECT_EQ(20, [ObjCClassSwizzlerTestOne function]);
 114     EXPECT_EQ(10, [ObjCClassSwizzlerTestTwo function]);
 115
 116     IMP original = swizzler.GetOriginalImplementation();
 117     id expected_result = reinterpret_cast<id>(10);
 118     EXPECT_EQ(expected_result, original(nil, @selector(function)));
 119   }
 120
 121   EXPECT_EQ(10, [ObjCClassSwizzlerTestOne function]);
 122   EXPECT_EQ(20, [ObjCClassSwizzlerTestTwo function]);
 123 }
 124
 125 TEST(ObjCClassSwizzlerTest, SwizzleViaCategory) {
 126   base::scoped_nsobject<ObjCClassSwizzlerTestOne> object_one(
 127       [[ObjCClassSwizzlerTestOne alloc] init]);
 128   EXPECT_EQ(3, [object_one method]);
 129
 130   {
 131     base::mac::ScopedObjCClassSwizzler swizzler(
 132         [ObjCClassSwizzlerTestOne class],
 133         @selector(method),
 134         @selector(alternate));
 135     EXPECT_EQ(9, [object_one method]);
 136
 137     IMP original = swizzler.GetOriginalImplementation();
 138     id expected_result = reinterpret_cast<id>(3);
 139     EXPECT_EQ(expected_result, original(object_one, @selector(method)));
 140   }
 141
 142   EXPECT_EQ(3, [object_one method]);
 143 }
 144
 145 TEST(ObjCClassSwizzlerTest, SwizzleViaInheritance) {
 146   base::scoped_nsobject<ObjCClassSwizzlerTestOneChild> child(
 147       [[ObjCClassSwizzlerTestOneChild alloc] init]);
 148   EXPECT_EQ(3, [child method]);
 149
 150   {
 151     base::mac::ScopedObjCClassSwizzler swizzler(
 152         [ObjCClassSwizzlerTestOneChild class],
 153         @selector(method),
 154         @selector(childAlternate));
 155     EXPECT_EQ(15, [child method]);
 156
 157     IMP original = swizzler.GetOriginalImplementation();
 158     id expected_result = reinterpret_cast<id>(3);
 159     EXPECT_EQ(expected_result, original(child, @selector(method)));
 160   }
 161
 162   EXPECT_EQ(3, [child method]);
 163 }
 164
 165 }  // namespace mac
 166 }  // namespace base