15
15
"""Tools for mocking parts of PyMongo to test other parts."""
18
from functools import partial
19
21
from pymongo import common
20
from pymongo import MongoClient, MongoReplicaSetClient
21
from pymongo.pool import Pool
22
from pymongo import MongoClient
23
from pymongo.errors import AutoReconnect, NetworkTimeout
24
from pymongo.ismaster import IsMaster
25
from pymongo.monitor import Monitor
26
from pymongo.pool import Pool, PoolOptions
27
from pymongo.server_description import ServerDescription
23
29
from test import host as default_host, port as default_port
24
from test.utils import my_partial
27
32
class MockPool(Pool):
28
33
def __init__(self, client, pair, *args, **kwargs):
29
# MockPool gets a 'client' arg, regular pools don't.
34
# MockPool gets a 'client' arg, regular pools don't. Weakref it to
35
# avoid cycle with __del__, causing ResourceWarnings in Python 3.3.
36
self.client = weakref.proxy(client)
31
37
self.mock_host, self.mock_port = pair
33
39
# Actually connect to the default server.
36
pair=(default_host, default_port),
41
(default_host, default_port),
42
PoolOptions(connect_timeout=20))
43
def get_socket(self, force=False):
44
@contextlib.contextmanager
45
def get_socket(self, all_credentials, checkout=False):
44
46
client = self.client
45
47
host_and_port = '%s:%s' % (self.mock_host, self.mock_port)
46
48
if host_and_port in client.mock_down_hosts:
47
raise socket.error('mock error')
49
raise AutoReconnect('mock error')
49
51
assert host_and_port in (
50
52
client.mock_standalones
51
53
+ client.mock_members
52
54
+ client.mock_mongoses), "bad host: %s" % host_and_port
54
sock_info = Pool.get_socket(self, force)
55
sock_info.mock_host = self.mock_host
56
sock_info.mock_port = self.mock_port
60
class MockClientBase(object):
61
def __init__(self, standalones, members, mongoses, config):
62
"""standalones, etc., are like ['a:1', 'b:2']"""
56
with Pool.get_socket(self, all_credentials) as sock_info:
57
sock_info.mock_host = self.mock_host
58
sock_info.mock_port = self.mock_port
62
class MockMonitor(Monitor):
70
# MockMonitor gets a 'client' arg, regular monitors don't.
79
def _check_once(self):
80
address = self._server_description.address
81
response, rtt = self.client.mock_is_master('%s:%d' % address)
82
return ServerDescription(address, IsMaster(response), rtt)
85
class MockClient(MongoClient):
87
self, standalones, members, mongoses, ismaster_hosts=None,
89
"""A MongoClient connected to the default server, with a mock topology.
91
standalones, members, mongoses determine the configuration of the
92
topology. They are formatted like ['a:1', 'b:2']. ismaster_hosts
93
provides an alternative host list for the server's mocked ismaster
94
response; see test_connect_with_internal_ips.
63
96
self.mock_standalones = standalones[:]
64
97
self.mock_members = members[:]
106
148
max_write_batch_size = self.mock_max_write_batch_sizes.get(
107
149
host, common.MAX_WRITE_BATCH_SIZE)
151
rtt = self.mock_rtts.get(host, 0)
109
153
# host is like 'a:1'.
110
154
if host in self.mock_down_hosts:
111
raise socket.timeout('mock timeout')
155
raise NetworkTimeout('mock timeout')
113
if host in self.mock_standalones:
157
elif host in self.mock_standalones:
115
160
'ismaster': True,
116
161
'minWireVersion': min_wire_version,
117
162
'maxWireVersion': max_wire_version,
118
163
'maxWriteBatchSize': max_write_batch_size}
120
if host in self.mock_members:
164
elif host in self.mock_members:
121
165
ismaster = (host == self.mock_primary)
123
167
# Simulate a replica set member.
125
170
'ismaster': ismaster,
126
171
'secondary': not ismaster,
133
178
if self.mock_primary:
134
179
response['primary'] = self.mock_primary
138
if host in self.mock_mongoses:
180
elif host in self.mock_mongoses:
140
183
'ismaster': True,
141
184
'minWireVersion': min_wire_version,
142
185
'maxWireVersion': max_wire_version,
143
186
'msg': 'isdbgrid',
144
187
'maxWriteBatchSize': max_write_batch_size}
146
# In test_internal_ips(), we try to connect to a host listed
147
# in ismaster['hosts'] but not publicly accessible.
148
raise socket.error('Unknown host: %s' % host)
150
def simple_command(self, sock_info, dbname, spec):
151
# __simple_command is also used for authentication, but in this
152
# test it's only used for ismaster.
153
assert spec == {'ismaster': 1}
154
response = self.mock_is_master(
155
'%s:%s' % (sock_info.mock_host, sock_info.mock_port))
158
return response, ping_time
161
class MockClient(MockClientBase, MongoClient):
163
self, standalones, members, mongoses, ismaster_hosts=None,
166
MockClientBase.__init__(
167
self, standalones, members, mongoses, ismaster_hosts)
169
kwargs['_pool_class'] = my_partial(MockPool, self)
170
MongoClient.__init__(self, *args, **kwargs)
172
def _MongoClient__simple_command(self, sock_info, dbname, spec):
173
return self.simple_command(sock_info, dbname, spec)
176
class MockReplicaSetClient(MockClientBase, MongoReplicaSetClient):
178
self, standalones, members, mongoses, ismaster_hosts=None,
181
MockClientBase.__init__(
182
self, standalones, members, mongoses, ismaster_hosts)
184
kwargs['_pool_class'] = my_partial(MockPool, self)
185
MongoReplicaSetClient.__init__(self, *args, **kwargs)
187
def _MongoReplicaSetClient__is_master(self, host):
188
response = self.mock_is_master('%s:%s' % host)
189
connection_pool = MockPool(self, host)
191
return response, connection_pool, ping_time
193
def _MongoReplicaSetClient__simple_command(self, sock_info, dbname, spec):
194
return self.simple_command(sock_info, dbname, spec)
189
# In test_internal_ips(), we try to connect to a host listed
190
# in ismaster['hosts'] but not publicly accessible.
191
raise AutoReconnect('Unknown host: %s' % host)
195
def _process_kill_cursors_queue(self):
196
# Avoid the background thread causing races, e.g. a surprising
197
# reconnect while we're trying to test a disconnected client.
added
removed
test/pymongo_mocks.py
