343 lines
13 KiB
Python
343 lines
13 KiB
Python
# Copyright 2014 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 AsyncIOMotorCollection."""
|
|
|
|
import asyncio
|
|
import sys
|
|
import traceback
|
|
import unittest
|
|
from test.asyncio_tests import AsyncIOTestCase, asyncio_test
|
|
from test.test_environment import env
|
|
from test.utils import ignore_deprecations
|
|
|
|
import pymongo
|
|
from bson import CodecOptions
|
|
from bson.binary import JAVA_LEGACY
|
|
from pymongo import ReadPreference, WriteConcern
|
|
from pymongo.encryption import Algorithm, QueryType
|
|
from pymongo.errors import BulkWriteError, DuplicateKeyError, OperationFailure
|
|
from pymongo.read_concern import ReadConcern
|
|
from pymongo.read_preferences import Secondary
|
|
|
|
from motor.motor_asyncio import AsyncIOMotorClientEncryption, AsyncIOMotorCollection
|
|
|
|
if pymongo.version_tuple >= (4, 4, 0):
|
|
from pymongo.encryption_options import RangeOpts
|
|
|
|
|
|
class TestAsyncIOCollection(AsyncIOTestCase):
|
|
@asyncio_test
|
|
async def test_collection(self):
|
|
# Test that we can create a collection directly, not just from
|
|
# database accessors.
|
|
collection = AsyncIOMotorCollection(self.db, "test_collection")
|
|
|
|
# Make sure we got the right collection and it can do an operation
|
|
self.assertEqual("test_collection", collection.name)
|
|
await collection.delete_many({})
|
|
await collection.insert_one({"_id": 1})
|
|
doc = await collection.find_one({"_id": 1})
|
|
self.assertEqual(1, doc["_id"])
|
|
|
|
# If you pass kwargs to PyMongo's Collection(), it calls
|
|
# db.create_collection(). Motor can't do I/O in a constructor
|
|
# so this is prohibited.
|
|
self.assertRaises(
|
|
TypeError, AsyncIOMotorCollection, self.db, "test_collection", capped=True
|
|
)
|
|
|
|
@asyncio_test
|
|
async def test_dotted_collection_name(self):
|
|
# Ensure that remove, insert, and find work on collections with dots
|
|
# in their names.
|
|
for coll in (self.db.foo.bar, self.db.foo.bar.baz):
|
|
await coll.delete_many({})
|
|
result = await coll.insert_one({"_id": "xyzzy"})
|
|
self.assertEqual("xyzzy", result.inserted_id)
|
|
result = await coll.find_one({"_id": "xyzzy"})
|
|
self.assertEqual(result["_id"], "xyzzy")
|
|
await coll.delete_many({})
|
|
resp = await coll.find_one({"_id": "xyzzy"})
|
|
self.assertEqual(None, resp)
|
|
|
|
def test_call(self):
|
|
# Prevents user error with nice message.
|
|
try:
|
|
self.db.foo()
|
|
except TypeError as e:
|
|
self.assertTrue("no such method exists" in str(e))
|
|
else:
|
|
self.fail("Expected TypeError")
|
|
|
|
@ignore_deprecations
|
|
@asyncio_test
|
|
async def test_update(self):
|
|
await self.collection.insert_one({"_id": 1})
|
|
result = await self.collection.update_one({"_id": 1}, {"$set": {"foo": "bar"}})
|
|
|
|
self.assertIsNone(result.upserted_id)
|
|
self.assertEqual(1, result.modified_count)
|
|
|
|
@ignore_deprecations
|
|
@asyncio_test
|
|
async def test_update_bad(self):
|
|
# Violate a unique index, make sure we handle error well
|
|
coll = self.db.unique_collection
|
|
await coll.create_index("s", unique=True)
|
|
|
|
try:
|
|
await coll.insert_many([{"s": 1}, {"s": 2}])
|
|
with self.assertRaises(DuplicateKeyError):
|
|
await coll.update_one({"s": 2}, {"$set": {"s": 1}})
|
|
|
|
finally:
|
|
await coll.drop()
|
|
|
|
@asyncio_test
|
|
async def test_insert_one(self):
|
|
collection = self.collection
|
|
result = await collection.insert_one({"_id": 201})
|
|
self.assertEqual(201, result.inserted_id)
|
|
|
|
@ignore_deprecations
|
|
@asyncio_test
|
|
async def test_insert_many_one_bad(self):
|
|
collection = self.collection
|
|
await collection.insert_one({"_id": 2})
|
|
|
|
# Violate a unique index in one of many updates, handle error.
|
|
with self.assertRaises(BulkWriteError):
|
|
await collection.insert_many([{"_id": 1}, {"_id": 2}, {"_id": 3}]) # Already exists
|
|
|
|
# First insert should have succeeded, but not second or third.
|
|
self.assertEqual(set([1, 2]), set(await collection.distinct("_id")))
|
|
|
|
@asyncio_test
|
|
async def test_delete_one(self):
|
|
# Remove a document twice, check that we get a success responses
|
|
# and n = 0 for the second time.
|
|
await self.collection.insert_one({"_id": 1})
|
|
result = await self.collection.delete_one({"_id": 1})
|
|
|
|
# First time we remove, n = 1
|
|
self.assertEqual(1, result.raw_result["n"])
|
|
self.assertEqual(1, result.raw_result["ok"])
|
|
self.assertEqual(None, result.raw_result.get("err"))
|
|
|
|
result = await self.collection.delete_one({"_id": 1})
|
|
|
|
# Second time, document is already gone, n = 0
|
|
self.assertEqual(0, result.raw_result["n"])
|
|
self.assertEqual(1, result.raw_result["ok"])
|
|
self.assertEqual(None, result.raw_result.get("err"))
|
|
|
|
@ignore_deprecations
|
|
@asyncio_test
|
|
async def test_unacknowledged_insert(self):
|
|
coll = self.db.test_unacknowledged_insert
|
|
await coll.with_options(write_concern=WriteConcern(0)).insert_one({"_id": 1})
|
|
|
|
# The insert is eventually executed.
|
|
while not (await coll.count_documents({})):
|
|
await asyncio.sleep(0.1)
|
|
|
|
@ignore_deprecations
|
|
@asyncio_test
|
|
async def test_unacknowledged_update(self):
|
|
coll = self.collection
|
|
|
|
await coll.insert_one({"_id": 1})
|
|
await coll.with_options(write_concern=WriteConcern(0)).update_one(
|
|
{"_id": 1}, {"$set": {"a": 1}}
|
|
)
|
|
|
|
while not (await coll.find_one({"a": 1})):
|
|
await asyncio.sleep(0.1)
|
|
|
|
@ignore_deprecations
|
|
@asyncio_test
|
|
async def test_indexes(self):
|
|
test_collection = self.collection
|
|
|
|
# Create an index
|
|
idx_name = await test_collection.create_index([("foo", 1)])
|
|
index_info = await test_collection.index_information()
|
|
self.assertEqual([("foo", 1)], index_info[idx_name]["key"])
|
|
|
|
# Don't test drop_index or drop_indexes -- Synchro tests them
|
|
|
|
async def _make_test_data(self, n):
|
|
await self.db.drop_collection("test")
|
|
await self.db.test.insert_many([{"_id": i} for i in range(n)])
|
|
expected_sum = sum(range(n))
|
|
return expected_sum
|
|
|
|
pipeline = [{"$project": {"_id": "$_id"}}]
|
|
|
|
@asyncio_test(timeout=30)
|
|
async def test_aggregation_cursor(self):
|
|
db = self.db
|
|
|
|
# A small collection which returns only an initial batch,
|
|
# and a larger one that requires a getMore.
|
|
for collection_size in (10, 1000):
|
|
expected_sum = await self._make_test_data(collection_size)
|
|
cursor = db.test.aggregate(self.pipeline)
|
|
docs = await cursor.to_list(collection_size)
|
|
self.assertEqual(expected_sum, sum(doc["_id"] for doc in docs))
|
|
|
|
@asyncio_test
|
|
async def test_aggregation_cursor_exc_info(self):
|
|
await self._make_test_data(200)
|
|
cursor = self.db.test.aggregate(self.pipeline)
|
|
await cursor.to_list(length=10)
|
|
await self.db.test.drop()
|
|
try:
|
|
await cursor.to_list(length=None)
|
|
except OperationFailure:
|
|
_, _, tb = sys.exc_info()
|
|
|
|
# The call tree should include PyMongo code we ran on a thread.
|
|
formatted = "\n".join(traceback.format_tb(tb))
|
|
self.assertTrue(
|
|
"_unpack_response" in formatted or "_check_command_response" in formatted
|
|
)
|
|
|
|
@asyncio_test
|
|
async def test_aggregate_cursor_del(self):
|
|
cursor = self.db.test.aggregate(self.pipeline)
|
|
del cursor
|
|
cursor = self.db.test.aggregate(self.pipeline)
|
|
await cursor.close()
|
|
del cursor
|
|
|
|
def test_with_options(self):
|
|
coll = self.db.test
|
|
codec_options = CodecOptions(tz_aware=True, uuid_representation=JAVA_LEGACY)
|
|
|
|
write_concern = WriteConcern(w=2, j=True)
|
|
coll2 = coll.with_options(codec_options, ReadPreference.SECONDARY, write_concern)
|
|
|
|
self.assertTrue(isinstance(coll2, AsyncIOMotorCollection))
|
|
self.assertEqual(codec_options, coll2.codec_options)
|
|
self.assertEqual(Secondary(), coll2.read_preference)
|
|
self.assertEqual(write_concern, coll2.write_concern)
|
|
|
|
pref = Secondary([{"dc": "sf"}])
|
|
coll2 = coll.with_options(read_preference=pref)
|
|
self.assertEqual(pref, coll2.read_preference)
|
|
self.assertEqual(coll.codec_options, coll2.codec_options)
|
|
self.assertEqual(coll.write_concern, coll2.write_concern)
|
|
|
|
def test_sub_collection(self):
|
|
# Verify that a collection with a dotted name inherits options from its
|
|
# parent collection.
|
|
write_concern = WriteConcern(w=2, j=True)
|
|
read_concern = ReadConcern("majority")
|
|
read_preference = Secondary([{"dc": "sf"}])
|
|
codec_options = CodecOptions(tz_aware=True, uuid_representation=JAVA_LEGACY)
|
|
|
|
coll1 = self.db.get_collection(
|
|
"test",
|
|
write_concern=write_concern,
|
|
read_concern=read_concern,
|
|
read_preference=read_preference,
|
|
codec_options=codec_options,
|
|
)
|
|
|
|
coll2 = coll1.subcollection
|
|
coll3 = coll1["subcollection"]
|
|
|
|
for c in [coll1, coll2, coll3]:
|
|
self.assertEqual(write_concern, c.write_concern)
|
|
self.assertEqual(read_concern, c.read_concern)
|
|
self.assertEqual(read_preference, c.read_preference)
|
|
self.assertEqual(codec_options, c.codec_options)
|
|
|
|
@env.require_version_min(7, 0, -1, -1)
|
|
@env.require_no_standalone
|
|
@asyncio_test
|
|
async def test_async_create_encrypted_collection(self):
|
|
c = self.cx
|
|
KMS_PROVIDERS = {"local": {"key": b"\x00" * 96}}
|
|
self.cx.drop_database("db")
|
|
async with AsyncIOMotorClientEncryption(
|
|
KMS_PROVIDERS, "keyvault.datakeys", c, CodecOptions()
|
|
) as client_encryption:
|
|
coll, ef = await client_encryption.create_encrypted_collection(
|
|
database=self.db,
|
|
name="testing1",
|
|
encrypted_fields={"fields": [{"path": "ssn", "bsonType": "string", "keyId": None}]},
|
|
kms_provider="local",
|
|
)
|
|
with self.assertRaises(pymongo.errors.WriteError) as exc:
|
|
await coll.insert_one({"ssn": "123-45-6789"})
|
|
self.assertEqual(exc.exception.code, 121)
|
|
await self.db.drop_collection("testing1", encrypted_fields=ef)
|
|
|
|
@env.require_version_min(8, 0, -1, -1)
|
|
@asyncio_test
|
|
async def test_async_encrypt_expression(self):
|
|
c = self.cx
|
|
KMS_PROVIDERS = {"local": {"key": b"\x00" * 96}}
|
|
self.cx.drop_database("db")
|
|
async with AsyncIOMotorClientEncryption(
|
|
KMS_PROVIDERS, "keyvault.datakeys", c, CodecOptions()
|
|
) as client_encryption:
|
|
data_key = await client_encryption.create_data_key(
|
|
"local", key_alt_names=["pymongo_encryption_example_1"]
|
|
)
|
|
name = "DoubleNoPrecision"
|
|
range_opts = RangeOpts(sparsity=1, trim_factor=1)
|
|
for i in [6.0, 30.0, 200.0]:
|
|
insert_payload = await client_encryption.encrypt(
|
|
float(i),
|
|
key_id=data_key,
|
|
algorithm=Algorithm.RANGE,
|
|
contention_factor=0,
|
|
range_opts=range_opts,
|
|
)
|
|
self.collection.insert_one(
|
|
{
|
|
f"encrypted{name}": insert_payload,
|
|
}
|
|
)
|
|
self.assertEqual(await client_encryption.decrypt(insert_payload), i)
|
|
|
|
find_payload = await client_encryption.encrypt_expression(
|
|
expression={
|
|
"$and": [
|
|
{f"encrypted{name}": {"$gte": 6.0}},
|
|
{f"encrypted{name}": {"$lte": 200.0}},
|
|
]
|
|
},
|
|
key_id=data_key,
|
|
algorithm=Algorithm.RANGE,
|
|
query_type=QueryType.RANGE,
|
|
contention_factor=0,
|
|
range_opts=range_opts,
|
|
)
|
|
|
|
sorted_find = sorted(
|
|
await self.collection.explicit_encryption.find(find_payload).to_list(3),
|
|
key=lambda x: x["_id"],
|
|
)
|
|
for elem, expected in zip(sorted_find, [6.0, 30.0, 200.0], strict=False):
|
|
self.assertEqual(elem[f"encrypted{name}"], expected)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|