-
Notifications
You must be signed in to change notification settings - Fork 31
/
models.py
35 lines (26 loc) · 808 Bytes
/
models.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
"""
mock_django.models
~~~~~~~~~~~~~~~~~~
:copyright: (c) 2012 DISQUS.
:license: Apache License 2.0, see LICENSE for more details.
"""
import mock
__all__ = ('ModelMock',)
# TODO: make foreignkey_id == foreignkey.id
class _ModelMock(mock.MagicMock):
def __init__(self, *args, **kwargs):
super(_ModelMock, self).__init__(*args, **kwargs)
# Django ORM needed state for write status
self._state = mock.Mock()
self._state.db = None
def _get_child_mock(self, **kwargs):
name = kwargs.get('name', '')
if name == 'pk':
return self.id
return super(_ModelMock, self)._get_child_mock(**kwargs)
def ModelMock(model):
"""
>>> Post = ModelMock(Post)
>>> assert post.pk == post.id
"""
return _ModelMock(spec=model())