213 lines
6.5 KiB
Python
213 lines
6.5 KiB
Python
# Copyright 2012-present MongoDB, Inc.
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
|
|
"""Utilities for testing pymongo that require synchronization."""
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import contextlib
|
|
import random
|
|
import threading # Used in the synchronized version of this file
|
|
import time
|
|
from asyncio import iscoroutinefunction
|
|
|
|
from bson.son import SON
|
|
from pymongo import AsyncMongoClient
|
|
from pymongo.errors import ConfigurationError
|
|
from pymongo.hello import HelloCompat
|
|
from pymongo.lock import _async_create_lock
|
|
from pymongo.operations import _Op
|
|
from pymongo.read_preferences import ReadPreference
|
|
from pymongo.server_selectors import any_server_selector, writable_server_selector
|
|
from pymongo.synchronous.pool import _CancellationContext, _PoolGeneration
|
|
|
|
_IS_SYNC = False
|
|
|
|
|
|
async def async_get_pool(client):
|
|
"""Get the standalone, primary, or mongos pool."""
|
|
topology = await client._get_topology()
|
|
server = await topology._select_server(writable_server_selector, _Op.TEST)
|
|
return server.pool
|
|
|
|
|
|
async def async_get_pools(client):
|
|
"""Get all pools."""
|
|
return [
|
|
server.pool
|
|
for server in await (await client._get_topology()).select_servers(
|
|
any_server_selector, _Op.TEST
|
|
)
|
|
]
|
|
|
|
|
|
async def async_wait_until(predicate, success_description, timeout=10):
|
|
"""Wait up to 10 seconds (by default) for predicate to be true.
|
|
|
|
E.g.:
|
|
|
|
wait_until(lambda: client.primary == ('a', 1),
|
|
'connect to the primary')
|
|
|
|
If the lambda-expression isn't true after 10 seconds, we raise
|
|
AssertionError("Didn't ever connect to the primary").
|
|
|
|
Returns the predicate's first true value.
|
|
"""
|
|
start = time.time()
|
|
interval = min(float(timeout) / 100, 0.1)
|
|
while True:
|
|
if iscoroutinefunction(predicate):
|
|
retval = await predicate()
|
|
else:
|
|
retval = predicate()
|
|
if retval:
|
|
return retval
|
|
|
|
if time.time() - start > timeout:
|
|
raise AssertionError("Didn't ever %s" % success_description)
|
|
|
|
await asyncio.sleep(interval)
|
|
|
|
|
|
async def async_is_mongos(client):
|
|
res = await client.admin.command(HelloCompat.LEGACY_CMD)
|
|
return res.get("msg", "") == "isdbgrid"
|
|
|
|
|
|
async def async_ensure_all_connected(client: AsyncMongoClient) -> None:
|
|
"""Ensure that the client's connection pool has socket connections to all
|
|
members of a replica set. Raises ConfigurationError when called with a
|
|
non-replica set client.
|
|
|
|
Depending on the use-case, the caller may need to clear any event listeners
|
|
that are configured on the client.
|
|
"""
|
|
hello: dict = await client.admin.command(HelloCompat.LEGACY_CMD)
|
|
if "setName" not in hello:
|
|
raise ConfigurationError("cluster is not a replica set")
|
|
|
|
target_host_list = set(hello["hosts"] + hello.get("passives", []))
|
|
connected_host_list = {hello["me"]}
|
|
|
|
# Run hello until we have connected to each host at least once.
|
|
async def discover():
|
|
i = 0
|
|
while i < 100 and connected_host_list != target_host_list:
|
|
hello: dict = await client.admin.command(
|
|
HelloCompat.LEGACY_CMD, read_preference=ReadPreference.SECONDARY
|
|
)
|
|
connected_host_list.update([hello["me"]])
|
|
i += 1
|
|
return connected_host_list
|
|
|
|
try:
|
|
|
|
async def predicate():
|
|
return target_host_list == await discover()
|
|
|
|
await async_wait_until(predicate, "connected to all hosts")
|
|
except AssertionError as exc:
|
|
raise AssertionError(
|
|
f"{exc}, {connected_host_list} != {target_host_list}, {client.topology_description}"
|
|
)
|
|
|
|
|
|
async def asyncAssertRaisesExactly(cls, fn, *args, **kwargs):
|
|
"""
|
|
Unlike the standard assertRaises, this checks that a function raises a
|
|
specific class of exception, and not a subclass. E.g., check that
|
|
MongoClient() raises ConnectionFailure but not its subclass, AutoReconnect.
|
|
"""
|
|
try:
|
|
await fn(*args, **kwargs)
|
|
except Exception as e:
|
|
assert e.__class__ == cls, f"got {e.__class__.__name__}, expected {cls.__name__}"
|
|
else:
|
|
raise AssertionError("%s not raised" % cls)
|
|
|
|
|
|
async def async_set_fail_point(client, command_args):
|
|
cmd = SON([("configureFailPoint", "failCommand")])
|
|
cmd.update(command_args)
|
|
await client.admin.command(cmd)
|
|
|
|
|
|
async def async_joinall(tasks):
|
|
"""Join threads with a 5-minute timeout, assert joins succeeded"""
|
|
if _IS_SYNC:
|
|
for t in tasks:
|
|
t.join(300)
|
|
assert not t.is_alive(), "Thread %s hung" % t
|
|
else:
|
|
await asyncio.wait([t.task for t in tasks if t is not None], timeout=300)
|
|
|
|
|
|
class AsyncMockConnection:
|
|
def __init__(self):
|
|
self.cancel_context = _CancellationContext()
|
|
self.more_to_come = False
|
|
self.id = random.randint(0, 100)
|
|
self.server_connection_id = random.randint(0, 100)
|
|
|
|
def close_conn(self, reason):
|
|
pass
|
|
|
|
def __aenter__(self):
|
|
return self
|
|
|
|
def __aexit__(self, exc_type, exc_val, exc_tb):
|
|
pass
|
|
|
|
|
|
class AsyncMockPool:
|
|
def __init__(self, address, options, handshake=True, client_id=None):
|
|
self.gen = _PoolGeneration()
|
|
self._lock = _async_create_lock()
|
|
self.opts = options
|
|
self.operation_count = 0
|
|
self.conns = []
|
|
|
|
def stale_generation(self, gen, service_id):
|
|
return self.gen.stale(gen, service_id)
|
|
|
|
@contextlib.asynccontextmanager
|
|
async def checkout(self, handler=None):
|
|
yield AsyncMockConnection()
|
|
|
|
async def checkin(self, *args, **kwargs):
|
|
pass
|
|
|
|
async def _reset(self, service_id=None):
|
|
async with self._lock:
|
|
self.gen.inc(service_id)
|
|
|
|
async def ready(self):
|
|
pass
|
|
|
|
async def reset(self, service_id=None, interrupt_connections=False):
|
|
await self._reset()
|
|
|
|
async def reset_without_pause(self):
|
|
await self._reset()
|
|
|
|
async def close(self):
|
|
await self._reset()
|
|
|
|
async def update_is_writable(self, is_writable):
|
|
pass
|
|
|
|
async def remove_stale_sockets(self, *args, **kwargs):
|
|
pass
|