1dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block# mock.py 2dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block# Test tools for mocking and patching. 3dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block# Copyright (C) 2007-2009 Michael Foord 4dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block# E-mail: fuzzyman AT voidspace DOT org DOT uk 5dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block 6dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block# mock 0.6.0 7dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block# http://www.voidspace.org.uk/python/mock/ 8dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block 9dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block# Released subject to the BSD License 10dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block# Please see http://www.voidspace.org.uk/python/license.shtml 11dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block 12dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block# 2009-11-25: Licence downloaded from above URL. 13dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block# BEGIN DOWNLOADED LICENSE 14dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block# 15dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block# Copyright (c) 2003-2009, Michael Foord 16dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block# All rights reserved. 17dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block# E-mail : fuzzyman AT voidspace DOT org DOT uk 18dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block# 19dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block# Redistribution and use in source and binary forms, with or without 20dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block# modification, are permitted provided that the following conditions are 21dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block# met: 22dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block# 23dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block# 24dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block# * Redistributions of source code must retain the above copyright 25dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block# notice, this list of conditions and the following disclaimer. 26dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block# 27dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block# * Redistributions in binary form must reproduce the above 28dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block# copyright notice, this list of conditions and the following 29dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block# disclaimer in the documentation and/or other materials provided 30dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block# with the distribution. 31dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block# 32dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block# * Neither the name of Michael Foord nor the name of Voidspace 33dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block# may be used to endorse or promote products derived from this 34dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block# software without specific prior written permission. 35dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block# 36dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS 37dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block# "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT 38dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block# LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR 39dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block# A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT 40dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block# OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, 41dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block# SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT 42dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block# LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, 43dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block# DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY 44dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block# THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT 45dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE 46dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. 47dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block# 48dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block# END DOWNLOADED LICENSE 49dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block 50dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block# Scripts maintained at http://www.voidspace.org.uk/python/index.shtml 51dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block# Comments, suggestions and bug reports welcome. 52dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block 53dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block 54dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block__all__ = ( 55dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block 'Mock', 56dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block 'patch', 57dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block 'patch_object', 58dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block 'sentinel', 59dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block 'DEFAULT' 60dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block) 61dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block 62dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block__version__ = '0.6.0' 63dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block 64dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Blockclass SentinelObject(object): 65dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block def __init__(self, name): 66dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block self.name = name 67dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block 68dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block def __repr__(self): 69dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block return '<SentinelObject "%s">' % self.name 70dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block 71dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block 72dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Blockclass Sentinel(object): 73dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block def __init__(self): 74dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block self._sentinels = {} 75dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block 76dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block def __getattr__(self, name): 77dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block return self._sentinels.setdefault(name, SentinelObject(name)) 78dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block 79dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block 80dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Blocksentinel = Sentinel() 81dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block 82dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve BlockDEFAULT = sentinel.DEFAULT 83dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block 84dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Blockclass OldStyleClass: 85dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block pass 86dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve BlockClassType = type(OldStyleClass) 87dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block 88dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Blockdef _is_magic(name): 89dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block return '__%s__' % name[2:-2] == name 90dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block 91dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Blockdef _copy(value): 92dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block if type(value) in (dict, list, tuple, set): 93dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block return type(value)(value) 94dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block return value 95dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block 96dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block 97dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Blockclass Mock(object): 98dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block 99dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block def __init__(self, spec=None, side_effect=None, return_value=DEFAULT, 100dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block name=None, parent=None, wraps=None): 101dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block self._parent = parent 102dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block self._name = name 103dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block if spec is not None and not isinstance(spec, list): 104dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block spec = [member for member in dir(spec) if not _is_magic(member)] 105dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block 106dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block self._methods = spec 107dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block self._children = {} 108dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block self._return_value = return_value 109dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block self.side_effect = side_effect 110dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block self._wraps = wraps 111dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block 112dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block self.reset_mock() 113dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block 114dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block 115dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block def reset_mock(self): 116dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block self.called = False 117dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block self.call_args = None 118dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block self.call_count = 0 119dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block self.call_args_list = [] 120dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block self.method_calls = [] 121dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block for child in self._children.itervalues(): 122dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block child.reset_mock() 123dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block if isinstance(self._return_value, Mock): 124dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block self._return_value.reset_mock() 125dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block 126dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block 127dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block def __get_return_value(self): 128dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block if self._return_value is DEFAULT: 129dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block self._return_value = Mock() 130dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block return self._return_value 131dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block 132dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block def __set_return_value(self, value): 133dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block self._return_value = value 134dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block 135dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block return_value = property(__get_return_value, __set_return_value) 136dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block 137dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block 138dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block def __call__(self, *args, **kwargs): 139dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block self.called = True 140dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block self.call_count += 1 141dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block self.call_args = (args, kwargs) 142dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block self.call_args_list.append((args, kwargs)) 143dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block 144dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block parent = self._parent 145dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block name = self._name 146dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block while parent is not None: 147dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block parent.method_calls.append((name, args, kwargs)) 148dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block if parent._parent is None: 149dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block break 150dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block name = parent._name + '.' + name 151dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block parent = parent._parent 152dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block 153dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block ret_val = DEFAULT 154dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block if self.side_effect is not None: 155dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block if (isinstance(self.side_effect, Exception) or 156dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block isinstance(self.side_effect, (type, ClassType)) and 157dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block issubclass(self.side_effect, Exception)): 158dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block raise self.side_effect 159dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block 160dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block ret_val = self.side_effect(*args, **kwargs) 161dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block if ret_val is DEFAULT: 162dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block ret_val = self.return_value 163dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block 164dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block if self._wraps is not None and self._return_value is DEFAULT: 165dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block return self._wraps(*args, **kwargs) 166dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block if ret_val is DEFAULT: 167dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block ret_val = self.return_value 168dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block return ret_val 169dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block 170dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block 171dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block def __getattr__(self, name): 172dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block if self._methods is not None: 173dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block if name not in self._methods: 174dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block raise AttributeError("Mock object has no attribute '%s'" % name) 175dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block elif _is_magic(name): 176dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block raise AttributeError(name) 177dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block 178dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block if name not in self._children: 179dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block wraps = None 180dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block if self._wraps is not None: 181dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block wraps = getattr(self._wraps, name) 182dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block self._children[name] = Mock(parent=self, name=name, wraps=wraps) 183dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block 184dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block return self._children[name] 185dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block 186dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block 187dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block def assert_called_with(self, *args, **kwargs): 188dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block assert self.call_args == (args, kwargs), 'Expected: %s\nCalled with: %s' % ((args, kwargs), self.call_args) 189dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block 190dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block 191dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Blockdef _dot_lookup(thing, comp, import_path): 192dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block try: 193dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block return getattr(thing, comp) 194dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block except AttributeError: 195dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block __import__(import_path) 196dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block return getattr(thing, comp) 197dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block 198dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block 199dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Blockdef _importer(target): 200dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block components = target.split('.') 201dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block import_path = components.pop(0) 202dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block thing = __import__(import_path) 203dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block 204dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block for comp in components: 205dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block import_path += ".%s" % comp 206dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block thing = _dot_lookup(thing, comp, import_path) 207dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block return thing 208dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block 209dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block 210dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Blockclass _patch(object): 211dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block def __init__(self, target, attribute, new, spec, create): 212dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block self.target = target 213dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block self.attribute = attribute 214dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block self.new = new 215dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block self.spec = spec 216dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block self.create = create 217dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block self.has_local = False 218dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block 219dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block 220dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block def __call__(self, func): 221dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block if hasattr(func, 'patchings'): 222dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block func.patchings.append(self) 223dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block return func 224dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block 225dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block def patched(*args, **keywargs): 226dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block # don't use a with here (backwards compatability with 2.5) 227dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block extra_args = [] 228dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block for patching in patched.patchings: 229dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block arg = patching.__enter__() 230dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block if patching.new is DEFAULT: 231dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block extra_args.append(arg) 232dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block args += tuple(extra_args) 233dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block try: 234dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block return func(*args, **keywargs) 235dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block finally: 236dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block for patching in getattr(patched, 'patchings', []): 237dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block patching.__exit__() 238dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block 239dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block patched.patchings = [self] 240dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block patched.__name__ = func.__name__ 241dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block patched.compat_co_firstlineno = getattr(func, "compat_co_firstlineno", 242dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block func.func_code.co_firstlineno) 243dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block return patched 244dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block 245dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block 246dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block def get_original(self): 247dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block target = self.target 248dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block name = self.attribute 249dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block create = self.create 250dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block 251dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block original = DEFAULT 252dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block if _has_local_attr(target, name): 253dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block try: 254dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block original = target.__dict__[name] 255dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block except AttributeError: 256dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block # for instances of classes with slots, they have no __dict__ 257dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block original = getattr(target, name) 258dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block elif not create and not hasattr(target, name): 259dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block raise AttributeError("%s does not have the attribute %r" % (target, name)) 260dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block return original 261dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block 262dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block 263dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block def __enter__(self): 264dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block new, spec, = self.new, self.spec 265dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block original = self.get_original() 266dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block if new is DEFAULT: 267dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block # XXXX what if original is DEFAULT - shouldn't use it as a spec 268dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block inherit = False 269dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block if spec == True: 270dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block # set spec to the object we are replacing 271dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block spec = original 272dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block if isinstance(spec, (type, ClassType)): 273dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block inherit = True 274dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block new = Mock(spec=spec) 275dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block if inherit: 276dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block new.return_value = Mock(spec=spec) 277dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block self.temp_original = original 278dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block setattr(self.target, self.attribute, new) 279dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block return new 280dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block 281dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block 282dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block def __exit__(self, *_): 283dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block if self.temp_original is not DEFAULT: 284dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block setattr(self.target, self.attribute, self.temp_original) 285dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block else: 286dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block delattr(self.target, self.attribute) 287dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block del self.temp_original 288dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block 289dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block 290dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Blockdef patch_object(target, attribute, new=DEFAULT, spec=None, create=False): 291dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block return _patch(target, attribute, new, spec, create) 292dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block 293dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block 294dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Blockdef patch(target, new=DEFAULT, spec=None, create=False): 295dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block try: 296dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block target, attribute = target.rsplit('.', 1) 297dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block except (TypeError, ValueError): 298dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block raise TypeError("Need a valid target to patch. You supplied: %r" % (target,)) 299dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block target = _importer(target) 300dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block return _patch(target, attribute, new, spec, create) 301dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block 302dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block 303dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block 304dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Blockdef _has_local_attr(obj, name): 305dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block try: 306dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block return name in vars(obj) 307dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block except TypeError: 308dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block # objects without a __dict__ 309dcc8cf2e65d1aa555cce12431a16547e66b469eeSteve Block return hasattr(obj, name) 310