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