blob: be96194793ef28fb57b11b33360be6ea390d16c0 [file] [log] [blame]
Michael Foord345266a2012-03-14 12:24:34 -07001# mock.py
2# Test tools for mocking and patching.
Michael Foordc17adf42012-03-14 13:30:29 -07003# Maintained by Michael Foord
4# Backport for other versions of Python available from
Ned Deily9d6d06e2018-06-11 00:45:50 -04005# https://pypi.org/project/mock
Michael Foord345266a2012-03-14 12:24:34 -07006
7__all__ = (
8 'Mock',
9 'MagicMock',
10 'patch',
11 'sentinel',
12 'DEFAULT',
13 'ANY',
14 'call',
15 'create_autospec',
Lisa Roach77b3b772019-05-20 09:19:53 -070016 'AsyncMock',
Michael Foord345266a2012-03-14 12:24:34 -070017 'FILTER_DIR',
18 'NonCallableMock',
19 'NonCallableMagicMock',
20 'mock_open',
21 'PropertyMock',
Mario Corchero552be9d2017-10-17 12:35:11 +010022 'seal',
Michael Foord345266a2012-03-14 12:24:34 -070023)
24
25
26__version__ = '1.0'
27
Lisa Roach77b3b772019-05-20 09:19:53 -070028import asyncio
Xtreak436c2b02019-05-28 12:37:39 +053029import contextlib
Rémi Lapeyre11a88322019-05-07 12:48:36 +020030import io
Michael Foord345266a2012-03-14 12:24:34 -070031import inspect
32import pprint
33import sys
Michael Foordfddcfa22014-04-14 16:25:20 -040034import builtins
Lisa Roach77b3b772019-05-20 09:19:53 -070035from types import CodeType, ModuleType, MethodType
Petter Strandmark47d94242018-10-28 21:37:10 +010036from unittest.util import safe_repr
Antoine Pitrou5c64df72013-02-03 00:23:58 +010037from functools import wraps, partial
Michael Foord345266a2012-03-14 12:24:34 -070038
39
Michael Foordfddcfa22014-04-14 16:25:20 -040040_builtins = {name for name in dir(builtins) if not name.startswith('_')}
41
Michael Foord345266a2012-03-14 12:24:34 -070042FILTER_DIR = True
43
Nick Coghlan0b43bcf2012-05-27 18:17:07 +100044# Workaround for issue #12370
45# Without this, the __class__ properties wouldn't be set correctly
46_safe_super = super
Michael Foord345266a2012-03-14 12:24:34 -070047
Lisa Roach77b3b772019-05-20 09:19:53 -070048def _is_async_obj(obj):
49 if getattr(obj, '__code__', None):
50 return asyncio.iscoroutinefunction(obj) or inspect.isawaitable(obj)
51 else:
52 return False
53
54
Xtreakff6b2e62019-05-27 18:26:23 +053055def _is_async_func(func):
56 if getattr(func, '__code__', None):
57 return asyncio.iscoroutinefunction(func)
58 else:
59 return False
60
61
Michael Foord345266a2012-03-14 12:24:34 -070062def _is_instance_mock(obj):
63 # can't use isinstance on Mock objects because they override __class__
64 # The base class for all mocks is NonCallableMock
65 return issubclass(type(obj), NonCallableMock)
66
67
68def _is_exception(obj):
69 return (
Chris Withers49e27f02019-05-01 08:48:44 +010070 isinstance(obj, BaseException) or
71 isinstance(obj, type) and issubclass(obj, BaseException)
Michael Foord345266a2012-03-14 12:24:34 -070072 )
73
74
Antoine Pitrou5c64df72013-02-03 00:23:58 +010075def _get_signature_object(func, as_instance, eat_self):
76 """
77 Given an arbitrary, possibly callable object, try to create a suitable
78 signature object.
79 Return a (reduced func, signature) tuple, or None.
80 """
81 if isinstance(func, type) and not as_instance:
82 # If it's a type and should be modelled as a type, use __init__.
Chris Withersadbf1782019-05-01 23:04:04 +010083 func = func.__init__
Antoine Pitrou5c64df72013-02-03 00:23:58 +010084 # Skip the `self` argument in __init__
85 eat_self = True
Michael Foord345266a2012-03-14 12:24:34 -070086 elif not isinstance(func, FunctionTypes):
Antoine Pitrou5c64df72013-02-03 00:23:58 +010087 # If we really want to model an instance of the passed type,
88 # __call__ should be looked up, not __init__.
Michael Foord345266a2012-03-14 12:24:34 -070089 try:
90 func = func.__call__
91 except AttributeError:
Antoine Pitrou5c64df72013-02-03 00:23:58 +010092 return None
93 if eat_self:
94 sig_func = partial(func, None)
95 else:
96 sig_func = func
Michael Foord345266a2012-03-14 12:24:34 -070097 try:
Antoine Pitrou5c64df72013-02-03 00:23:58 +010098 return func, inspect.signature(sig_func)
99 except ValueError:
100 # Certain callable types are not supported by inspect.signature()
101 return None
Michael Foord345266a2012-03-14 12:24:34 -0700102
103
104def _check_signature(func, mock, skipfirst, instance=False):
Antoine Pitrou5c64df72013-02-03 00:23:58 +0100105 sig = _get_signature_object(func, instance, skipfirst)
106 if sig is None:
Michael Foord345266a2012-03-14 12:24:34 -0700107 return
Antoine Pitrou5c64df72013-02-03 00:23:58 +0100108 func, sig = sig
109 def checksig(_mock_self, *args, **kwargs):
110 sig.bind(*args, **kwargs)
Michael Foord345266a2012-03-14 12:24:34 -0700111 _copy_func_details(func, checksig)
112 type(mock)._mock_check_sig = checksig
Xtreakf7fa62e2018-12-12 13:24:54 +0530113 type(mock).__signature__ = sig
Michael Foord345266a2012-03-14 12:24:34 -0700114
115
116def _copy_func_details(func, funcopy):
Michael Foord345266a2012-03-14 12:24:34 -0700117 # we explicitly don't copy func.__dict__ into this copy as it would
118 # expose original attributes that should be mocked
Berker Peksag161a4dd2016-12-15 05:21:44 +0300119 for attribute in (
120 '__name__', '__doc__', '__text_signature__',
121 '__module__', '__defaults__', '__kwdefaults__',
122 ):
123 try:
124 setattr(funcopy, attribute, getattr(func, attribute))
125 except AttributeError:
126 pass
Michael Foord345266a2012-03-14 12:24:34 -0700127
128
129def _callable(obj):
130 if isinstance(obj, type):
131 return True
Xtreak9b218562019-04-22 08:00:23 +0530132 if isinstance(obj, (staticmethod, classmethod, MethodType)):
133 return _callable(obj.__func__)
Michael Foord345266a2012-03-14 12:24:34 -0700134 if getattr(obj, '__call__', None) is not None:
135 return True
136 return False
137
138
139def _is_list(obj):
140 # checks for list or tuples
141 # XXXX badly named!
142 return type(obj) in (list, tuple)
143
144
145def _instance_callable(obj):
146 """Given an object, return True if the object is callable.
147 For classes, return True if instances would be callable."""
148 if not isinstance(obj, type):
149 # already an instance
150 return getattr(obj, '__call__', None) is not None
151
Michael Foorda74b3aa2012-03-14 14:40:22 -0700152 # *could* be broken by a class overriding __mro__ or __dict__ via
153 # a metaclass
154 for base in (obj,) + obj.__mro__:
155 if base.__dict__.get('__call__') is not None:
Michael Foord345266a2012-03-14 12:24:34 -0700156 return True
157 return False
158
159
160def _set_signature(mock, original, instance=False):
161 # creates a function with signature (*args, **kwargs) that delegates to a
162 # mock. It still does signature checking by calling a lambda with the same
163 # signature as the original.
Michael Foord345266a2012-03-14 12:24:34 -0700164
165 skipfirst = isinstance(original, type)
Antoine Pitrou5c64df72013-02-03 00:23:58 +0100166 result = _get_signature_object(original, instance, skipfirst)
Michael Foord345266a2012-03-14 12:24:34 -0700167 if result is None:
Aaron Gallagher856cbcc2017-07-19 17:01:14 -0700168 return mock
Antoine Pitrou5c64df72013-02-03 00:23:58 +0100169 func, sig = result
170 def checksig(*args, **kwargs):
171 sig.bind(*args, **kwargs)
Michael Foord345266a2012-03-14 12:24:34 -0700172 _copy_func_details(func, checksig)
173
174 name = original.__name__
175 if not name.isidentifier():
176 name = 'funcopy'
Michael Foord313f85f2012-03-25 18:16:07 +0100177 context = {'_checksig_': checksig, 'mock': mock}
Michael Foord345266a2012-03-14 12:24:34 -0700178 src = """def %s(*args, **kwargs):
Michael Foord313f85f2012-03-25 18:16:07 +0100179 _checksig_(*args, **kwargs)
Michael Foord345266a2012-03-14 12:24:34 -0700180 return mock(*args, **kwargs)""" % name
181 exec (src, context)
182 funcopy = context[name]
Xtreakf7fa62e2018-12-12 13:24:54 +0530183 _setup_func(funcopy, mock, sig)
Michael Foord345266a2012-03-14 12:24:34 -0700184 return funcopy
185
186
Xtreakf7fa62e2018-12-12 13:24:54 +0530187def _setup_func(funcopy, mock, sig):
Michael Foord345266a2012-03-14 12:24:34 -0700188 funcopy.mock = mock
189
Michael Foord345266a2012-03-14 12:24:34 -0700190 def assert_called_with(*args, **kwargs):
191 return mock.assert_called_with(*args, **kwargs)
Gregory P. Smithac5084b2016-10-06 14:31:23 -0700192 def assert_called(*args, **kwargs):
193 return mock.assert_called(*args, **kwargs)
194 def assert_not_called(*args, **kwargs):
195 return mock.assert_not_called(*args, **kwargs)
196 def assert_called_once(*args, **kwargs):
197 return mock.assert_called_once(*args, **kwargs)
Michael Foord345266a2012-03-14 12:24:34 -0700198 def assert_called_once_with(*args, **kwargs):
199 return mock.assert_called_once_with(*args, **kwargs)
200 def assert_has_calls(*args, **kwargs):
201 return mock.assert_has_calls(*args, **kwargs)
202 def assert_any_call(*args, **kwargs):
203 return mock.assert_any_call(*args, **kwargs)
204 def reset_mock():
205 funcopy.method_calls = _CallList()
206 funcopy.mock_calls = _CallList()
207 mock.reset_mock()
208 ret = funcopy.return_value
209 if _is_instance_mock(ret) and not ret is mock:
210 ret.reset_mock()
211
212 funcopy.called = False
213 funcopy.call_count = 0
214 funcopy.call_args = None
215 funcopy.call_args_list = _CallList()
216 funcopy.method_calls = _CallList()
217 funcopy.mock_calls = _CallList()
218
219 funcopy.return_value = mock.return_value
220 funcopy.side_effect = mock.side_effect
221 funcopy._mock_children = mock._mock_children
222
223 funcopy.assert_called_with = assert_called_with
224 funcopy.assert_called_once_with = assert_called_once_with
225 funcopy.assert_has_calls = assert_has_calls
226 funcopy.assert_any_call = assert_any_call
227 funcopy.reset_mock = reset_mock
Gregory P. Smithac5084b2016-10-06 14:31:23 -0700228 funcopy.assert_called = assert_called
229 funcopy.assert_not_called = assert_not_called
230 funcopy.assert_called_once = assert_called_once
Xtreakf7fa62e2018-12-12 13:24:54 +0530231 funcopy.__signature__ = sig
Michael Foord345266a2012-03-14 12:24:34 -0700232
233 mock._mock_delegate = funcopy
234
235
Xtreakff6b2e62019-05-27 18:26:23 +0530236def _setup_async_mock(mock):
237 mock._is_coroutine = asyncio.coroutines._is_coroutine
238 mock.await_count = 0
239 mock.await_args = None
240 mock.await_args_list = _CallList()
241 mock.awaited = _AwaitEvent(mock)
242
243 # Mock is not configured yet so the attributes are set
244 # to a function and then the corresponding mock helper function
245 # is called when the helper is accessed similar to _setup_func.
246 def wrapper(attr, *args, **kwargs):
247 return getattr(mock.mock, attr)(*args, **kwargs)
248
249 for attribute in ('assert_awaited',
250 'assert_awaited_once',
251 'assert_awaited_with',
252 'assert_awaited_once_with',
253 'assert_any_await',
254 'assert_has_awaits',
255 'assert_not_awaited'):
256
257 # setattr(mock, attribute, wrapper) causes late binding
258 # hence attribute will always be the last value in the loop
259 # Use partial(wrapper, attribute) to ensure the attribute is bound
260 # correctly.
261 setattr(mock, attribute, partial(wrapper, attribute))
262
263
Michael Foord345266a2012-03-14 12:24:34 -0700264def _is_magic(name):
265 return '__%s__' % name[2:-2] == name
266
267
268class _SentinelObject(object):
269 "A unique, named, sentinel object."
270 def __init__(self, name):
271 self.name = name
272
273 def __repr__(self):
274 return 'sentinel.%s' % self.name
275
Serhiy Storchakad9c956f2017-01-11 20:13:03 +0200276 def __reduce__(self):
277 return 'sentinel.%s' % self.name
278
Michael Foord345266a2012-03-14 12:24:34 -0700279
280class _Sentinel(object):
281 """Access attributes to return a named object, usable as a sentinel."""
282 def __init__(self):
283 self._sentinels = {}
284
285 def __getattr__(self, name):
286 if name == '__bases__':
287 # Without this help(unittest.mock) raises an exception
288 raise AttributeError
289 return self._sentinels.setdefault(name, _SentinelObject(name))
290
Serhiy Storchakad9c956f2017-01-11 20:13:03 +0200291 def __reduce__(self):
292 return 'sentinel'
293
Michael Foord345266a2012-03-14 12:24:34 -0700294
295sentinel = _Sentinel()
296
297DEFAULT = sentinel.DEFAULT
298_missing = sentinel.MISSING
299_deleted = sentinel.DELETED
300
301
Serhiy Storchakac02d1882014-12-11 10:28:14 +0200302_allowed_names = {
303 'return_value', '_mock_return_value', 'side_effect',
304 '_mock_side_effect', '_mock_parent', '_mock_new_parent',
305 '_mock_name', '_mock_new_name'
306}
Michael Foord345266a2012-03-14 12:24:34 -0700307
308
309def _delegating_property(name):
310 _allowed_names.add(name)
311 _the_name = '_mock_' + name
312 def _get(self, name=name, _the_name=_the_name):
313 sig = self._mock_delegate
314 if sig is None:
315 return getattr(self, _the_name)
316 return getattr(sig, name)
317 def _set(self, value, name=name, _the_name=_the_name):
318 sig = self._mock_delegate
319 if sig is None:
320 self.__dict__[_the_name] = value
321 else:
322 setattr(sig, name, value)
323
324 return property(_get, _set)
325
326
327
328class _CallList(list):
329
330 def __contains__(self, value):
331 if not isinstance(value, list):
332 return list.__contains__(self, value)
333 len_value = len(value)
334 len_self = len(self)
335 if len_value > len_self:
336 return False
337
338 for i in range(0, len_self - len_value + 1):
339 sub_list = self[i:i+len_value]
340 if sub_list == value:
341 return True
342 return False
343
344 def __repr__(self):
345 return pprint.pformat(list(self))
346
347
348def _check_and_set_parent(parent, value, name, new_name):
Xtreak9c3f2842019-02-26 03:16:34 +0530349 # function passed to create_autospec will have mock
350 # attribute attached to which parent must be set
351 if isinstance(value, FunctionTypes):
352 try:
353 value = value.mock
354 except AttributeError:
355 pass
356
Michael Foord345266a2012-03-14 12:24:34 -0700357 if not _is_instance_mock(value):
358 return False
359 if ((value._mock_name or value._mock_new_name) or
360 (value._mock_parent is not None) or
361 (value._mock_new_parent is not None)):
362 return False
363
364 _parent = parent
365 while _parent is not None:
366 # setting a mock (value) as a child or return value of itself
367 # should not modify the mock
368 if _parent is value:
369 return False
370 _parent = _parent._mock_new_parent
371
372 if new_name:
373 value._mock_new_parent = parent
374 value._mock_new_name = new_name
375 if name:
376 value._mock_parent = parent
377 value._mock_name = name
378 return True
379
Michael Foord01bafdc2014-04-14 16:09:42 -0400380# Internal class to identify if we wrapped an iterator object or not.
381class _MockIter(object):
382 def __init__(self, obj):
383 self.obj = iter(obj)
Michael Foord01bafdc2014-04-14 16:09:42 -0400384 def __next__(self):
385 return next(self.obj)
Michael Foord345266a2012-03-14 12:24:34 -0700386
387class Base(object):
388 _mock_return_value = DEFAULT
389 _mock_side_effect = None
390 def __init__(self, *args, **kwargs):
391 pass
392
393
394
395class NonCallableMock(Base):
396 """A non-callable version of `Mock`"""
397
398 def __new__(cls, *args, **kw):
399 # every instance has its own class
400 # so we can create magic methods on the
401 # class without stomping on other mocks
Lisa Roach77b3b772019-05-20 09:19:53 -0700402 bases = (cls,)
403 if not issubclass(cls, AsyncMock):
404 # Check if spec is an async object or function
405 sig = inspect.signature(NonCallableMock.__init__)
406 bound_args = sig.bind_partial(cls, *args, **kw).arguments
407 spec_arg = [
408 arg for arg in bound_args.keys()
409 if arg.startswith('spec')
410 ]
411 if spec_arg:
412 # what if spec_set is different than spec?
413 if _is_async_obj(bound_args[spec_arg[0]]):
414 bases = (AsyncMockMixin, cls,)
415 new = type(cls.__name__, bases, {'__doc__': cls.__doc__})
Michael Foord345266a2012-03-14 12:24:34 -0700416 instance = object.__new__(new)
417 return instance
418
419
420 def __init__(
421 self, spec=None, wraps=None, name=None, spec_set=None,
422 parent=None, _spec_state=None, _new_name='', _new_parent=None,
Kushal Das8c145342014-04-16 23:32:21 +0530423 _spec_as_instance=False, _eat_self=None, unsafe=False, **kwargs
Michael Foord345266a2012-03-14 12:24:34 -0700424 ):
425 if _new_parent is None:
426 _new_parent = parent
427
428 __dict__ = self.__dict__
429 __dict__['_mock_parent'] = parent
430 __dict__['_mock_name'] = name
431 __dict__['_mock_new_name'] = _new_name
432 __dict__['_mock_new_parent'] = _new_parent
Mario Corchero552be9d2017-10-17 12:35:11 +0100433 __dict__['_mock_sealed'] = False
Michael Foord345266a2012-03-14 12:24:34 -0700434
435 if spec_set is not None:
436 spec = spec_set
437 spec_set = True
Antoine Pitrou5c64df72013-02-03 00:23:58 +0100438 if _eat_self is None:
439 _eat_self = parent is not None
Michael Foord345266a2012-03-14 12:24:34 -0700440
Antoine Pitrou5c64df72013-02-03 00:23:58 +0100441 self._mock_add_spec(spec, spec_set, _spec_as_instance, _eat_self)
Michael Foord345266a2012-03-14 12:24:34 -0700442
443 __dict__['_mock_children'] = {}
444 __dict__['_mock_wraps'] = wraps
445 __dict__['_mock_delegate'] = None
446
447 __dict__['_mock_called'] = False
448 __dict__['_mock_call_args'] = None
449 __dict__['_mock_call_count'] = 0
450 __dict__['_mock_call_args_list'] = _CallList()
451 __dict__['_mock_mock_calls'] = _CallList()
452
453 __dict__['method_calls'] = _CallList()
Kushal Das8c145342014-04-16 23:32:21 +0530454 __dict__['_mock_unsafe'] = unsafe
Michael Foord345266a2012-03-14 12:24:34 -0700455
456 if kwargs:
457 self.configure_mock(**kwargs)
458
Nick Coghlan0b43bcf2012-05-27 18:17:07 +1000459 _safe_super(NonCallableMock, self).__init__(
Michael Foord345266a2012-03-14 12:24:34 -0700460 spec, wraps, name, spec_set, parent,
461 _spec_state
462 )
463
464
465 def attach_mock(self, mock, attribute):
466 """
467 Attach a mock as an attribute of this one, replacing its name and
468 parent. Calls to the attached mock will be recorded in the
469 `method_calls` and `mock_calls` attributes of this one."""
470 mock._mock_parent = None
471 mock._mock_new_parent = None
472 mock._mock_name = ''
473 mock._mock_new_name = None
474
475 setattr(self, attribute, mock)
476
477
478 def mock_add_spec(self, spec, spec_set=False):
479 """Add a spec to a mock. `spec` can either be an object or a
480 list of strings. Only attributes on the `spec` can be fetched as
481 attributes from the mock.
482
483 If `spec_set` is True then only attributes on the spec can be set."""
484 self._mock_add_spec(spec, spec_set)
485
486
Antoine Pitrou5c64df72013-02-03 00:23:58 +0100487 def _mock_add_spec(self, spec, spec_set, _spec_as_instance=False,
488 _eat_self=False):
Michael Foord345266a2012-03-14 12:24:34 -0700489 _spec_class = None
Antoine Pitrou5c64df72013-02-03 00:23:58 +0100490 _spec_signature = None
Lisa Roach77b3b772019-05-20 09:19:53 -0700491 _spec_asyncs = []
492
493 for attr in dir(spec):
494 if asyncio.iscoroutinefunction(getattr(spec, attr, None)):
495 _spec_asyncs.append(attr)
Michael Foord345266a2012-03-14 12:24:34 -0700496
497 if spec is not None and not _is_list(spec):
498 if isinstance(spec, type):
499 _spec_class = spec
500 else:
Chris Withersadbf1782019-05-01 23:04:04 +0100501 _spec_class = type(spec)
Antoine Pitrou5c64df72013-02-03 00:23:58 +0100502 res = _get_signature_object(spec,
503 _spec_as_instance, _eat_self)
504 _spec_signature = res and res[1]
Michael Foord345266a2012-03-14 12:24:34 -0700505
506 spec = dir(spec)
507
508 __dict__ = self.__dict__
509 __dict__['_spec_class'] = _spec_class
510 __dict__['_spec_set'] = spec_set
Antoine Pitrou5c64df72013-02-03 00:23:58 +0100511 __dict__['_spec_signature'] = _spec_signature
Michael Foord345266a2012-03-14 12:24:34 -0700512 __dict__['_mock_methods'] = spec
Lisa Roach77b3b772019-05-20 09:19:53 -0700513 __dict__['_spec_asyncs'] = _spec_asyncs
Michael Foord345266a2012-03-14 12:24:34 -0700514
515 def __get_return_value(self):
516 ret = self._mock_return_value
517 if self._mock_delegate is not None:
518 ret = self._mock_delegate.return_value
519
520 if ret is DEFAULT:
521 ret = self._get_child_mock(
522 _new_parent=self, _new_name='()'
523 )
524 self.return_value = ret
525 return ret
526
527
528 def __set_return_value(self, value):
529 if self._mock_delegate is not None:
530 self._mock_delegate.return_value = value
531 else:
532 self._mock_return_value = value
533 _check_and_set_parent(self, value, None, '()')
534
535 __return_value_doc = "The value to be returned when the mock is called."
536 return_value = property(__get_return_value, __set_return_value,
537 __return_value_doc)
538
539
540 @property
541 def __class__(self):
542 if self._spec_class is None:
543 return type(self)
544 return self._spec_class
545
546 called = _delegating_property('called')
547 call_count = _delegating_property('call_count')
548 call_args = _delegating_property('call_args')
549 call_args_list = _delegating_property('call_args_list')
550 mock_calls = _delegating_property('mock_calls')
551
552
553 def __get_side_effect(self):
554 delegated = self._mock_delegate
555 if delegated is None:
556 return self._mock_side_effect
Michael Foord01bafdc2014-04-14 16:09:42 -0400557 sf = delegated.side_effect
Robert Collinsf58f88c2015-07-14 13:51:40 +1200558 if (sf is not None and not callable(sf)
559 and not isinstance(sf, _MockIter) and not _is_exception(sf)):
Michael Foord01bafdc2014-04-14 16:09:42 -0400560 sf = _MockIter(sf)
561 delegated.side_effect = sf
562 return sf
Michael Foord345266a2012-03-14 12:24:34 -0700563
564 def __set_side_effect(self, value):
565 value = _try_iter(value)
566 delegated = self._mock_delegate
567 if delegated is None:
568 self._mock_side_effect = value
569 else:
570 delegated.side_effect = value
571
572 side_effect = property(__get_side_effect, __set_side_effect)
573
574
Kushal Das9cd39a12016-06-02 10:20:16 -0700575 def reset_mock(self, visited=None,*, return_value=False, side_effect=False):
Michael Foord345266a2012-03-14 12:24:34 -0700576 "Restore the mock object to its initial state."
Robert Collinsb37f43f2015-07-15 11:42:28 +1200577 if visited is None:
578 visited = []
579 if id(self) in visited:
580 return
581 visited.append(id(self))
582
Michael Foord345266a2012-03-14 12:24:34 -0700583 self.called = False
584 self.call_args = None
585 self.call_count = 0
586 self.mock_calls = _CallList()
587 self.call_args_list = _CallList()
588 self.method_calls = _CallList()
589
Kushal Das9cd39a12016-06-02 10:20:16 -0700590 if return_value:
591 self._mock_return_value = DEFAULT
592 if side_effect:
593 self._mock_side_effect = None
594
Michael Foord345266a2012-03-14 12:24:34 -0700595 for child in self._mock_children.values():
Xtreakedeca922018-12-01 15:33:54 +0530596 if isinstance(child, _SpecState) or child is _deleted:
Michael Foord75963642012-06-09 17:31:59 +0100597 continue
Robert Collinsb37f43f2015-07-15 11:42:28 +1200598 child.reset_mock(visited)
Michael Foord345266a2012-03-14 12:24:34 -0700599
600 ret = self._mock_return_value
601 if _is_instance_mock(ret) and ret is not self:
Robert Collinsb37f43f2015-07-15 11:42:28 +1200602 ret.reset_mock(visited)
Michael Foord345266a2012-03-14 12:24:34 -0700603
604
605 def configure_mock(self, **kwargs):
606 """Set attributes on the mock through keyword arguments.
607
608 Attributes plus return values and side effects can be set on child
609 mocks using standard dot notation and unpacking a dictionary in the
610 method call:
611
612 >>> attrs = {'method.return_value': 3, 'other.side_effect': KeyError}
613 >>> mock.configure_mock(**attrs)"""
614 for arg, val in sorted(kwargs.items(),
615 # we sort on the number of dots so that
616 # attributes are set before we set attributes on
617 # attributes
618 key=lambda entry: entry[0].count('.')):
619 args = arg.split('.')
620 final = args.pop()
621 obj = self
622 for entry in args:
623 obj = getattr(obj, entry)
624 setattr(obj, final, val)
625
626
627 def __getattr__(self, name):
Kushal Das8c145342014-04-16 23:32:21 +0530628 if name in {'_mock_methods', '_mock_unsafe'}:
Michael Foord345266a2012-03-14 12:24:34 -0700629 raise AttributeError(name)
630 elif self._mock_methods is not None:
631 if name not in self._mock_methods or name in _all_magics:
632 raise AttributeError("Mock object has no attribute %r" % name)
633 elif _is_magic(name):
634 raise AttributeError(name)
Kushal Das8c145342014-04-16 23:32:21 +0530635 if not self._mock_unsafe:
636 if name.startswith(('assert', 'assret')):
Zackery Spytzb9b08cd2019-05-08 11:32:24 -0600637 raise AttributeError("Attributes cannot start with 'assert' "
638 "or 'assret'")
Michael Foord345266a2012-03-14 12:24:34 -0700639
640 result = self._mock_children.get(name)
641 if result is _deleted:
642 raise AttributeError(name)
643 elif result is None:
644 wraps = None
645 if self._mock_wraps is not None:
646 # XXXX should we get the attribute without triggering code
647 # execution?
648 wraps = getattr(self._mock_wraps, name)
649
650 result = self._get_child_mock(
651 parent=self, name=name, wraps=wraps, _new_name=name,
652 _new_parent=self
653 )
654 self._mock_children[name] = result
655
656 elif isinstance(result, _SpecState):
657 result = create_autospec(
658 result.spec, result.spec_set, result.instance,
659 result.parent, result.name
660 )
661 self._mock_children[name] = result
662
663 return result
664
665
Mario Corchero552be9d2017-10-17 12:35:11 +0100666 def _extract_mock_name(self):
Michael Foord345266a2012-03-14 12:24:34 -0700667 _name_list = [self._mock_new_name]
668 _parent = self._mock_new_parent
669 last = self
670
671 dot = '.'
672 if _name_list == ['()']:
673 dot = ''
Chris Withersadbf1782019-05-01 23:04:04 +0100674
Michael Foord345266a2012-03-14 12:24:34 -0700675 while _parent is not None:
676 last = _parent
677
678 _name_list.append(_parent._mock_new_name + dot)
679 dot = '.'
680 if _parent._mock_new_name == '()':
681 dot = ''
682
683 _parent = _parent._mock_new_parent
684
Michael Foord345266a2012-03-14 12:24:34 -0700685 _name_list = list(reversed(_name_list))
686 _first = last._mock_name or 'mock'
687 if len(_name_list) > 1:
688 if _name_list[1] not in ('()', '().'):
689 _first += '.'
690 _name_list[0] = _first
Mario Corchero552be9d2017-10-17 12:35:11 +0100691 return ''.join(_name_list)
692
693 def __repr__(self):
694 name = self._extract_mock_name()
Michael Foord345266a2012-03-14 12:24:34 -0700695
696 name_string = ''
697 if name not in ('mock', 'mock.'):
698 name_string = ' name=%r' % name
699
700 spec_string = ''
701 if self._spec_class is not None:
702 spec_string = ' spec=%r'
703 if self._spec_set:
704 spec_string = ' spec_set=%r'
705 spec_string = spec_string % self._spec_class.__name__
706 return "<%s%s%s id='%s'>" % (
707 type(self).__name__,
708 name_string,
709 spec_string,
710 id(self)
711 )
712
713
714 def __dir__(self):
Michael Foordd7c65e22012-03-14 14:56:54 -0700715 """Filter the output of `dir(mock)` to only useful members."""
Michael Foord313f85f2012-03-25 18:16:07 +0100716 if not FILTER_DIR:
717 return object.__dir__(self)
718
Michael Foord345266a2012-03-14 12:24:34 -0700719 extras = self._mock_methods or []
720 from_type = dir(type(self))
721 from_dict = list(self.__dict__)
Mario Corchero0df635c2019-04-30 19:56:36 +0100722 from_child_mocks = [
723 m_name for m_name, m_value in self._mock_children.items()
724 if m_value is not _deleted]
Michael Foord345266a2012-03-14 12:24:34 -0700725
Michael Foord313f85f2012-03-25 18:16:07 +0100726 from_type = [e for e in from_type if not e.startswith('_')]
727 from_dict = [e for e in from_dict if not e.startswith('_') or
728 _is_magic(e)]
Mario Corchero0df635c2019-04-30 19:56:36 +0100729 return sorted(set(extras + from_type + from_dict + from_child_mocks))
Michael Foord345266a2012-03-14 12:24:34 -0700730
731
732 def __setattr__(self, name, value):
733 if name in _allowed_names:
734 # property setters go through here
735 return object.__setattr__(self, name, value)
736 elif (self._spec_set and self._mock_methods is not None and
737 name not in self._mock_methods and
738 name not in self.__dict__):
739 raise AttributeError("Mock object has no attribute '%s'" % name)
740 elif name in _unsupported_magics:
741 msg = 'Attempting to set unsupported magic method %r.' % name
742 raise AttributeError(msg)
743 elif name in _all_magics:
744 if self._mock_methods is not None and name not in self._mock_methods:
745 raise AttributeError("Mock object has no attribute '%s'" % name)
746
747 if not _is_instance_mock(value):
748 setattr(type(self), name, _get_method(name, value))
749 original = value
750 value = lambda *args, **kw: original(self, *args, **kw)
751 else:
752 # only set _new_name and not name so that mock_calls is tracked
753 # but not method calls
754 _check_and_set_parent(self, value, None, name)
755 setattr(type(self), name, value)
Michael Foord75963642012-06-09 17:31:59 +0100756 self._mock_children[name] = value
Michael Foord345266a2012-03-14 12:24:34 -0700757 elif name == '__class__':
758 self._spec_class = value
759 return
760 else:
761 if _check_and_set_parent(self, value, name, name):
762 self._mock_children[name] = value
Mario Corchero552be9d2017-10-17 12:35:11 +0100763
764 if self._mock_sealed and not hasattr(self, name):
765 mock_name = f'{self._extract_mock_name()}.{name}'
766 raise AttributeError(f'Cannot set {mock_name}')
767
Michael Foord345266a2012-03-14 12:24:34 -0700768 return object.__setattr__(self, name, value)
769
770
771 def __delattr__(self, name):
772 if name in _all_magics and name in type(self).__dict__:
773 delattr(type(self), name)
774 if name not in self.__dict__:
775 # for magic methods that are still MagicProxy objects and
776 # not set on the instance itself
777 return
778
Michael Foord345266a2012-03-14 12:24:34 -0700779 obj = self._mock_children.get(name, _missing)
Pablo Galindo222d3032019-01-21 08:57:46 +0000780 if name in self.__dict__:
Xtreak830b43d2019-04-14 00:42:33 +0530781 _safe_super(NonCallableMock, self).__delattr__(name)
Pablo Galindo222d3032019-01-21 08:57:46 +0000782 elif obj is _deleted:
Michael Foord345266a2012-03-14 12:24:34 -0700783 raise AttributeError(name)
784 if obj is not _missing:
785 del self._mock_children[name]
786 self._mock_children[name] = _deleted
787
788
Michael Foord345266a2012-03-14 12:24:34 -0700789 def _format_mock_call_signature(self, args, kwargs):
790 name = self._mock_name or 'mock'
791 return _format_call_signature(name, args, kwargs)
792
793
Xtreak0ae022c2019-05-29 12:32:26 +0530794 def _format_mock_failure_message(self, args, kwargs, action='call'):
795 message = 'expected %s not found.\nExpected: %s\nActual: %s'
Michael Foord345266a2012-03-14 12:24:34 -0700796 expected_string = self._format_mock_call_signature(args, kwargs)
797 call_args = self.call_args
Michael Foord345266a2012-03-14 12:24:34 -0700798 actual_string = self._format_mock_call_signature(*call_args)
Xtreak0ae022c2019-05-29 12:32:26 +0530799 return message % (action, expected_string, actual_string)
Michael Foord345266a2012-03-14 12:24:34 -0700800
801
Antoine Pitrou5c64df72013-02-03 00:23:58 +0100802 def _call_matcher(self, _call):
803 """
Martin Panter204bf0b2016-07-11 07:51:37 +0000804 Given a call (or simply an (args, kwargs) tuple), return a
Antoine Pitrou5c64df72013-02-03 00:23:58 +0100805 comparison key suitable for matching with other calls.
806 This is a best effort method which relies on the spec's signature,
807 if available, or falls back on the arguments themselves.
808 """
809 sig = self._spec_signature
810 if sig is not None:
811 if len(_call) == 2:
812 name = ''
813 args, kwargs = _call
814 else:
815 name, args, kwargs = _call
816 try:
817 return name, sig.bind(*args, **kwargs)
818 except TypeError as e:
819 return e.with_traceback(None)
820 else:
821 return _call
822
Kushal Das68290f42014-04-17 01:54:07 +0530823 def assert_not_called(_mock_self):
Kushal Das8af9db32014-04-17 01:36:14 +0530824 """assert that the mock was never called.
825 """
826 self = _mock_self
827 if self.call_count != 0:
Petter Strandmark47d94242018-10-28 21:37:10 +0100828 msg = ("Expected '%s' to not have been called. Called %s times.%s"
829 % (self._mock_name or 'mock',
830 self.call_count,
831 self._calls_repr()))
Kushal Das8af9db32014-04-17 01:36:14 +0530832 raise AssertionError(msg)
Antoine Pitrou5c64df72013-02-03 00:23:58 +0100833
Victor Stinner2c2a4e62016-03-11 22:17:48 +0100834 def assert_called(_mock_self):
835 """assert that the mock was called at least once
836 """
837 self = _mock_self
838 if self.call_count == 0:
839 msg = ("Expected '%s' to have been called." %
840 self._mock_name or 'mock')
841 raise AssertionError(msg)
842
843 def assert_called_once(_mock_self):
844 """assert that the mock was called only once.
845 """
846 self = _mock_self
847 if not self.call_count == 1:
Petter Strandmark47d94242018-10-28 21:37:10 +0100848 msg = ("Expected '%s' to have been called once. Called %s times.%s"
849 % (self._mock_name or 'mock',
850 self.call_count,
851 self._calls_repr()))
Victor Stinner2c2a4e62016-03-11 22:17:48 +0100852 raise AssertionError(msg)
853
Michael Foord345266a2012-03-14 12:24:34 -0700854 def assert_called_with(_mock_self, *args, **kwargs):
855 """assert that the mock was called with the specified arguments.
856
857 Raises an AssertionError if the args and keyword args passed in are
858 different to the last call to the mock."""
859 self = _mock_self
860 if self.call_args is None:
861 expected = self._format_mock_call_signature(args, kwargs)
Susan Su2bdd5852019-02-13 18:22:29 -0800862 actual = 'not called.'
863 error_message = ('expected call not found.\nExpected: %s\nActual: %s'
864 % (expected, actual))
865 raise AssertionError(error_message)
Michael Foord345266a2012-03-14 12:24:34 -0700866
Antoine Pitrou5c64df72013-02-03 00:23:58 +0100867 def _error_message():
Michael Foord345266a2012-03-14 12:24:34 -0700868 msg = self._format_mock_failure_message(args, kwargs)
Antoine Pitrou5c64df72013-02-03 00:23:58 +0100869 return msg
870 expected = self._call_matcher((args, kwargs))
871 actual = self._call_matcher(self.call_args)
872 if expected != actual:
873 cause = expected if isinstance(expected, Exception) else None
874 raise AssertionError(_error_message()) from cause
Michael Foord345266a2012-03-14 12:24:34 -0700875
876
877 def assert_called_once_with(_mock_self, *args, **kwargs):
Arne de Laat324c5d82017-02-23 15:57:25 +0100878 """assert that the mock was called exactly once and that that call was
879 with the specified arguments."""
Michael Foord345266a2012-03-14 12:24:34 -0700880 self = _mock_self
881 if not self.call_count == 1:
Petter Strandmark47d94242018-10-28 21:37:10 +0100882 msg = ("Expected '%s' to be called once. Called %s times.%s"
883 % (self._mock_name or 'mock',
884 self.call_count,
885 self._calls_repr()))
Michael Foord345266a2012-03-14 12:24:34 -0700886 raise AssertionError(msg)
887 return self.assert_called_with(*args, **kwargs)
888
889
890 def assert_has_calls(self, calls, any_order=False):
891 """assert the mock has been called with the specified calls.
892 The `mock_calls` list is checked for the calls.
893
894 If `any_order` is False (the default) then the calls must be
895 sequential. There can be extra calls before or after the
896 specified calls.
897
898 If `any_order` is True then the calls can be in any order, but
899 they must all appear in `mock_calls`."""
Antoine Pitrou5c64df72013-02-03 00:23:58 +0100900 expected = [self._call_matcher(c) for c in calls]
901 cause = expected if isinstance(expected, Exception) else None
902 all_calls = _CallList(self._call_matcher(c) for c in self.mock_calls)
Michael Foord345266a2012-03-14 12:24:34 -0700903 if not any_order:
Antoine Pitrou5c64df72013-02-03 00:23:58 +0100904 if expected not in all_calls:
Michael Foord345266a2012-03-14 12:24:34 -0700905 raise AssertionError(
Petter Strandmark47d94242018-10-28 21:37:10 +0100906 'Calls not found.\nExpected: %r%s'
907 % (_CallList(calls), self._calls_repr(prefix="Actual"))
Antoine Pitrou5c64df72013-02-03 00:23:58 +0100908 ) from cause
Michael Foord345266a2012-03-14 12:24:34 -0700909 return
910
Antoine Pitrou5c64df72013-02-03 00:23:58 +0100911 all_calls = list(all_calls)
Michael Foord345266a2012-03-14 12:24:34 -0700912
913 not_found = []
Antoine Pitrou5c64df72013-02-03 00:23:58 +0100914 for kall in expected:
Michael Foord345266a2012-03-14 12:24:34 -0700915 try:
916 all_calls.remove(kall)
917 except ValueError:
918 not_found.append(kall)
919 if not_found:
920 raise AssertionError(
davidair2b32da22018-08-17 15:09:58 -0400921 '%r does not contain all of %r in its call list, '
922 'found %r instead' % (self._mock_name or 'mock',
923 tuple(not_found), all_calls)
Antoine Pitrou5c64df72013-02-03 00:23:58 +0100924 ) from cause
Michael Foord345266a2012-03-14 12:24:34 -0700925
926
927 def assert_any_call(self, *args, **kwargs):
928 """assert the mock has been called with the specified arguments.
929
930 The assert passes if the mock has *ever* been called, unlike
931 `assert_called_with` and `assert_called_once_with` that only pass if
932 the call is the most recent one."""
Antoine Pitrou5c64df72013-02-03 00:23:58 +0100933 expected = self._call_matcher((args, kwargs))
934 actual = [self._call_matcher(c) for c in self.call_args_list]
935 if expected not in actual:
936 cause = expected if isinstance(expected, Exception) else None
Michael Foord345266a2012-03-14 12:24:34 -0700937 expected_string = self._format_mock_call_signature(args, kwargs)
938 raise AssertionError(
939 '%s call not found' % expected_string
Antoine Pitrou5c64df72013-02-03 00:23:58 +0100940 ) from cause
Michael Foord345266a2012-03-14 12:24:34 -0700941
942
943 def _get_child_mock(self, **kw):
944 """Create the child mocks for attributes and return value.
945 By default child mocks will be the same type as the parent.
946 Subclasses of Mock may want to override this to customize the way
947 child mocks are made.
948
949 For non-callable mocks the callable variant will be used (rather than
950 any custom subclass)."""
Lisa Roach77b3b772019-05-20 09:19:53 -0700951 _new_name = kw.get("_new_name")
952 if _new_name in self.__dict__['_spec_asyncs']:
953 return AsyncMock(**kw)
954
Michael Foord345266a2012-03-14 12:24:34 -0700955 _type = type(self)
Lisa Roach77b3b772019-05-20 09:19:53 -0700956 if issubclass(_type, MagicMock) and _new_name in _async_method_magics:
957 klass = AsyncMock
958 if issubclass(_type, AsyncMockMixin):
959 klass = MagicMock
Michael Foord345266a2012-03-14 12:24:34 -0700960 if not issubclass(_type, CallableMixin):
961 if issubclass(_type, NonCallableMagicMock):
962 klass = MagicMock
963 elif issubclass(_type, NonCallableMock) :
964 klass = Mock
965 else:
966 klass = _type.__mro__[1]
Mario Corchero552be9d2017-10-17 12:35:11 +0100967
968 if self._mock_sealed:
969 attribute = "." + kw["name"] if "name" in kw else "()"
970 mock_name = self._extract_mock_name() + attribute
971 raise AttributeError(mock_name)
972
Michael Foord345266a2012-03-14 12:24:34 -0700973 return klass(**kw)
974
975
Petter Strandmark47d94242018-10-28 21:37:10 +0100976 def _calls_repr(self, prefix="Calls"):
977 """Renders self.mock_calls as a string.
978
979 Example: "\nCalls: [call(1), call(2)]."
980
981 If self.mock_calls is empty, an empty string is returned. The
982 output will be truncated if very long.
983 """
984 if not self.mock_calls:
985 return ""
986 return f"\n{prefix}: {safe_repr(self.mock_calls)}."
987
988
Michael Foord345266a2012-03-14 12:24:34 -0700989
990def _try_iter(obj):
991 if obj is None:
992 return obj
993 if _is_exception(obj):
994 return obj
995 if _callable(obj):
996 return obj
997 try:
998 return iter(obj)
999 except TypeError:
1000 # XXXX backwards compatibility
1001 # but this will blow up on first call - so maybe we should fail early?
1002 return obj
1003
1004
Michael Foord345266a2012-03-14 12:24:34 -07001005class CallableMixin(Base):
1006
1007 def __init__(self, spec=None, side_effect=None, return_value=DEFAULT,
1008 wraps=None, name=None, spec_set=None, parent=None,
1009 _spec_state=None, _new_name='', _new_parent=None, **kwargs):
1010 self.__dict__['_mock_return_value'] = return_value
Nick Coghlan0b43bcf2012-05-27 18:17:07 +10001011 _safe_super(CallableMixin, self).__init__(
Michael Foord345266a2012-03-14 12:24:34 -07001012 spec, wraps, name, spec_set, parent,
1013 _spec_state, _new_name, _new_parent, **kwargs
1014 )
1015
1016 self.side_effect = side_effect
1017
1018
1019 def _mock_check_sig(self, *args, **kwargs):
1020 # stub method that can be replaced with one with a specific signature
1021 pass
1022
1023
1024 def __call__(_mock_self, *args, **kwargs):
1025 # can't use self in-case a function / method we are mocking uses self
1026 # in the signature
1027 _mock_self._mock_check_sig(*args, **kwargs)
1028 return _mock_self._mock_call(*args, **kwargs)
1029
1030
1031 def _mock_call(_mock_self, *args, **kwargs):
1032 self = _mock_self
1033 self.called = True
1034 self.call_count += 1
Antoine Pitrou5c64df72013-02-03 00:23:58 +01001035
Chris Withers8ca0fa92018-12-03 21:31:37 +00001036 # handle call_args
Antoine Pitrou5c64df72013-02-03 00:23:58 +01001037 _call = _Call((args, kwargs), two=True)
1038 self.call_args = _call
1039 self.call_args_list.append(_call)
Michael Foord345266a2012-03-14 12:24:34 -07001040
Chris Withers8ca0fa92018-12-03 21:31:37 +00001041 # initial stuff for method_calls:
Michael Foord345266a2012-03-14 12:24:34 -07001042 do_method_calls = self._mock_parent is not None
Chris Withers8ca0fa92018-12-03 21:31:37 +00001043 method_call_name = self._mock_name
1044
1045 # initial stuff for mock_calls:
1046 mock_call_name = self._mock_new_name
1047 is_a_call = mock_call_name == '()'
1048 self.mock_calls.append(_Call(('', args, kwargs)))
1049
1050 # follow up the chain of mocks:
1051 _new_parent = self._mock_new_parent
Michael Foord345266a2012-03-14 12:24:34 -07001052 while _new_parent is not None:
Michael Foord345266a2012-03-14 12:24:34 -07001053
Chris Withers8ca0fa92018-12-03 21:31:37 +00001054 # handle method_calls:
Michael Foord345266a2012-03-14 12:24:34 -07001055 if do_method_calls:
Chris Withers8ca0fa92018-12-03 21:31:37 +00001056 _new_parent.method_calls.append(_Call((method_call_name, args, kwargs)))
Michael Foord345266a2012-03-14 12:24:34 -07001057 do_method_calls = _new_parent._mock_parent is not None
1058 if do_method_calls:
Chris Withers8ca0fa92018-12-03 21:31:37 +00001059 method_call_name = _new_parent._mock_name + '.' + method_call_name
Michael Foord345266a2012-03-14 12:24:34 -07001060
Chris Withers8ca0fa92018-12-03 21:31:37 +00001061 # handle mock_calls:
1062 this_mock_call = _Call((mock_call_name, args, kwargs))
Michael Foord345266a2012-03-14 12:24:34 -07001063 _new_parent.mock_calls.append(this_mock_call)
Chris Withers8ca0fa92018-12-03 21:31:37 +00001064
1065 if _new_parent._mock_new_name:
1066 if is_a_call:
1067 dot = ''
1068 else:
1069 dot = '.'
1070 is_a_call = _new_parent._mock_new_name == '()'
1071 mock_call_name = _new_parent._mock_new_name + dot + mock_call_name
1072
1073 # follow the parental chain:
Michael Foord345266a2012-03-14 12:24:34 -07001074 _new_parent = _new_parent._mock_new_parent
1075
Michael Foord345266a2012-03-14 12:24:34 -07001076 effect = self.side_effect
1077 if effect is not None:
1078 if _is_exception(effect):
1079 raise effect
Mario Corcherof05df0a2018-12-08 11:25:02 +00001080 elif not _callable(effect):
Michael Foord2cd48732012-04-21 15:52:11 +01001081 result = next(effect)
1082 if _is_exception(result):
1083 raise result
Mario Corcherof05df0a2018-12-08 11:25:02 +00001084 else:
1085 result = effect(*args, **kwargs)
1086
1087 if result is not DEFAULT:
Michael Foord2cd48732012-04-21 15:52:11 +01001088 return result
Michael Foord345266a2012-03-14 12:24:34 -07001089
Mario Corcherof05df0a2018-12-08 11:25:02 +00001090 if self._mock_return_value is not DEFAULT:
1091 return self.return_value
Michael Foord345266a2012-03-14 12:24:34 -07001092
Mario Corcherof05df0a2018-12-08 11:25:02 +00001093 if self._mock_wraps is not None:
Michael Foord345266a2012-03-14 12:24:34 -07001094 return self._mock_wraps(*args, **kwargs)
Mario Corcherof05df0a2018-12-08 11:25:02 +00001095
1096 return self.return_value
Michael Foord345266a2012-03-14 12:24:34 -07001097
1098
1099
1100class Mock(CallableMixin, NonCallableMock):
1101 """
1102 Create a new `Mock` object. `Mock` takes several optional arguments
1103 that specify the behaviour of the Mock object:
1104
1105 * `spec`: This can be either a list of strings or an existing object (a
1106 class or instance) that acts as the specification for the mock object. If
1107 you pass in an object then a list of strings is formed by calling dir on
1108 the object (excluding unsupported magic attributes and methods). Accessing
1109 any attribute not in this list will raise an `AttributeError`.
1110
1111 If `spec` is an object (rather than a list of strings) then
1112 `mock.__class__` returns the class of the spec object. This allows mocks
1113 to pass `isinstance` tests.
1114
1115 * `spec_set`: A stricter variant of `spec`. If used, attempting to *set*
1116 or get an attribute on the mock that isn't on the object passed as
1117 `spec_set` will raise an `AttributeError`.
1118
1119 * `side_effect`: A function to be called whenever the Mock is called. See
1120 the `side_effect` attribute. Useful for raising exceptions or
1121 dynamically changing return values. The function is called with the same
1122 arguments as the mock, and unless it returns `DEFAULT`, the return
1123 value of this function is used as the return value.
1124
Michael Foord2cd48732012-04-21 15:52:11 +01001125 If `side_effect` is an iterable then each call to the mock will return
1126 the next value from the iterable. If any of the members of the iterable
1127 are exceptions they will be raised instead of returned.
Michael Foord345266a2012-03-14 12:24:34 -07001128
Michael Foord345266a2012-03-14 12:24:34 -07001129 * `return_value`: The value returned when the mock is called. By default
1130 this is a new Mock (created on first access). See the
1131 `return_value` attribute.
1132
Michael Foord0682a0c2012-04-13 20:51:20 +01001133 * `wraps`: Item for the mock object to wrap. If `wraps` is not None then
1134 calling the Mock will pass the call through to the wrapped object
1135 (returning the real result). Attribute access on the mock will return a
1136 Mock object that wraps the corresponding attribute of the wrapped object
1137 (so attempting to access an attribute that doesn't exist will raise an
1138 `AttributeError`).
Michael Foord345266a2012-03-14 12:24:34 -07001139
1140 If the mock has an explicit `return_value` set then calls are not passed
1141 to the wrapped object and the `return_value` is returned instead.
1142
1143 * `name`: If the mock has a name then it will be used in the repr of the
1144 mock. This can be useful for debugging. The name is propagated to child
1145 mocks.
1146
1147 Mocks can also be called with arbitrary keyword arguments. These will be
1148 used to set attributes on the mock after it is created.
1149 """
1150
1151
Michael Foord345266a2012-03-14 12:24:34 -07001152def _dot_lookup(thing, comp, import_path):
1153 try:
1154 return getattr(thing, comp)
1155 except AttributeError:
1156 __import__(import_path)
1157 return getattr(thing, comp)
1158
1159
1160def _importer(target):
1161 components = target.split('.')
1162 import_path = components.pop(0)
1163 thing = __import__(import_path)
1164
1165 for comp in components:
1166 import_path += ".%s" % comp
1167 thing = _dot_lookup(thing, comp, import_path)
1168 return thing
1169
1170
1171def _is_started(patcher):
1172 # XXXX horrible
1173 return hasattr(patcher, 'is_local')
1174
1175
1176class _patch(object):
1177
1178 attribute_name = None
Michael Foordebc1a302014-04-15 17:21:08 -04001179 _active_patches = []
Michael Foord345266a2012-03-14 12:24:34 -07001180
1181 def __init__(
1182 self, getter, attribute, new, spec, create,
1183 spec_set, autospec, new_callable, kwargs
1184 ):
1185 if new_callable is not None:
1186 if new is not DEFAULT:
1187 raise ValueError(
1188 "Cannot use 'new' and 'new_callable' together"
1189 )
Michael Foord50a8c0e2012-03-25 18:57:58 +01001190 if autospec is not None:
Michael Foord345266a2012-03-14 12:24:34 -07001191 raise ValueError(
1192 "Cannot use 'autospec' and 'new_callable' together"
1193 )
1194
1195 self.getter = getter
1196 self.attribute = attribute
1197 self.new = new
1198 self.new_callable = new_callable
1199 self.spec = spec
1200 self.create = create
1201 self.has_local = False
1202 self.spec_set = spec_set
1203 self.autospec = autospec
1204 self.kwargs = kwargs
1205 self.additional_patchers = []
1206
1207
1208 def copy(self):
1209 patcher = _patch(
1210 self.getter, self.attribute, self.new, self.spec,
1211 self.create, self.spec_set,
1212 self.autospec, self.new_callable, self.kwargs
1213 )
1214 patcher.attribute_name = self.attribute_name
1215 patcher.additional_patchers = [
1216 p.copy() for p in self.additional_patchers
1217 ]
1218 return patcher
1219
1220
1221 def __call__(self, func):
1222 if isinstance(func, type):
1223 return self.decorate_class(func)
Xtreak436c2b02019-05-28 12:37:39 +05301224 if inspect.iscoroutinefunction(func):
1225 return self.decorate_async_callable(func)
Michael Foord345266a2012-03-14 12:24:34 -07001226 return self.decorate_callable(func)
1227
1228
1229 def decorate_class(self, klass):
1230 for attr in dir(klass):
1231 if not attr.startswith(patch.TEST_PREFIX):
1232 continue
1233
1234 attr_value = getattr(klass, attr)
1235 if not hasattr(attr_value, "__call__"):
1236 continue
1237
1238 patcher = self.copy()
1239 setattr(klass, attr, patcher(attr_value))
1240 return klass
1241
1242
Xtreak436c2b02019-05-28 12:37:39 +05301243 @contextlib.contextmanager
1244 def decoration_helper(self, patched, args, keywargs):
1245 extra_args = []
1246 entered_patchers = []
1247 patching = None
1248
1249 exc_info = tuple()
1250 try:
1251 for patching in patched.patchings:
1252 arg = patching.__enter__()
1253 entered_patchers.append(patching)
1254 if patching.attribute_name is not None:
1255 keywargs.update(arg)
1256 elif patching.new is DEFAULT:
1257 extra_args.append(arg)
1258
1259 args += tuple(extra_args)
1260 yield (args, keywargs)
1261 except:
1262 if (patching not in entered_patchers and
1263 _is_started(patching)):
1264 # the patcher may have been started, but an exception
1265 # raised whilst entering one of its additional_patchers
1266 entered_patchers.append(patching)
1267 # Pass the exception to __exit__
1268 exc_info = sys.exc_info()
1269 # re-raise the exception
1270 raise
1271 finally:
1272 for patching in reversed(entered_patchers):
1273 patching.__exit__(*exc_info)
1274
1275
Michael Foord345266a2012-03-14 12:24:34 -07001276 def decorate_callable(self, func):
Xtreak436c2b02019-05-28 12:37:39 +05301277 # NB. Keep the method in sync with decorate_async_callable()
Michael Foord345266a2012-03-14 12:24:34 -07001278 if hasattr(func, 'patchings'):
1279 func.patchings.append(self)
1280 return func
1281
1282 @wraps(func)
1283 def patched(*args, **keywargs):
Xtreak436c2b02019-05-28 12:37:39 +05301284 with self.decoration_helper(patched,
1285 args,
1286 keywargs) as (newargs, newkeywargs):
1287 return func(*newargs, **newkeywargs)
Michael Foord345266a2012-03-14 12:24:34 -07001288
Xtreak436c2b02019-05-28 12:37:39 +05301289 patched.patchings = [self]
1290 return patched
Michael Foord345266a2012-03-14 12:24:34 -07001291
Xtreak436c2b02019-05-28 12:37:39 +05301292
1293 def decorate_async_callable(self, func):
1294 # NB. Keep the method in sync with decorate_callable()
1295 if hasattr(func, 'patchings'):
1296 func.patchings.append(self)
1297 return func
1298
1299 @wraps(func)
1300 async def patched(*args, **keywargs):
1301 with self.decoration_helper(patched,
1302 args,
1303 keywargs) as (newargs, newkeywargs):
1304 return await func(*newargs, **newkeywargs)
Michael Foord345266a2012-03-14 12:24:34 -07001305
1306 patched.patchings = [self]
Michael Foord345266a2012-03-14 12:24:34 -07001307 return patched
1308
1309
1310 def get_original(self):
1311 target = self.getter()
1312 name = self.attribute
1313
1314 original = DEFAULT
1315 local = False
1316
1317 try:
1318 original = target.__dict__[name]
1319 except (AttributeError, KeyError):
1320 original = getattr(target, name, DEFAULT)
1321 else:
1322 local = True
1323
Michael Foordfddcfa22014-04-14 16:25:20 -04001324 if name in _builtins and isinstance(target, ModuleType):
1325 self.create = True
1326
Michael Foord345266a2012-03-14 12:24:34 -07001327 if not self.create and original is DEFAULT:
1328 raise AttributeError(
1329 "%s does not have the attribute %r" % (target, name)
1330 )
1331 return original, local
1332
1333
1334 def __enter__(self):
1335 """Perform the patch."""
1336 new, spec, spec_set = self.new, self.spec, self.spec_set
1337 autospec, kwargs = self.autospec, self.kwargs
1338 new_callable = self.new_callable
1339 self.target = self.getter()
1340
Michael Foord50a8c0e2012-03-25 18:57:58 +01001341 # normalise False to None
1342 if spec is False:
1343 spec = None
1344 if spec_set is False:
1345 spec_set = None
1346 if autospec is False:
1347 autospec = None
1348
1349 if spec is not None and autospec is not None:
1350 raise TypeError("Can't specify spec and autospec")
1351 if ((spec is not None or autospec is not None) and
1352 spec_set not in (True, None)):
1353 raise TypeError("Can't provide explicit spec_set *and* spec or autospec")
1354
Michael Foord345266a2012-03-14 12:24:34 -07001355 original, local = self.get_original()
1356
Michael Foord50a8c0e2012-03-25 18:57:58 +01001357 if new is DEFAULT and autospec is None:
Michael Foord345266a2012-03-14 12:24:34 -07001358 inherit = False
Michael Foord50a8c0e2012-03-25 18:57:58 +01001359 if spec is True:
Michael Foord345266a2012-03-14 12:24:34 -07001360 # set spec to the object we are replacing
1361 spec = original
Michael Foord50a8c0e2012-03-25 18:57:58 +01001362 if spec_set is True:
1363 spec_set = original
1364 spec = None
1365 elif spec is not None:
1366 if spec_set is True:
1367 spec_set = spec
1368 spec = None
1369 elif spec_set is True:
1370 spec_set = original
Michael Foord345266a2012-03-14 12:24:34 -07001371
Michael Foord50a8c0e2012-03-25 18:57:58 +01001372 if spec is not None or spec_set is not None:
1373 if original is DEFAULT:
1374 raise TypeError("Can't use 'spec' with create=True")
Michael Foord345266a2012-03-14 12:24:34 -07001375 if isinstance(original, type):
1376 # If we're patching out a class and there is a spec
1377 inherit = True
Lisa Roach77b3b772019-05-20 09:19:53 -07001378 if spec is None and _is_async_obj(original):
1379 Klass = AsyncMock
1380 else:
1381 Klass = MagicMock
Michael Foord345266a2012-03-14 12:24:34 -07001382 _kwargs = {}
1383 if new_callable is not None:
1384 Klass = new_callable
Michael Foord50a8c0e2012-03-25 18:57:58 +01001385 elif spec is not None or spec_set is not None:
Michael Foorde58a5622012-03-25 19:53:18 +01001386 this_spec = spec
1387 if spec_set is not None:
1388 this_spec = spec_set
1389 if _is_list(this_spec):
1390 not_callable = '__call__' not in this_spec
1391 else:
1392 not_callable = not callable(this_spec)
Lisa Roach77b3b772019-05-20 09:19:53 -07001393 if _is_async_obj(this_spec):
1394 Klass = AsyncMock
1395 elif not_callable:
Michael Foord345266a2012-03-14 12:24:34 -07001396 Klass = NonCallableMagicMock
1397
1398 if spec is not None:
1399 _kwargs['spec'] = spec
1400 if spec_set is not None:
1401 _kwargs['spec_set'] = spec_set
1402
1403 # add a name to mocks
1404 if (isinstance(Klass, type) and
1405 issubclass(Klass, NonCallableMock) and self.attribute):
1406 _kwargs['name'] = self.attribute
1407
1408 _kwargs.update(kwargs)
1409 new = Klass(**_kwargs)
1410
1411 if inherit and _is_instance_mock(new):
1412 # we can only tell if the instance should be callable if the
1413 # spec is not a list
Michael Foord50a8c0e2012-03-25 18:57:58 +01001414 this_spec = spec
1415 if spec_set is not None:
1416 this_spec = spec_set
1417 if (not _is_list(this_spec) and not
1418 _instance_callable(this_spec)):
Michael Foord345266a2012-03-14 12:24:34 -07001419 Klass = NonCallableMagicMock
1420
1421 _kwargs.pop('name')
1422 new.return_value = Klass(_new_parent=new, _new_name='()',
1423 **_kwargs)
Michael Foord50a8c0e2012-03-25 18:57:58 +01001424 elif autospec is not None:
Michael Foord345266a2012-03-14 12:24:34 -07001425 # spec is ignored, new *must* be default, spec_set is treated
1426 # as a boolean. Should we check spec is not None and that spec_set
1427 # is a bool?
1428 if new is not DEFAULT:
1429 raise TypeError(
1430 "autospec creates the mock for you. Can't specify "
1431 "autospec and new."
1432 )
Michael Foord50a8c0e2012-03-25 18:57:58 +01001433 if original is DEFAULT:
Michael Foord99254732012-03-25 19:07:33 +01001434 raise TypeError("Can't use 'autospec' with create=True")
Michael Foord345266a2012-03-14 12:24:34 -07001435 spec_set = bool(spec_set)
1436 if autospec is True:
1437 autospec = original
1438
1439 new = create_autospec(autospec, spec_set=spec_set,
1440 _name=self.attribute, **kwargs)
1441 elif kwargs:
1442 # can't set keyword args when we aren't creating the mock
1443 # XXXX If new is a Mock we could call new.configure_mock(**kwargs)
1444 raise TypeError("Can't pass kwargs to a mock we aren't creating")
1445
1446 new_attr = new
1447
1448 self.temp_original = original
1449 self.is_local = local
1450 setattr(self.target, self.attribute, new_attr)
1451 if self.attribute_name is not None:
1452 extra_args = {}
1453 if self.new is DEFAULT:
1454 extra_args[self.attribute_name] = new
1455 for patching in self.additional_patchers:
1456 arg = patching.__enter__()
1457 if patching.new is DEFAULT:
1458 extra_args.update(arg)
1459 return extra_args
1460
1461 return new
1462
1463
Michael Foord50a8c0e2012-03-25 18:57:58 +01001464 def __exit__(self, *exc_info):
Michael Foord345266a2012-03-14 12:24:34 -07001465 """Undo the patch."""
1466 if not _is_started(self):
Xtreak02b84cb2019-03-29 02:38:43 +05301467 return
Michael Foord345266a2012-03-14 12:24:34 -07001468
1469 if self.is_local and self.temp_original is not DEFAULT:
1470 setattr(self.target, self.attribute, self.temp_original)
1471 else:
1472 delattr(self.target, self.attribute)
Senthil Kumaran81bc9272016-01-08 23:43:29 -08001473 if not self.create and (not hasattr(self.target, self.attribute) or
1474 self.attribute in ('__doc__', '__module__',
1475 '__defaults__', '__annotations__',
1476 '__kwdefaults__')):
Michael Foord345266a2012-03-14 12:24:34 -07001477 # needed for proxy objects like django settings
1478 setattr(self.target, self.attribute, self.temp_original)
1479
1480 del self.temp_original
1481 del self.is_local
1482 del self.target
1483 for patcher in reversed(self.additional_patchers):
1484 if _is_started(patcher):
Michael Foord50a8c0e2012-03-25 18:57:58 +01001485 patcher.__exit__(*exc_info)
Michael Foord345266a2012-03-14 12:24:34 -07001486
Michael Foordf7c41582012-06-10 20:36:32 +01001487
1488 def start(self):
1489 """Activate a patch, returning any created mock."""
1490 result = self.__enter__()
Michael Foordebc1a302014-04-15 17:21:08 -04001491 self._active_patches.append(self)
Michael Foordf7c41582012-06-10 20:36:32 +01001492 return result
1493
1494
1495 def stop(self):
1496 """Stop an active patch."""
Michael Foordebc1a302014-04-15 17:21:08 -04001497 try:
1498 self._active_patches.remove(self)
1499 except ValueError:
1500 # If the patch hasn't been started this will fail
1501 pass
1502
Michael Foordf7c41582012-06-10 20:36:32 +01001503 return self.__exit__()
Michael Foord345266a2012-03-14 12:24:34 -07001504
1505
1506
1507def _get_target(target):
1508 try:
1509 target, attribute = target.rsplit('.', 1)
1510 except (TypeError, ValueError):
1511 raise TypeError("Need a valid target to patch. You supplied: %r" %
1512 (target,))
1513 getter = lambda: _importer(target)
1514 return getter, attribute
1515
1516
1517def _patch_object(
1518 target, attribute, new=DEFAULT, spec=None,
Michael Foord50a8c0e2012-03-25 18:57:58 +01001519 create=False, spec_set=None, autospec=None,
Michael Foord345266a2012-03-14 12:24:34 -07001520 new_callable=None, **kwargs
1521 ):
1522 """
Michael Foord345266a2012-03-14 12:24:34 -07001523 patch the named member (`attribute`) on an object (`target`) with a mock
1524 object.
1525
1526 `patch.object` can be used as a decorator, class decorator or a context
1527 manager. Arguments `new`, `spec`, `create`, `spec_set`,
1528 `autospec` and `new_callable` have the same meaning as for `patch`. Like
1529 `patch`, `patch.object` takes arbitrary keyword arguments for configuring
1530 the mock object it creates.
1531
1532 When used as a class decorator `patch.object` honours `patch.TEST_PREFIX`
1533 for choosing which methods to wrap.
1534 """
1535 getter = lambda: target
1536 return _patch(
1537 getter, attribute, new, spec, create,
1538 spec_set, autospec, new_callable, kwargs
1539 )
1540
1541
Michael Foord50a8c0e2012-03-25 18:57:58 +01001542def _patch_multiple(target, spec=None, create=False, spec_set=None,
1543 autospec=None, new_callable=None, **kwargs):
Michael Foord345266a2012-03-14 12:24:34 -07001544 """Perform multiple patches in a single call. It takes the object to be
1545 patched (either as an object or a string to fetch the object by importing)
1546 and keyword arguments for the patches::
1547
1548 with patch.multiple(settings, FIRST_PATCH='one', SECOND_PATCH='two'):
1549 ...
1550
1551 Use `DEFAULT` as the value if you want `patch.multiple` to create
1552 mocks for you. In this case the created mocks are passed into a decorated
1553 function by keyword, and a dictionary is returned when `patch.multiple` is
1554 used as a context manager.
1555
1556 `patch.multiple` can be used as a decorator, class decorator or a context
1557 manager. The arguments `spec`, `spec_set`, `create`,
1558 `autospec` and `new_callable` have the same meaning as for `patch`. These
1559 arguments will be applied to *all* patches done by `patch.multiple`.
1560
1561 When used as a class decorator `patch.multiple` honours `patch.TEST_PREFIX`
1562 for choosing which methods to wrap.
1563 """
1564 if type(target) is str:
1565 getter = lambda: _importer(target)
1566 else:
1567 getter = lambda: target
1568
1569 if not kwargs:
1570 raise ValueError(
1571 'Must supply at least one keyword argument with patch.multiple'
1572 )
1573 # need to wrap in a list for python 3, where items is a view
1574 items = list(kwargs.items())
1575 attribute, new = items[0]
1576 patcher = _patch(
1577 getter, attribute, new, spec, create, spec_set,
1578 autospec, new_callable, {}
1579 )
1580 patcher.attribute_name = attribute
1581 for attribute, new in items[1:]:
1582 this_patcher = _patch(
1583 getter, attribute, new, spec, create, spec_set,
1584 autospec, new_callable, {}
1585 )
1586 this_patcher.attribute_name = attribute
1587 patcher.additional_patchers.append(this_patcher)
1588 return patcher
1589
1590
1591def patch(
1592 target, new=DEFAULT, spec=None, create=False,
Michael Foord50a8c0e2012-03-25 18:57:58 +01001593 spec_set=None, autospec=None, new_callable=None, **kwargs
Michael Foord345266a2012-03-14 12:24:34 -07001594 ):
1595 """
1596 `patch` acts as a function decorator, class decorator or a context
1597 manager. Inside the body of the function or with statement, the `target`
Michael Foord54b3db82012-03-28 15:08:08 +01001598 is patched with a `new` object. When the function/with statement exits
1599 the patch is undone.
Michael Foord345266a2012-03-14 12:24:34 -07001600
Michael Foord54b3db82012-03-28 15:08:08 +01001601 If `new` is omitted, then the target is replaced with a
1602 `MagicMock`. If `patch` is used as a decorator and `new` is
1603 omitted, the created mock is passed in as an extra argument to the
1604 decorated function. If `patch` is used as a context manager the created
1605 mock is returned by the context manager.
Michael Foord345266a2012-03-14 12:24:34 -07001606
Michael Foord54b3db82012-03-28 15:08:08 +01001607 `target` should be a string in the form `'package.module.ClassName'`. The
1608 `target` is imported and the specified object replaced with the `new`
1609 object, so the `target` must be importable from the environment you are
1610 calling `patch` from. The target is imported when the decorated function
1611 is executed, not at decoration time.
Michael Foord345266a2012-03-14 12:24:34 -07001612
1613 The `spec` and `spec_set` keyword arguments are passed to the `MagicMock`
1614 if patch is creating one for you.
1615
1616 In addition you can pass `spec=True` or `spec_set=True`, which causes
1617 patch to pass in the object being mocked as the spec/spec_set object.
1618
1619 `new_callable` allows you to specify a different class, or callable object,
1620 that will be called to create the `new` object. By default `MagicMock` is
1621 used.
1622
1623 A more powerful form of `spec` is `autospec`. If you set `autospec=True`
Robert Collins92b3e0652015-07-17 21:58:36 +12001624 then the mock will be created with a spec from the object being replaced.
Michael Foord345266a2012-03-14 12:24:34 -07001625 All attributes of the mock will also have the spec of the corresponding
1626 attribute of the object being replaced. Methods and functions being
1627 mocked will have their arguments checked and will raise a `TypeError` if
1628 they are called with the wrong signature. For mocks replacing a class,
1629 their return value (the 'instance') will have the same spec as the class.
1630
1631 Instead of `autospec=True` you can pass `autospec=some_object` to use an
1632 arbitrary object as the spec instead of the one being replaced.
1633
1634 By default `patch` will fail to replace attributes that don't exist. If
1635 you pass in `create=True`, and the attribute doesn't exist, patch will
1636 create the attribute for you when the patched function is called, and
1637 delete it again afterwards. This is useful for writing tests against
Terry Jan Reedy0f847642013-03-11 18:34:00 -04001638 attributes that your production code creates at runtime. It is off by
Michael Foord345266a2012-03-14 12:24:34 -07001639 default because it can be dangerous. With it switched on you can write
1640 passing tests against APIs that don't actually exist!
1641
1642 Patch can be used as a `TestCase` class decorator. It works by
1643 decorating each test method in the class. This reduces the boilerplate
1644 code when your test methods share a common patchings set. `patch` finds
1645 tests by looking for method names that start with `patch.TEST_PREFIX`.
1646 By default this is `test`, which matches the way `unittest` finds tests.
1647 You can specify an alternative prefix by setting `patch.TEST_PREFIX`.
1648
1649 Patch can be used as a context manager, with the with statement. Here the
1650 patching applies to the indented block after the with statement. If you
1651 use "as" then the patched object will be bound to the name after the
1652 "as"; very useful if `patch` is creating a mock object for you.
1653
1654 `patch` takes arbitrary keyword arguments. These will be passed to
1655 the `Mock` (or `new_callable`) on construction.
1656
1657 `patch.dict(...)`, `patch.multiple(...)` and `patch.object(...)` are
1658 available for alternate use-cases.
1659 """
1660 getter, attribute = _get_target(target)
1661 return _patch(
1662 getter, attribute, new, spec, create,
1663 spec_set, autospec, new_callable, kwargs
1664 )
1665
1666
1667class _patch_dict(object):
1668 """
1669 Patch a dictionary, or dictionary like object, and restore the dictionary
1670 to its original state after the test.
1671
1672 `in_dict` can be a dictionary or a mapping like container. If it is a
1673 mapping then it must at least support getting, setting and deleting items
1674 plus iterating over keys.
1675
1676 `in_dict` can also be a string specifying the name of the dictionary, which
1677 will then be fetched by importing it.
1678
1679 `values` can be a dictionary of values to set in the dictionary. `values`
1680 can also be an iterable of `(key, value)` pairs.
1681
1682 If `clear` is True then the dictionary will be cleared before the new
1683 values are set.
1684
1685 `patch.dict` can also be called with arbitrary keyword arguments to set
1686 values in the dictionary::
1687
1688 with patch.dict('sys.modules', mymodule=Mock(), other_module=Mock()):
1689 ...
1690
1691 `patch.dict` can be used as a context manager, decorator or class
1692 decorator. When used as a class decorator `patch.dict` honours
1693 `patch.TEST_PREFIX` for choosing which methods to wrap.
1694 """
1695
1696 def __init__(self, in_dict, values=(), clear=False, **kwargs):
Michael Foord345266a2012-03-14 12:24:34 -07001697 self.in_dict = in_dict
1698 # support any argument supported by dict(...) constructor
1699 self.values = dict(values)
1700 self.values.update(kwargs)
1701 self.clear = clear
1702 self._original = None
1703
1704
1705 def __call__(self, f):
1706 if isinstance(f, type):
1707 return self.decorate_class(f)
1708 @wraps(f)
1709 def _inner(*args, **kw):
1710 self._patch_dict()
1711 try:
1712 return f(*args, **kw)
1713 finally:
1714 self._unpatch_dict()
1715
1716 return _inner
1717
1718
1719 def decorate_class(self, klass):
1720 for attr in dir(klass):
1721 attr_value = getattr(klass, attr)
1722 if (attr.startswith(patch.TEST_PREFIX) and
1723 hasattr(attr_value, "__call__")):
1724 decorator = _patch_dict(self.in_dict, self.values, self.clear)
1725 decorated = decorator(attr_value)
1726 setattr(klass, attr, decorated)
1727 return klass
1728
1729
1730 def __enter__(self):
1731 """Patch the dict."""
1732 self._patch_dict()
Mario Corchero04530812019-05-28 13:53:31 +01001733 return self.in_dict
Michael Foord345266a2012-03-14 12:24:34 -07001734
1735
1736 def _patch_dict(self):
1737 values = self.values
Xtreaka875ea52019-02-25 00:24:49 +05301738 if isinstance(self.in_dict, str):
1739 self.in_dict = _importer(self.in_dict)
Michael Foord345266a2012-03-14 12:24:34 -07001740 in_dict = self.in_dict
1741 clear = self.clear
1742
1743 try:
1744 original = in_dict.copy()
1745 except AttributeError:
1746 # dict like object with no copy method
1747 # must support iteration over keys
1748 original = {}
1749 for key in in_dict:
1750 original[key] = in_dict[key]
1751 self._original = original
1752
1753 if clear:
1754 _clear_dict(in_dict)
1755
1756 try:
1757 in_dict.update(values)
1758 except AttributeError:
1759 # dict like object with no update method
1760 for key in values:
1761 in_dict[key] = values[key]
1762
1763
1764 def _unpatch_dict(self):
1765 in_dict = self.in_dict
1766 original = self._original
1767
1768 _clear_dict(in_dict)
1769
1770 try:
1771 in_dict.update(original)
1772 except AttributeError:
1773 for key in original:
1774 in_dict[key] = original[key]
1775
1776
1777 def __exit__(self, *args):
1778 """Unpatch the dict."""
1779 self._unpatch_dict()
1780 return False
1781
1782 start = __enter__
1783 stop = __exit__
1784
1785
1786def _clear_dict(in_dict):
1787 try:
1788 in_dict.clear()
1789 except AttributeError:
1790 keys = list(in_dict)
1791 for key in keys:
1792 del in_dict[key]
1793
1794
Michael Foordf7c41582012-06-10 20:36:32 +01001795def _patch_stopall():
Michael Foordebc1a302014-04-15 17:21:08 -04001796 """Stop all active patches. LIFO to unroll nested patches."""
1797 for patch in reversed(_patch._active_patches):
Michael Foordf7c41582012-06-10 20:36:32 +01001798 patch.stop()
1799
1800
Michael Foord345266a2012-03-14 12:24:34 -07001801patch.object = _patch_object
1802patch.dict = _patch_dict
1803patch.multiple = _patch_multiple
Michael Foordf7c41582012-06-10 20:36:32 +01001804patch.stopall = _patch_stopall
Michael Foord345266a2012-03-14 12:24:34 -07001805patch.TEST_PREFIX = 'test'
1806
1807magic_methods = (
1808 "lt le gt ge eq ne "
1809 "getitem setitem delitem "
1810 "len contains iter "
1811 "hash str sizeof "
1812 "enter exit "
Berker Peksaga785dec2015-03-15 01:51:56 +02001813 # we added divmod and rdivmod here instead of numerics
1814 # because there is no idivmod
1815 "divmod rdivmod neg pos abs invert "
Michael Foord345266a2012-03-14 12:24:34 -07001816 "complex int float index "
John Reese6c4fab02018-05-22 13:01:10 -07001817 "round trunc floor ceil "
Michael Foord345266a2012-03-14 12:24:34 -07001818 "bool next "
Max Bélanger6c83d9f2018-10-25 14:48:58 -07001819 "fspath "
Michael Foord345266a2012-03-14 12:24:34 -07001820)
1821
Michael Foordd2623d72014-04-14 11:23:48 -04001822numerics = (
Berker Peksag9bd8af72015-03-12 20:42:48 +02001823 "add sub mul matmul div floordiv mod lshift rshift and xor or pow truediv"
Michael Foordd2623d72014-04-14 11:23:48 -04001824)
Michael Foord345266a2012-03-14 12:24:34 -07001825inplace = ' '.join('i%s' % n for n in numerics.split())
1826right = ' '.join('r%s' % n for n in numerics.split())
1827
1828# not including __prepare__, __instancecheck__, __subclasscheck__
1829# (as they are metaclass methods)
1830# __del__ is not supported at all as it causes problems if it exists
1831
Serhiy Storchakac02d1882014-12-11 10:28:14 +02001832_non_defaults = {
1833 '__get__', '__set__', '__delete__', '__reversed__', '__missing__',
1834 '__reduce__', '__reduce_ex__', '__getinitargs__', '__getnewargs__',
1835 '__getstate__', '__setstate__', '__getformat__', '__setformat__',
1836 '__repr__', '__dir__', '__subclasses__', '__format__',
Lisa Roach77b3b772019-05-20 09:19:53 -07001837 '__getnewargs_ex__', '__aenter__', '__aexit__', '__anext__', '__aiter__',
Victor Stinnereb1a9952014-12-11 22:25:49 +01001838}
Michael Foord345266a2012-03-14 12:24:34 -07001839
1840
1841def _get_method(name, func):
1842 "Turns a callable object (like a mock) into a real function"
1843 def method(self, *args, **kw):
1844 return func(self, *args, **kw)
1845 method.__name__ = name
1846 return method
1847
1848
Serhiy Storchakac02d1882014-12-11 10:28:14 +02001849_magics = {
Michael Foord345266a2012-03-14 12:24:34 -07001850 '__%s__' % method for method in
1851 ' '.join([magic_methods, numerics, inplace, right]).split()
Serhiy Storchakac02d1882014-12-11 10:28:14 +02001852}
Michael Foord345266a2012-03-14 12:24:34 -07001853
Lisa Roach77b3b772019-05-20 09:19:53 -07001854# Magic methods used for async `with` statements
1855_async_method_magics = {"__aenter__", "__aexit__", "__anext__"}
1856# `__aiter__` is a plain function but used with async calls
1857_async_magics = _async_method_magics | {"__aiter__"}
1858
Michael Foord345266a2012-03-14 12:24:34 -07001859_all_magics = _magics | _non_defaults
1860
Serhiy Storchakac02d1882014-12-11 10:28:14 +02001861_unsupported_magics = {
Michael Foord345266a2012-03-14 12:24:34 -07001862 '__getattr__', '__setattr__',
Serhiy Storchaka34fd4c22018-11-05 16:20:25 +02001863 '__init__', '__new__', '__prepare__',
Michael Foord345266a2012-03-14 12:24:34 -07001864 '__instancecheck__', '__subclasscheck__',
1865 '__del__'
Serhiy Storchakac02d1882014-12-11 10:28:14 +02001866}
Michael Foord345266a2012-03-14 12:24:34 -07001867
1868_calculate_return_value = {
1869 '__hash__': lambda self: object.__hash__(self),
1870 '__str__': lambda self: object.__str__(self),
1871 '__sizeof__': lambda self: object.__sizeof__(self),
Max Bélanger6c83d9f2018-10-25 14:48:58 -07001872 '__fspath__': lambda self: f"{type(self).__name__}/{self._extract_mock_name()}/{id(self)}",
Michael Foord345266a2012-03-14 12:24:34 -07001873}
1874
1875_return_values = {
Michael Foord313f85f2012-03-25 18:16:07 +01001876 '__lt__': NotImplemented,
1877 '__gt__': NotImplemented,
1878 '__le__': NotImplemented,
1879 '__ge__': NotImplemented,
Michael Foord345266a2012-03-14 12:24:34 -07001880 '__int__': 1,
1881 '__contains__': False,
1882 '__len__': 0,
1883 '__exit__': False,
1884 '__complex__': 1j,
1885 '__float__': 1.0,
1886 '__bool__': True,
1887 '__index__': 1,
Lisa Roach77b3b772019-05-20 09:19:53 -07001888 '__aexit__': False,
Michael Foord345266a2012-03-14 12:24:34 -07001889}
1890
1891
1892def _get_eq(self):
1893 def __eq__(other):
1894 ret_val = self.__eq__._mock_return_value
1895 if ret_val is not DEFAULT:
1896 return ret_val
Serhiy Storchaka362f0582017-01-21 23:12:58 +02001897 if self is other:
1898 return True
1899 return NotImplemented
Michael Foord345266a2012-03-14 12:24:34 -07001900 return __eq__
1901
1902def _get_ne(self):
1903 def __ne__(other):
1904 if self.__ne__._mock_return_value is not DEFAULT:
1905 return DEFAULT
Serhiy Storchaka362f0582017-01-21 23:12:58 +02001906 if self is other:
1907 return False
1908 return NotImplemented
Michael Foord345266a2012-03-14 12:24:34 -07001909 return __ne__
1910
1911def _get_iter(self):
1912 def __iter__():
1913 ret_val = self.__iter__._mock_return_value
1914 if ret_val is DEFAULT:
1915 return iter([])
1916 # if ret_val was already an iterator, then calling iter on it should
1917 # return the iterator unchanged
1918 return iter(ret_val)
1919 return __iter__
1920
Lisa Roach77b3b772019-05-20 09:19:53 -07001921def _get_async_iter(self):
1922 def __aiter__():
1923 ret_val = self.__aiter__._mock_return_value
1924 if ret_val is DEFAULT:
1925 return _AsyncIterator(iter([]))
1926 return _AsyncIterator(iter(ret_val))
1927 return __aiter__
1928
Michael Foord345266a2012-03-14 12:24:34 -07001929_side_effect_methods = {
1930 '__eq__': _get_eq,
1931 '__ne__': _get_ne,
1932 '__iter__': _get_iter,
Lisa Roach77b3b772019-05-20 09:19:53 -07001933 '__aiter__': _get_async_iter
Michael Foord345266a2012-03-14 12:24:34 -07001934}
1935
1936
1937
1938def _set_return_value(mock, method, name):
1939 fixed = _return_values.get(name, DEFAULT)
1940 if fixed is not DEFAULT:
1941 method.return_value = fixed
1942 return
1943
1944 return_calulator = _calculate_return_value.get(name)
1945 if return_calulator is not None:
Chris Withersadbf1782019-05-01 23:04:04 +01001946 return_value = return_calulator(mock)
Michael Foord345266a2012-03-14 12:24:34 -07001947 method.return_value = return_value
1948 return
1949
1950 side_effector = _side_effect_methods.get(name)
1951 if side_effector is not None:
1952 method.side_effect = side_effector(mock)
1953
1954
1955
1956class MagicMixin(object):
1957 def __init__(self, *args, **kw):
Łukasz Langaa468db92015-04-13 23:12:42 -07001958 self._mock_set_magics() # make magic work for kwargs in init
Nick Coghlan0b43bcf2012-05-27 18:17:07 +10001959 _safe_super(MagicMixin, self).__init__(*args, **kw)
Łukasz Langaa468db92015-04-13 23:12:42 -07001960 self._mock_set_magics() # fix magic broken by upper level init
Michael Foord345266a2012-03-14 12:24:34 -07001961
1962
1963 def _mock_set_magics(self):
1964 these_magics = _magics
1965
Łukasz Langaa468db92015-04-13 23:12:42 -07001966 if getattr(self, "_mock_methods", None) is not None:
Michael Foord345266a2012-03-14 12:24:34 -07001967 these_magics = _magics.intersection(self._mock_methods)
1968
1969 remove_magics = set()
1970 remove_magics = _magics - these_magics
1971
1972 for entry in remove_magics:
1973 if entry in type(self).__dict__:
1974 # remove unneeded magic methods
1975 delattr(self, entry)
1976
1977 # don't overwrite existing attributes if called a second time
1978 these_magics = these_magics - set(type(self).__dict__)
1979
1980 _type = type(self)
1981 for entry in these_magics:
1982 setattr(_type, entry, MagicProxy(entry, self))
1983
1984
1985
1986class NonCallableMagicMock(MagicMixin, NonCallableMock):
1987 """A version of `MagicMock` that isn't callable."""
1988 def mock_add_spec(self, spec, spec_set=False):
1989 """Add a spec to a mock. `spec` can either be an object or a
1990 list of strings. Only attributes on the `spec` can be fetched as
1991 attributes from the mock.
1992
1993 If `spec_set` is True then only attributes on the spec can be set."""
1994 self._mock_add_spec(spec, spec_set)
1995 self._mock_set_magics()
1996
1997
Lisa Roach77b3b772019-05-20 09:19:53 -07001998class AsyncMagicMixin:
1999 def __init__(self, *args, **kw):
2000 self._mock_set_async_magics() # make magic work for kwargs in init
2001 _safe_super(AsyncMagicMixin, self).__init__(*args, **kw)
2002 self._mock_set_async_magics() # fix magic broken by upper level init
Michael Foord345266a2012-03-14 12:24:34 -07002003
Lisa Roach77b3b772019-05-20 09:19:53 -07002004 def _mock_set_async_magics(self):
2005 these_magics = _async_magics
2006
2007 if getattr(self, "_mock_methods", None) is not None:
2008 these_magics = _async_magics.intersection(self._mock_methods)
2009 remove_magics = _async_magics - these_magics
2010
2011 for entry in remove_magics:
2012 if entry in type(self).__dict__:
2013 # remove unneeded magic methods
2014 delattr(self, entry)
2015
2016 # don't overwrite existing attributes if called a second time
2017 these_magics = these_magics - set(type(self).__dict__)
2018
2019 _type = type(self)
2020 for entry in these_magics:
2021 setattr(_type, entry, MagicProxy(entry, self))
2022
2023
2024class MagicMock(MagicMixin, AsyncMagicMixin, Mock):
Michael Foord345266a2012-03-14 12:24:34 -07002025 """
2026 MagicMock is a subclass of Mock with default implementations
2027 of most of the magic methods. You can use MagicMock without having to
2028 configure the magic methods yourself.
2029
2030 If you use the `spec` or `spec_set` arguments then *only* magic
2031 methods that exist in the spec will be created.
2032
2033 Attributes and the return value of a `MagicMock` will also be `MagicMocks`.
2034 """
2035 def mock_add_spec(self, spec, spec_set=False):
2036 """Add a spec to a mock. `spec` can either be an object or a
2037 list of strings. Only attributes on the `spec` can be fetched as
2038 attributes from the mock.
2039
2040 If `spec_set` is True then only attributes on the spec can be set."""
2041 self._mock_add_spec(spec, spec_set)
2042 self._mock_set_magics()
2043
2044
2045
2046class MagicProxy(object):
2047 def __init__(self, name, parent):
2048 self.name = name
2049 self.parent = parent
2050
Michael Foord345266a2012-03-14 12:24:34 -07002051 def create_mock(self):
2052 entry = self.name
2053 parent = self.parent
2054 m = parent._get_child_mock(name=entry, _new_name=entry,
2055 _new_parent=parent)
2056 setattr(parent, entry, m)
2057 _set_return_value(parent, m, entry)
2058 return m
2059
2060 def __get__(self, obj, _type=None):
2061 return self.create_mock()
2062
2063
Lisa Roach77b3b772019-05-20 09:19:53 -07002064class AsyncMockMixin(Base):
2065 awaited = _delegating_property('awaited')
2066 await_count = _delegating_property('await_count')
2067 await_args = _delegating_property('await_args')
2068 await_args_list = _delegating_property('await_args_list')
2069
2070 def __init__(self, *args, **kwargs):
2071 super().__init__(*args, **kwargs)
2072 # asyncio.iscoroutinefunction() checks _is_coroutine property to say if an
2073 # object is a coroutine. Without this check it looks to see if it is a
2074 # function/method, which in this case it is not (since it is an
2075 # AsyncMock).
2076 # It is set through __dict__ because when spec_set is True, this
2077 # attribute is likely undefined.
2078 self.__dict__['_is_coroutine'] = asyncio.coroutines._is_coroutine
2079 self.__dict__['_mock_awaited'] = _AwaitEvent(self)
2080 self.__dict__['_mock_await_count'] = 0
2081 self.__dict__['_mock_await_args'] = None
2082 self.__dict__['_mock_await_args_list'] = _CallList()
2083 code_mock = NonCallableMock(spec_set=CodeType)
2084 code_mock.co_flags = inspect.CO_COROUTINE
2085 self.__dict__['__code__'] = code_mock
2086
2087 async def _mock_call(_mock_self, *args, **kwargs):
2088 self = _mock_self
2089 try:
2090 result = super()._mock_call(*args, **kwargs)
2091 except (BaseException, StopIteration) as e:
2092 side_effect = self.side_effect
2093 if side_effect is not None and not callable(side_effect):
2094 raise
2095 return await _raise(e)
2096
2097 _call = self.call_args
2098
2099 async def proxy():
2100 try:
2101 if inspect.isawaitable(result):
2102 return await result
2103 else:
2104 return result
2105 finally:
2106 self.await_count += 1
2107 self.await_args = _call
2108 self.await_args_list.append(_call)
2109 await self.awaited._notify()
2110
2111 return await proxy()
2112
2113 def assert_awaited(_mock_self):
2114 """
2115 Assert that the mock was awaited at least once.
2116 """
2117 self = _mock_self
2118 if self.await_count == 0:
2119 msg = f"Expected {self._mock_name or 'mock'} to have been awaited."
2120 raise AssertionError(msg)
2121
2122 def assert_awaited_once(_mock_self):
2123 """
2124 Assert that the mock was awaited exactly once.
2125 """
2126 self = _mock_self
2127 if not self.await_count == 1:
2128 msg = (f"Expected {self._mock_name or 'mock'} to have been awaited once."
2129 f" Awaited {self.await_count} times.")
2130 raise AssertionError(msg)
2131
2132 def assert_awaited_with(_mock_self, *args, **kwargs):
2133 """
2134 Assert that the last await was with the specified arguments.
2135 """
2136 self = _mock_self
2137 if self.await_args is None:
2138 expected = self._format_mock_call_signature(args, kwargs)
2139 raise AssertionError(f'Expected await: {expected}\nNot awaited')
2140
2141 def _error_message():
Xtreak0ae022c2019-05-29 12:32:26 +05302142 msg = self._format_mock_failure_message(args, kwargs, action='await')
Lisa Roach77b3b772019-05-20 09:19:53 -07002143 return msg
2144
2145 expected = self._call_matcher((args, kwargs))
2146 actual = self._call_matcher(self.await_args)
2147 if expected != actual:
2148 cause = expected if isinstance(expected, Exception) else None
2149 raise AssertionError(_error_message()) from cause
2150
2151 def assert_awaited_once_with(_mock_self, *args, **kwargs):
2152 """
2153 Assert that the mock was awaited exactly once and with the specified
2154 arguments.
2155 """
2156 self = _mock_self
2157 if not self.await_count == 1:
2158 msg = (f"Expected {self._mock_name or 'mock'} to have been awaited once."
2159 f" Awaited {self.await_count} times.")
2160 raise AssertionError(msg)
2161 return self.assert_awaited_with(*args, **kwargs)
2162
2163 def assert_any_await(_mock_self, *args, **kwargs):
2164 """
2165 Assert the mock has ever been awaited with the specified arguments.
2166 """
2167 self = _mock_self
2168 expected = self._call_matcher((args, kwargs))
2169 actual = [self._call_matcher(c) for c in self.await_args_list]
2170 if expected not in actual:
2171 cause = expected if isinstance(expected, Exception) else None
2172 expected_string = self._format_mock_call_signature(args, kwargs)
2173 raise AssertionError(
2174 '%s await not found' % expected_string
2175 ) from cause
2176
2177 def assert_has_awaits(_mock_self, calls, any_order=False):
2178 """
2179 Assert the mock has been awaited with the specified calls.
2180 The :attr:`await_args_list` list is checked for the awaits.
2181
2182 If `any_order` is False (the default) then the awaits must be
2183 sequential. There can be extra calls before or after the
2184 specified awaits.
2185
2186 If `any_order` is True then the awaits can be in any order, but
2187 they must all appear in :attr:`await_args_list`.
2188 """
2189 self = _mock_self
2190 expected = [self._call_matcher(c) for c in calls]
2191 cause = expected if isinstance(expected, Exception) else None
2192 all_awaits = _CallList(self._call_matcher(c) for c in self.await_args_list)
2193 if not any_order:
2194 if expected not in all_awaits:
2195 raise AssertionError(
Xtreak0ae022c2019-05-29 12:32:26 +05302196 f'Awaits not found.\nExpected: {_CallList(calls)}\n'
Lisa Roach77b3b772019-05-20 09:19:53 -07002197 f'Actual: {self.await_args_list}'
2198 ) from cause
2199 return
2200
2201 all_awaits = list(all_awaits)
2202
2203 not_found = []
2204 for kall in expected:
2205 try:
2206 all_awaits.remove(kall)
2207 except ValueError:
2208 not_found.append(kall)
2209 if not_found:
2210 raise AssertionError(
2211 '%r not all found in await list' % (tuple(not_found),)
2212 ) from cause
2213
2214 def assert_not_awaited(_mock_self):
2215 """
2216 Assert that the mock was never awaited.
2217 """
2218 self = _mock_self
2219 if self.await_count != 0:
Xtreakff6b2e62019-05-27 18:26:23 +05302220 msg = (f"Expected {self._mock_name or 'mock'} to not have been awaited."
Lisa Roach77b3b772019-05-20 09:19:53 -07002221 f" Awaited {self.await_count} times.")
2222 raise AssertionError(msg)
2223
2224 def reset_mock(self, *args, **kwargs):
2225 """
2226 See :func:`.Mock.reset_mock()`
2227 """
2228 super().reset_mock(*args, **kwargs)
2229 self.await_count = 0
2230 self.await_args = None
2231 self.await_args_list = _CallList()
2232
2233
2234class AsyncMock(AsyncMockMixin, AsyncMagicMixin, Mock):
2235 """
2236 Enhance :class:`Mock` with features allowing to mock
2237 an async function.
2238
2239 The :class:`AsyncMock` object will behave so the object is
2240 recognized as an async function, and the result of a call is an awaitable:
2241
2242 >>> mock = AsyncMock()
2243 >>> asyncio.iscoroutinefunction(mock)
2244 True
2245 >>> inspect.isawaitable(mock())
2246 True
2247
2248
2249 The result of ``mock()`` is an async function which will have the outcome
2250 of ``side_effect`` or ``return_value``:
2251
2252 - if ``side_effect`` is a function, the async function will return the
2253 result of that function,
2254 - if ``side_effect`` is an exception, the async function will raise the
2255 exception,
2256 - if ``side_effect`` is an iterable, the async function will return the
2257 next value of the iterable, however, if the sequence of result is
2258 exhausted, ``StopIteration`` is raised immediately,
2259 - if ``side_effect`` is not defined, the async function will return the
2260 value defined by ``return_value``, hence, by default, the async function
2261 returns a new :class:`AsyncMock` object.
2262
2263 If the outcome of ``side_effect`` or ``return_value`` is an async function,
2264 the mock async function obtained when the mock object is called will be this
2265 async function itself (and not an async function returning an async
2266 function).
2267
2268 The test author can also specify a wrapped object with ``wraps``. In this
2269 case, the :class:`Mock` object behavior is the same as with an
2270 :class:`.Mock` object: the wrapped object may have methods
2271 defined as async function functions.
2272
Xtreake7cb23b2019-05-21 14:17:17 +05302273 Based on Martin Richard's asynctest project.
Lisa Roach77b3b772019-05-20 09:19:53 -07002274 """
2275
Michael Foord345266a2012-03-14 12:24:34 -07002276
2277class _ANY(object):
2278 "A helper object that compares equal to everything."
2279
2280 def __eq__(self, other):
2281 return True
2282
2283 def __ne__(self, other):
2284 return False
2285
2286 def __repr__(self):
2287 return '<ANY>'
2288
2289ANY = _ANY()
2290
2291
2292
2293def _format_call_signature(name, args, kwargs):
2294 message = '%s(%%s)' % name
2295 formatted_args = ''
2296 args_string = ', '.join([repr(arg) for arg in args])
2297 kwargs_string = ', '.join([
Kushal Das047f14c2014-06-09 13:45:56 +05302298 '%s=%r' % (key, value) for key, value in sorted(kwargs.items())
Michael Foord345266a2012-03-14 12:24:34 -07002299 ])
2300 if args_string:
2301 formatted_args = args_string
2302 if kwargs_string:
2303 if formatted_args:
2304 formatted_args += ', '
2305 formatted_args += kwargs_string
2306
2307 return message % formatted_args
2308
2309
2310
2311class _Call(tuple):
2312 """
2313 A tuple for holding the results of a call to a mock, either in the form
2314 `(args, kwargs)` or `(name, args, kwargs)`.
2315
2316 If args or kwargs are empty then a call tuple will compare equal to
2317 a tuple without those values. This makes comparisons less verbose::
2318
2319 _Call(('name', (), {})) == ('name',)
2320 _Call(('name', (1,), {})) == ('name', (1,))
2321 _Call(((), {'a': 'b'})) == ({'a': 'b'},)
2322
2323 The `_Call` object provides a useful shortcut for comparing with call::
2324
2325 _Call(((1, 2), {'a': 3})) == call(1, 2, a=3)
2326 _Call(('foo', (1, 2), {'a': 3})) == call.foo(1, 2, a=3)
2327
2328 If the _Call has no name then it will match any name.
2329 """
Victor Stinner84b6fb02017-01-06 18:15:51 +01002330 def __new__(cls, value=(), name='', parent=None, two=False,
Michael Foord345266a2012-03-14 12:24:34 -07002331 from_kall=True):
Michael Foord345266a2012-03-14 12:24:34 -07002332 args = ()
2333 kwargs = {}
2334 _len = len(value)
2335 if _len == 3:
2336 name, args, kwargs = value
2337 elif _len == 2:
2338 first, second = value
2339 if isinstance(first, str):
2340 name = first
2341 if isinstance(second, tuple):
2342 args = second
2343 else:
2344 kwargs = second
2345 else:
2346 args, kwargs = first, second
2347 elif _len == 1:
2348 value, = value
2349 if isinstance(value, str):
2350 name = value
2351 elif isinstance(value, tuple):
2352 args = value
2353 else:
2354 kwargs = value
2355
2356 if two:
2357 return tuple.__new__(cls, (args, kwargs))
2358
2359 return tuple.__new__(cls, (name, args, kwargs))
2360
2361
2362 def __init__(self, value=(), name=None, parent=None, two=False,
2363 from_kall=True):
Andrew Dunaie63e6172018-12-04 11:08:45 +02002364 self._mock_name = name
2365 self._mock_parent = parent
2366 self._mock_from_kall = from_kall
Michael Foord345266a2012-03-14 12:24:34 -07002367
2368
2369 def __eq__(self, other):
2370 if other is ANY:
2371 return True
2372 try:
2373 len_other = len(other)
2374 except TypeError:
2375 return False
2376
2377 self_name = ''
2378 if len(self) == 2:
2379 self_args, self_kwargs = self
2380 else:
2381 self_name, self_args, self_kwargs = self
2382
Andrew Dunaie63e6172018-12-04 11:08:45 +02002383 if (getattr(self, '_mock_parent', None) and getattr(other, '_mock_parent', None)
2384 and self._mock_parent != other._mock_parent):
Chris Withers8ca0fa92018-12-03 21:31:37 +00002385 return False
2386
Michael Foord345266a2012-03-14 12:24:34 -07002387 other_name = ''
2388 if len_other == 0:
2389 other_args, other_kwargs = (), {}
2390 elif len_other == 3:
2391 other_name, other_args, other_kwargs = other
2392 elif len_other == 1:
2393 value, = other
2394 if isinstance(value, tuple):
2395 other_args = value
2396 other_kwargs = {}
2397 elif isinstance(value, str):
2398 other_name = value
2399 other_args, other_kwargs = (), {}
2400 else:
2401 other_args = ()
2402 other_kwargs = value
Berker Peksag3fc536f2015-09-09 23:35:25 +03002403 elif len_other == 2:
Michael Foord345266a2012-03-14 12:24:34 -07002404 # could be (name, args) or (name, kwargs) or (args, kwargs)
2405 first, second = other
2406 if isinstance(first, str):
2407 other_name = first
2408 if isinstance(second, tuple):
2409 other_args, other_kwargs = second, {}
2410 else:
2411 other_args, other_kwargs = (), second
2412 else:
2413 other_args, other_kwargs = first, second
Berker Peksag3fc536f2015-09-09 23:35:25 +03002414 else:
2415 return False
Michael Foord345266a2012-03-14 12:24:34 -07002416
2417 if self_name and other_name != self_name:
2418 return False
2419
2420 # this order is important for ANY to work!
2421 return (other_args, other_kwargs) == (self_args, self_kwargs)
2422
2423
Berker Peksagce913872016-03-28 00:30:02 +03002424 __ne__ = object.__ne__
2425
2426
Michael Foord345266a2012-03-14 12:24:34 -07002427 def __call__(self, *args, **kwargs):
Andrew Dunaie63e6172018-12-04 11:08:45 +02002428 if self._mock_name is None:
Michael Foord345266a2012-03-14 12:24:34 -07002429 return _Call(('', args, kwargs), name='()')
2430
Andrew Dunaie63e6172018-12-04 11:08:45 +02002431 name = self._mock_name + '()'
2432 return _Call((self._mock_name, args, kwargs), name=name, parent=self)
Michael Foord345266a2012-03-14 12:24:34 -07002433
2434
2435 def __getattr__(self, attr):
Andrew Dunaie63e6172018-12-04 11:08:45 +02002436 if self._mock_name is None:
Michael Foord345266a2012-03-14 12:24:34 -07002437 return _Call(name=attr, from_kall=False)
Andrew Dunaie63e6172018-12-04 11:08:45 +02002438 name = '%s.%s' % (self._mock_name, attr)
Michael Foord345266a2012-03-14 12:24:34 -07002439 return _Call(name=name, parent=self, from_kall=False)
2440
2441
Kushal Dasa37b9582014-09-16 18:33:37 +05302442 def count(self, *args, **kwargs):
2443 return self.__getattr__('count')(*args, **kwargs)
2444
2445 def index(self, *args, **kwargs):
2446 return self.__getattr__('index')(*args, **kwargs)
2447
Kumar Akshayb0df45e2019-03-22 13:40:40 +05302448 def _get_call_arguments(self):
2449 if len(self) == 2:
2450 args, kwargs = self
2451 else:
2452 name, args, kwargs = self
2453
2454 return args, kwargs
2455
2456 @property
2457 def args(self):
2458 return self._get_call_arguments()[0]
2459
2460 @property
2461 def kwargs(self):
2462 return self._get_call_arguments()[1]
2463
Michael Foord345266a2012-03-14 12:24:34 -07002464 def __repr__(self):
Andrew Dunaie63e6172018-12-04 11:08:45 +02002465 if not self._mock_from_kall:
2466 name = self._mock_name or 'call'
Michael Foord345266a2012-03-14 12:24:34 -07002467 if name.startswith('()'):
2468 name = 'call%s' % name
2469 return name
2470
2471 if len(self) == 2:
2472 name = 'call'
2473 args, kwargs = self
2474 else:
2475 name, args, kwargs = self
2476 if not name:
2477 name = 'call'
2478 elif not name.startswith('()'):
2479 name = 'call.%s' % name
2480 else:
2481 name = 'call%s' % name
2482 return _format_call_signature(name, args, kwargs)
2483
2484
2485 def call_list(self):
2486 """For a call object that represents multiple calls, `call_list`
2487 returns a list of all the intermediate calls as well as the
2488 final call."""
2489 vals = []
2490 thing = self
2491 while thing is not None:
Andrew Dunaie63e6172018-12-04 11:08:45 +02002492 if thing._mock_from_kall:
Michael Foord345266a2012-03-14 12:24:34 -07002493 vals.append(thing)
Andrew Dunaie63e6172018-12-04 11:08:45 +02002494 thing = thing._mock_parent
Michael Foord345266a2012-03-14 12:24:34 -07002495 return _CallList(reversed(vals))
2496
2497
2498call = _Call(from_kall=False)
2499
2500
Michael Foord345266a2012-03-14 12:24:34 -07002501def create_autospec(spec, spec_set=False, instance=False, _parent=None,
2502 _name=None, **kwargs):
2503 """Create a mock object using another object as a spec. Attributes on the
2504 mock will use the corresponding attribute on the `spec` object as their
2505 spec.
2506
2507 Functions or methods being mocked will have their arguments checked
2508 to check that they are called with the correct signature.
2509
2510 If `spec_set` is True then attempting to set attributes that don't exist
2511 on the spec object will raise an `AttributeError`.
2512
2513 If a class is used as a spec then the return value of the mock (the
2514 instance of the class) will have the same spec. You can use a class as the
2515 spec for an instance object by passing `instance=True`. The returned mock
2516 will only be callable if instances of the mock are callable.
2517
2518 `create_autospec` also takes arbitrary keyword arguments that are passed to
2519 the constructor of the created mock."""
2520 if _is_list(spec):
2521 # can't pass a list instance to the mock constructor as it will be
2522 # interpreted as a list of strings
2523 spec = type(spec)
2524
2525 is_type = isinstance(spec, type)
Xtreakff6b2e62019-05-27 18:26:23 +05302526 is_async_func = _is_async_func(spec)
Michael Foord345266a2012-03-14 12:24:34 -07002527 _kwargs = {'spec': spec}
2528 if spec_set:
2529 _kwargs = {'spec_set': spec}
2530 elif spec is None:
2531 # None we mock with a normal mock without a spec
2532 _kwargs = {}
Antoine Pitrou5c64df72013-02-03 00:23:58 +01002533 if _kwargs and instance:
2534 _kwargs['_spec_as_instance'] = True
Michael Foord345266a2012-03-14 12:24:34 -07002535
2536 _kwargs.update(kwargs)
2537
2538 Klass = MagicMock
Gregory P. Smithd4583d72016-08-15 23:23:40 -07002539 if inspect.isdatadescriptor(spec):
Michael Foord345266a2012-03-14 12:24:34 -07002540 # descriptors don't have a spec
2541 # because we don't know what type they return
2542 _kwargs = {}
Lisa Roach77b3b772019-05-20 09:19:53 -07002543 elif is_async_func:
2544 if instance:
2545 raise RuntimeError("Instance can not be True when create_autospec "
2546 "is mocking an async function")
2547 Klass = AsyncMock
Michael Foord345266a2012-03-14 12:24:34 -07002548 elif not _callable(spec):
2549 Klass = NonCallableMagicMock
2550 elif is_type and instance and not _instance_callable(spec):
2551 Klass = NonCallableMagicMock
2552
Kushal Das484f8a82014-04-16 01:05:50 +05302553 _name = _kwargs.pop('name', _name)
2554
Michael Foord345266a2012-03-14 12:24:34 -07002555 _new_name = _name
2556 if _parent is None:
2557 # for a top level object no _new_name should be set
2558 _new_name = ''
2559
2560 mock = Klass(parent=_parent, _new_parent=_parent, _new_name=_new_name,
2561 name=_name, **_kwargs)
2562
2563 if isinstance(spec, FunctionTypes):
2564 # should only happen at the top level because we don't
2565 # recurse for functions
2566 mock = _set_signature(mock, spec)
Lisa Roach77b3b772019-05-20 09:19:53 -07002567 if is_async_func:
Xtreakff6b2e62019-05-27 18:26:23 +05302568 _setup_async_mock(mock)
Michael Foord345266a2012-03-14 12:24:34 -07002569 else:
2570 _check_signature(spec, mock, is_type, instance)
2571
2572 if _parent is not None and not instance:
2573 _parent._mock_children[_name] = mock
2574
2575 if is_type and not instance and 'return_value' not in kwargs:
Michael Foord345266a2012-03-14 12:24:34 -07002576 mock.return_value = create_autospec(spec, spec_set, instance=True,
2577 _name='()', _parent=mock)
2578
2579 for entry in dir(spec):
2580 if _is_magic(entry):
2581 # MagicMock already does the useful magic methods for us
2582 continue
2583
Michael Foord345266a2012-03-14 12:24:34 -07002584 # XXXX do we need a better way of getting attributes without
2585 # triggering code execution (?) Probably not - we need the actual
2586 # object to mock it so we would rather trigger a property than mock
2587 # the property descriptor. Likewise we want to mock out dynamically
2588 # provided attributes.
Michael Foord656319e2012-04-13 17:39:16 +01002589 # XXXX what about attributes that raise exceptions other than
2590 # AttributeError on being fetched?
Michael Foord345266a2012-03-14 12:24:34 -07002591 # we could be resilient against it, or catch and propagate the
2592 # exception when the attribute is fetched from the mock
Michael Foord656319e2012-04-13 17:39:16 +01002593 try:
2594 original = getattr(spec, entry)
2595 except AttributeError:
2596 continue
Michael Foord345266a2012-03-14 12:24:34 -07002597
2598 kwargs = {'spec': original}
2599 if spec_set:
2600 kwargs = {'spec_set': original}
2601
2602 if not isinstance(original, FunctionTypes):
2603 new = _SpecState(original, spec_set, mock, entry, instance)
2604 mock._mock_children[entry] = new
2605 else:
2606 parent = mock
2607 if isinstance(spec, FunctionTypes):
2608 parent = mock.mock
2609
Michael Foord345266a2012-03-14 12:24:34 -07002610 skipfirst = _must_skip(spec, entry, is_type)
Antoine Pitrou5c64df72013-02-03 00:23:58 +01002611 kwargs['_eat_self'] = skipfirst
Lisa Roach77b3b772019-05-20 09:19:53 -07002612 if asyncio.iscoroutinefunction(original):
2613 child_klass = AsyncMock
2614 else:
2615 child_klass = MagicMock
2616 new = child_klass(parent=parent, name=entry, _new_name=entry,
2617 _new_parent=parent,
2618 **kwargs)
Antoine Pitrou5c64df72013-02-03 00:23:58 +01002619 mock._mock_children[entry] = new
Michael Foord345266a2012-03-14 12:24:34 -07002620 _check_signature(original, new, skipfirst=skipfirst)
2621
2622 # so functions created with _set_signature become instance attributes,
2623 # *plus* their underlying mock exists in _mock_children of the parent
2624 # mock. Adding to _mock_children may be unnecessary where we are also
2625 # setting as an instance attribute?
2626 if isinstance(new, FunctionTypes):
2627 setattr(mock, entry, new)
2628
2629 return mock
2630
2631
2632def _must_skip(spec, entry, is_type):
Antoine Pitrou5c64df72013-02-03 00:23:58 +01002633 """
2634 Return whether we should skip the first argument on spec's `entry`
2635 attribute.
2636 """
Michael Foord345266a2012-03-14 12:24:34 -07002637 if not isinstance(spec, type):
2638 if entry in getattr(spec, '__dict__', {}):
2639 # instance attribute - shouldn't skip
2640 return False
Michael Foord345266a2012-03-14 12:24:34 -07002641 spec = spec.__class__
Michael Foord345266a2012-03-14 12:24:34 -07002642
2643 for klass in spec.__mro__:
2644 result = klass.__dict__.get(entry, DEFAULT)
2645 if result is DEFAULT:
2646 continue
2647 if isinstance(result, (staticmethod, classmethod)):
2648 return False
Antoine Pitrou5c64df72013-02-03 00:23:58 +01002649 elif isinstance(getattr(result, '__get__', None), MethodWrapperTypes):
2650 # Normal method => skip if looked up on type
2651 # (if looked up on instance, self is already skipped)
2652 return is_type
2653 else:
2654 return False
Michael Foord345266a2012-03-14 12:24:34 -07002655
Chris Withersadbf1782019-05-01 23:04:04 +01002656 # function is a dynamically provided attribute
Michael Foord345266a2012-03-14 12:24:34 -07002657 return is_type
2658
2659
Michael Foord345266a2012-03-14 12:24:34 -07002660class _SpecState(object):
2661
2662 def __init__(self, spec, spec_set=False, parent=None,
2663 name=None, ids=None, instance=False):
2664 self.spec = spec
2665 self.ids = ids
2666 self.spec_set = spec_set
2667 self.parent = parent
2668 self.instance = instance
2669 self.name = name
2670
2671
2672FunctionTypes = (
2673 # python function
2674 type(create_autospec),
2675 # instance method
2676 type(ANY.__eq__),
Michael Foord345266a2012-03-14 12:24:34 -07002677)
2678
Antoine Pitrou5c64df72013-02-03 00:23:58 +01002679MethodWrapperTypes = (
2680 type(ANY.__eq__.__get__),
2681)
2682
Michael Foord345266a2012-03-14 12:24:34 -07002683
Michael Foorda74561a2012-03-25 19:03:13 +01002684file_spec = None
Michael Foord345266a2012-03-14 12:24:34 -07002685
Michael Foord04cbe0c2013-03-19 17:22:51 -07002686
Rémi Lapeyre11a88322019-05-07 12:48:36 +02002687def _to_stream(read_data):
2688 if isinstance(read_data, bytes):
2689 return io.BytesIO(read_data)
Michael Foord04cbe0c2013-03-19 17:22:51 -07002690 else:
Rémi Lapeyre11a88322019-05-07 12:48:36 +02002691 return io.StringIO(read_data)
Michael Foord0dccf652012-03-25 19:11:50 +01002692
Robert Collins5329aaa2015-07-17 20:08:45 +12002693
Michael Foord0dccf652012-03-25 19:11:50 +01002694def mock_open(mock=None, read_data=''):
Michael Foord99254732012-03-25 19:07:33 +01002695 """
2696 A helper function to create a mock to replace the use of `open`. It works
2697 for `open` called directly or used as a context manager.
2698
2699 The `mock` argument is the mock object to configure. If `None` (the
2700 default) then a `MagicMock` will be created for you, with the API limited
2701 to methods or attributes available on standard file handles.
2702
Xtreak71f82a22018-12-20 21:30:21 +05302703 `read_data` is a string for the `read`, `readline` and `readlines` of the
Michael Foord04cbe0c2013-03-19 17:22:51 -07002704 file handle to return. This is an empty string by default.
Michael Foord99254732012-03-25 19:07:33 +01002705 """
Rémi Lapeyre11a88322019-05-07 12:48:36 +02002706 _read_data = _to_stream(read_data)
2707 _state = [_read_data, None]
2708
Robert Collinsca647ef2015-07-24 03:48:20 +12002709 def _readlines_side_effect(*args, **kwargs):
2710 if handle.readlines.return_value is not None:
2711 return handle.readlines.return_value
Rémi Lapeyre11a88322019-05-07 12:48:36 +02002712 return _state[0].readlines(*args, **kwargs)
Robert Collinsca647ef2015-07-24 03:48:20 +12002713
2714 def _read_side_effect(*args, **kwargs):
2715 if handle.read.return_value is not None:
2716 return handle.read.return_value
Rémi Lapeyre11a88322019-05-07 12:48:36 +02002717 return _state[0].read(*args, **kwargs)
Robert Collinsca647ef2015-07-24 03:48:20 +12002718
Rémi Lapeyre11a88322019-05-07 12:48:36 +02002719 def _readline_side_effect(*args, **kwargs):
Tony Flury20870232018-09-12 23:21:16 +01002720 yield from _iter_side_effect()
2721 while True:
Rémi Lapeyre11a88322019-05-07 12:48:36 +02002722 yield _state[0].readline(*args, **kwargs)
Tony Flury20870232018-09-12 23:21:16 +01002723
2724 def _iter_side_effect():
Robert Collinsca647ef2015-07-24 03:48:20 +12002725 if handle.readline.return_value is not None:
2726 while True:
2727 yield handle.readline.return_value
2728 for line in _state[0]:
2729 yield line
Robert Collinsca647ef2015-07-24 03:48:20 +12002730
Damien Nadé394119a2019-05-23 12:03:25 +02002731 def _next_side_effect():
2732 if handle.readline.return_value is not None:
2733 return handle.readline.return_value
2734 return next(_state[0])
2735
Michael Foorda74561a2012-03-25 19:03:13 +01002736 global file_spec
2737 if file_spec is None:
2738 import _io
2739 file_spec = list(set(dir(_io.TextIOWrapper)).union(set(dir(_io.BytesIO))))
2740
Michael Foord345266a2012-03-14 12:24:34 -07002741 if mock is None:
Michael Foord0dccf652012-03-25 19:11:50 +01002742 mock = MagicMock(name='open', spec=open)
Michael Foord345266a2012-03-14 12:24:34 -07002743
Robert Collinsca647ef2015-07-24 03:48:20 +12002744 handle = MagicMock(spec=file_spec)
2745 handle.__enter__.return_value = handle
Michael Foord04cbe0c2013-03-19 17:22:51 -07002746
Robert Collinsca647ef2015-07-24 03:48:20 +12002747 handle.write.return_value = None
2748 handle.read.return_value = None
2749 handle.readline.return_value = None
2750 handle.readlines.return_value = None
Michael Foord04cbe0c2013-03-19 17:22:51 -07002751
Robert Collinsca647ef2015-07-24 03:48:20 +12002752 handle.read.side_effect = _read_side_effect
2753 _state[1] = _readline_side_effect()
2754 handle.readline.side_effect = _state[1]
2755 handle.readlines.side_effect = _readlines_side_effect
Tony Flury20870232018-09-12 23:21:16 +01002756 handle.__iter__.side_effect = _iter_side_effect
Damien Nadé394119a2019-05-23 12:03:25 +02002757 handle.__next__.side_effect = _next_side_effect
Michael Foord345266a2012-03-14 12:24:34 -07002758
Robert Collinsca647ef2015-07-24 03:48:20 +12002759 def reset_data(*args, **kwargs):
Rémi Lapeyre11a88322019-05-07 12:48:36 +02002760 _state[0] = _to_stream(read_data)
Robert Collinsca647ef2015-07-24 03:48:20 +12002761 if handle.readline.side_effect == _state[1]:
2762 # Only reset the side effect if the user hasn't overridden it.
2763 _state[1] = _readline_side_effect()
2764 handle.readline.side_effect = _state[1]
2765 return DEFAULT
Robert Collins5329aaa2015-07-17 20:08:45 +12002766
Robert Collinsca647ef2015-07-24 03:48:20 +12002767 mock.side_effect = reset_data
2768 mock.return_value = handle
Michael Foord345266a2012-03-14 12:24:34 -07002769 return mock
2770
2771
2772class PropertyMock(Mock):
Michael Foord99254732012-03-25 19:07:33 +01002773 """
2774 A mock intended to be used as a property, or other descriptor, on a class.
2775 `PropertyMock` provides `__get__` and `__set__` methods so you can specify
2776 a return value when it is fetched.
2777
2778 Fetching a `PropertyMock` instance from an object calls the mock, with
2779 no args. Setting it calls the mock with the value being set.
2780 """
Michael Foordc2870622012-04-13 16:57:22 +01002781 def _get_child_mock(self, **kwargs):
2782 return MagicMock(**kwargs)
2783
Michael Foord345266a2012-03-14 12:24:34 -07002784 def __get__(self, obj, obj_type):
2785 return self()
2786 def __set__(self, obj, val):
2787 self(val)
Mario Corchero552be9d2017-10-17 12:35:11 +01002788
2789
2790def seal(mock):
Mario Corchero96200eb2018-10-19 22:57:37 +01002791 """Disable the automatic generation of child mocks.
Mario Corchero552be9d2017-10-17 12:35:11 +01002792
2793 Given an input Mock, seals it to ensure no further mocks will be generated
2794 when accessing an attribute that was not already defined.
2795
Mario Corchero96200eb2018-10-19 22:57:37 +01002796 The operation recursively seals the mock passed in, meaning that
2797 the mock itself, any mocks generated by accessing one of its attributes,
2798 and all assigned mocks without a name or spec will be sealed.
Mario Corchero552be9d2017-10-17 12:35:11 +01002799 """
2800 mock._mock_sealed = True
2801 for attr in dir(mock):
2802 try:
2803 m = getattr(mock, attr)
2804 except AttributeError:
2805 continue
2806 if not isinstance(m, NonCallableMock):
2807 continue
2808 if m._mock_new_parent is mock:
2809 seal(m)
Lisa Roach77b3b772019-05-20 09:19:53 -07002810
2811
2812async def _raise(exception):
2813 raise exception
2814
2815
2816class _AsyncIterator:
2817 """
2818 Wraps an iterator in an asynchronous iterator.
2819 """
2820 def __init__(self, iterator):
2821 self.iterator = iterator
2822 code_mock = NonCallableMock(spec_set=CodeType)
2823 code_mock.co_flags = inspect.CO_ITERABLE_COROUTINE
2824 self.__dict__['__code__'] = code_mock
2825
2826 def __aiter__(self):
2827 return self
2828
2829 async def __anext__(self):
2830 try:
2831 return next(self.iterator)
2832 except StopIteration:
2833 pass
2834 raise StopAsyncIteration
2835
2836
2837class _AwaitEvent:
2838 def __init__(self, mock):
2839 self._mock = mock
2840 self._condition = None
2841
2842 async def _notify(self):
2843 condition = self._get_condition()
2844 try:
2845 await condition.acquire()
2846 condition.notify_all()
2847 finally:
2848 condition.release()
2849
2850 def _get_condition(self):
2851 """
2852 Creation of condition is delayed, to minimize the chance of using the
2853 wrong loop.
2854 A user may create a mock with _AwaitEvent before selecting the
2855 execution loop. Requiring a user to delay creation is error-prone and
2856 inflexible. Instead, condition is created when user actually starts to
2857 use the mock.
2858 """
2859 # No synchronization is needed:
2860 # - asyncio is thread unsafe
2861 # - there are no awaits here, method will be executed without
2862 # switching asyncio context.
2863 if self._condition is None:
2864 self._condition = asyncio.Condition()
2865
2866 return self._condition