mongo-python-driver/test/mockupdb/test_query_read_pref_sharded.py
2024-10-09 09:58:20 -07:00

89 lines
2.6 KiB
Python

# Copyright 2015 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.
"""Test PyMongo query and read preference with a sharded cluster."""
from __future__ import annotations
import unittest
from test import PyMongoTestCase
import pytest
try:
from mockupdb import MockupDB, OpMsg, going
_HAVE_MOCKUPDB = True
except ImportError:
_HAVE_MOCKUPDB = False
from bson import SON
from pymongo.common import MIN_SUPPORTED_WIRE_VERSION
from pymongo.read_preferences import (
Nearest,
Primary,
PrimaryPreferred,
Secondary,
SecondaryPreferred,
)
pytestmark = pytest.mark.mockupdb
class TestQueryAndReadModeSharded(PyMongoTestCase):
def test_query_and_read_mode_sharded_op_msg(self):
"""Test OP_MSG sends non-primary $readPreference and never $query."""
server = MockupDB()
server.autoresponds(
"ismaster",
ismaster=True,
msg="isdbgrid",
minWireVersion=2,
maxWireVersion=MIN_SUPPORTED_WIRE_VERSION,
)
server.run()
self.addCleanup(server.stop)
client = self.simple_client(server.uri)
read_prefs = (
Primary(),
SecondaryPreferred(),
PrimaryPreferred(),
Secondary(),
Nearest(),
SecondaryPreferred([{"tag": "value"}]),
)
for query in (
{"a": 1},
{"$query": {"a": 1}},
):
for pref in read_prefs:
collection = client.db.get_collection("test", read_preference=pref)
cursor = collection.find(query.copy())
with going(next, cursor):
request = server.receives()
# Command is not nested in $query.
expected_cmd = SON([("find", "test"), ("filter", {"a": 1})])
if pref.mode:
expected_cmd["$readPreference"] = pref.document
request.assert_matches(OpMsg(expected_cmd))
request.replies({"cursor": {"id": 0, "firstBatch": [{}]}})
if __name__ == "__main__":
unittest.main()