1 #!/usr/bin/python2.4 2 # 3 # Copyright 2008 Google Inc. 4 # 5 # Licensed under the Apache License, Version 2.0 (the "License"); 6 # you may not use this file except in compliance with the License. 7 # You may obtain a copy of the License at 8 # 9 # http://www.apache.org/licenses/LICENSE-2.0 10 # 11 # Unless required by applicable law or agreed to in writing, software 12 # distributed under the License is distributed on an "AS IS" BASIS, 13 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 14 # See the License for the specific language governing permissions and 15 # limitations under the License. 16 17 # This file is used for testing. The original is at: 18 # http://code.google.com/p/pymox/ 19 20 class StubOutForTesting: 21 """Sample Usage: 22 You want os.path.exists() to always return true during testing. 23 24 stubs = StubOutForTesting() 25 stubs.Set(os.path, 'exists', lambda x: 1) 26 ... 27 stubs.UnsetAll() 28 29 The above changes os.path.exists into a lambda that returns 1. Once 30 the ... part of the code finishes, the UnsetAll() looks up the old value 31 of os.path.exists and restores it. 32 33 """ 34 def __init__(self): 35 self.cache = [] 36 self.stubs = [] 37 38 def __del__(self): 39 self.SmartUnsetAll() 40 self.UnsetAll() 41 42 def SmartSet(self, obj, attr_name, new_attr): 43 """Replace obj.attr_name with new_attr. This method is smart and works 44 at the module, class, and instance level while preserving proper 45 inheritance. It will not stub out C types however unless that has been 46 explicitly allowed by the type. 47 48 This method supports the case where attr_name is a staticmethod or a 49 classmethod of obj. 50 51 Notes: 52 - If obj is an instance, then it is its class that will actually be 53 stubbed. Note that the method Set() does not do that: if obj is 54 an instance, it (and not its class) will be stubbed. 55 - The stubbing is using the builtin getattr and setattr. So, the __get__ 56 and __set__ will be called when stubbing (TODO: A better idea would 57 probably be to manipulate obj.__dict__ instead of getattr() and 58 setattr()). 59 60 Raises AttributeError if the attribute cannot be found. 61 """ 62 if (inspect.ismodule(obj) or 63 (not inspect.isclass(obj) and obj.__dict__.has_key(attr_name))): 64 orig_obj = obj 65 orig_attr = getattr(obj, attr_name) 66 67 else: 68 if not inspect.isclass(obj): 69 mro = list(inspect.getmro(obj.__class__)) 70 else: 71 mro = list(inspect.getmro(obj)) 72 73 mro.reverse() 74 75 orig_attr = None 76 77 for cls in mro: 78 try: 79 orig_obj = cls 80 orig_attr = getattr(obj, attr_name) 81 except AttributeError: 82 continue 83 84 if orig_attr is None: 85 raise AttributeError("Attribute not found.") 86 87 # Calling getattr() on a staticmethod transforms it to a 'normal' function. 88 # We need to ensure that we put it back as a staticmethod. 89 old_attribute = obj.__dict__.get(attr_name) 90 if old_attribute is not None and isinstance(old_attribute, staticmethod): 91 orig_attr = staticmethod(orig_attr) 92 93 self.stubs.append((orig_obj, attr_name, orig_attr)) 94 setattr(orig_obj, attr_name, new_attr) 95 96 def SmartUnsetAll(self): 97 """Reverses all the SmartSet() calls, restoring things to their original 98 definition. Its okay to call SmartUnsetAll() repeatedly, as later calls 99 have no effect if no SmartSet() calls have been made. 100 101 """ 102 self.stubs.reverse() 103 104 for args in self.stubs: 105 setattr(*args) 106 107 self.stubs = [] 108 109 def Set(self, parent, child_name, new_child): 110 """Replace child_name's old definition with new_child, in the context 111 of the given parent. The parent could be a module when the child is a 112 function at module scope. Or the parent could be a class when a class' 113 method is being replaced. The named child is set to new_child, while 114 the prior definition is saved away for later, when UnsetAll() is called. 115 116 This method supports the case where child_name is a staticmethod or a 117 classmethod of parent. 118 """ 119 old_child = getattr(parent, child_name) 120 121 old_attribute = parent.__dict__.get(child_name) 122 if old_attribute is not None and isinstance(old_attribute, staticmethod): 123 old_child = staticmethod(old_child) 124 125 self.cache.append((parent, old_child, child_name)) 126 setattr(parent, child_name, new_child) 127 128 def UnsetAll(self): 129 """Reverses all the Set() calls, restoring things to their original 130 definition. Its okay to call UnsetAll() repeatedly, as later calls have 131 no effect if no Set() calls have been made. 132 133 """ 134 # Undo calls to Set() in reverse order, in case Set() was called on the 135 # same arguments repeatedly (want the original call to be last one undone) 136 self.cache.reverse() 137 138 for (parent, old_child, child_name) in self.cache: 139 setattr(parent, child_name, old_child) 140 self.cache = [] 141
