async def test_find_many_skip(preset_documents):
q = Sample.find_many(Sample.integer > 1, skip=2)
assert q.skip_number == 2
q = Sample.find_many(Sample.integer > 1).skip(2)
assert q.skip_number == 2
result = (
await Sample.find_many(Sample.increment > 2)
.find_many(Sample.nested.optional == None)
.skip(1)
.to_list()
)
assert len(result) == 3
for sample in result:
assert sample.increment > 2
assert sample.nested.optional is None
len_result = 0
async for sample in Sample.find_many(Sample.increment > 2).find_many(
Sample.nested.optional == None
).skip(
1
): # noqa
assert sample in result
len_result += 1
assert len_result == len(result)
async def test_delete_many_with_session(preset_documents, session):
count_before = await Sample.count()
count_find = (
await Sample.find_many(Sample.integer > 1)
.find_many(Sample.nested.optional == None)
.count()
) # noqa
q = (
Sample.find_many(Sample.integer > 1)
.find_many(Sample.nested.optional == None)
.delete(session=session)
) # noqa
assert q.session == session
q = (
Sample.find_many(Sample.integer > 1)
.find_many(Sample.nested.optional == None)
.delete()
.set_session(session=session)
)
assert q.session == session
await q
count_after = await Sample.count()
assert count_before - count_find == count_after
async def test_find_many_limit(preset_documents):
q = Sample.find_many(Sample.integer > 1, limit=2)
assert q.limit_number == 2
q = Sample.find_many(Sample.integer > 1).limit(2)
assert q.limit_number == 2
result = (await
Sample.find_many(Sample.increment > 2
).find_many(Sample.nested.optional == None
).sort(Sample.increment
).limit(2).to_list()) # noqa
assert len(result) == 2
for a in result:
assert a.increment > 2
assert a.nested.optional is None
len_result = 0
async for a in Sample.find_many(Sample.increment > 2).find(
Sample.nested.optional == None).sort(
Sample.increment).limit(2): # noqa
assert a in result
len_result += 1
assert len_result == len(result)
async def test_find_many_with_session(preset_documents, session):
q_1 = (
Sample.find_many(Sample.integer > 1)
.find_many(Sample.nested.optional == None)
.set_session(session)
)
assert q_1.session == session
q_2 = Sample.find_many(Sample.integer > 1).find_many(
Sample.nested.optional == None, session=session
)
assert q_2.session == session
result = await q_2.to_list()
assert len(result) == 2
for a in result:
assert a.integer > 1
assert a.nested.optional is None
len_result = 0
async for a in Sample.find_many(Sample.integer > 1).find_many(
Sample.nested.optional == None
): # noqa
assert a in result
len_result += 1
assert len_result == len(result)
def test_nesting():
assert Sample.id == "_id"
q = Sample.find_many(Sample.integer == 1)
assert q.get_filter_query() == {"integer": 1}
assert Sample.integer == "integer"
q = Sample.find_many(Sample.nested.integer == 1)
assert q.get_filter_query() == {"nested.integer": 1}
assert Sample.nested.integer == "nested.integer"
q = Sample.find_many(Sample.union.s == "test")
assert q.get_filter_query() == {"union.s": "test"}
assert Sample.union.s == "union.s"
q = Sample.find_many(Sample.nested.optional == None) # noqa
assert q.get_filter_query() == {"nested.optional": None}
assert Sample.nested.optional == "nested.optional"
q = Sample.find_many(Sample.nested.integer == 1).find_many(
Sample.nested.union.s == "test")
assert q.get_filter_query() == {
"$and": [{
"nested.integer": 1
}, {
"nested.union.s": "test"
}]
}
async def test_aggregate_with_session(preset_documents, session):
q = Sample.find(Sample.increment >= 4).aggregate(
[{
"$group": {
"_id": "$string",
"total": {
"$sum": "$integer"
}
}
}],
session=session,
)
assert q.session == session
q = Sample.find(Sample.increment >= 4, session=session).aggregate([{
"$group": {
"_id": "$string",
"total": {
"$sum": "$integer"
}
}
}])
assert q.session == session
result = await q.to_list()
assert len(result) == 3
assert {"_id": "test_1", "total": 2} in result
assert {"_id": "test_2", "total": 6} in result
assert {"_id": "test_3", "total": 3} in result
async def test_find_query():
q = Sample.find_many(Sample.integer == 1).get_filter_query()
assert q == {"integer": 1}
q = Sample.find_many(Sample.integer == 1,
Sample.nested.integer >= 2).get_filter_query()
assert q == {"$and": [{"integer": 1}, {"nested.integer": {"$gte": 2}}]}
q = (Sample.find_many(Sample.integer == 1).find_many(
Sample.nested.integer >= 2).get_filter_query())
assert q == {"$and": [{"integer": 1}, {"nested.integer": {"$gte": 2}}]}
q = Sample.find().get_filter_query()
assert q == {}
async def test_aggregate_with_filter(preset_documents):
q = Sample.find(Sample.increment >= 4).aggregate([{
"$group": {
"_id": "$string",
"total": {
"$sum": "$integer"
}
}
}])
assert q.get_aggregation_pipeline() == [
{
"$match": {
"increment": {
"$gte": 4
}
}
},
{
"$group": {
"_id": "$string",
"total": {
"$sum": "$integer"
}
}
},
]
result = await q.to_list()
assert len(result) == 3
assert {"_id": "test_1", "total": 2} in result
assert {"_id": "test_2", "total": 6} in result
assert {"_id": "test_3", "total": 3} in result
async def test_aggregate_with_projection_model(preset_documents):
class OutputItem(BaseModel):
id: str = Field(None, alias="_id")
total: int
ids = []
q = Sample.find(Sample.increment >= 4).aggregate(
[{"$group": {"_id": "$string", "total": {"$sum": "$integer"}}}],
projection_model=OutputItem,
)
assert q.get_aggregation_pipeline() == [
{"$match": {"increment": {"$gte": 4}}},
{"$group": {"_id": "$string", "total": {"$sum": "$integer"}}},
{"$project": {"_id": 1, "total": 1}},
]
async for i in q:
if i.id == "test_1":
assert i.total == 2
elif i.id == "test_2":
assert i.total == 6
elif i.id == "test_3":
assert i.total == 3
else:
raise KeyError
ids.append(i.id)
assert set(ids) == {"test_1", "test_2", "test_3"}
async def test_update_query():
q = (Sample.find_many(Sample.integer == 1).update(Set({Sample.integer:
10})).update_query)
assert q == {"$set": {"integer": 10}}
q = (Sample.find_many(Sample.integer == 1).update(
Max({Sample.integer: 10}), Set({Sample.optional: None})).update_query)
assert q == {"$max": {"integer": 10}, "$set": {"optional": None}}
q = (Sample.find_many(Sample.integer == 1).update(
Set({Sample.integer: 10}), Set({Sample.optional: None})).update_query)
assert q == {"$set": {"optional": None}}
q = (Sample.find_many(Sample.integer == 1).update(Max(
{Sample.integer: 10})).update(Set({Sample.optional:
None})).update_query)
assert q == {"$max": {"integer": 10}, "$set": {"optional": None}}
q = (Sample.find_many(Sample.integer == 1).update(Set(
{Sample.integer: 10})).update(Set({Sample.optional:
None})).update_query)
assert q == {"$set": {"optional": None}}
with pytest.raises(TypeError):
Sample.find_many(Sample.integer == 1).update(40).update_query
async def test_find_all(preset_documents):
result = await Sample.find_all().to_list()
assert len(result) == 10
len_result = 0
async for a in Sample.find_all():
assert a in result
len_result += 1
assert len_result == len(result)
async def test_update_many_with_session(preset_documents, session):
q = (Sample.find_many(
Sample.increment > 4).find_many(Sample.nested.optional == None).update(
Set({Sample.increment: 100})).set_session(session=session))
assert q.session == session
q = (Sample.find_many(Sample.increment > 4).find_many(
Sample.nested.optional == None).update(Set({Sample.increment: 100}),
session=session))
assert q.session == session
q = (Sample.find_many(Sample.increment > 4).find_many(
Sample.nested.optional == None,
session=session).update(Set({Sample.increment: 100})))
assert q.session == session
await q # noqa
result = await Sample.find_many(Sample.increment == 100).to_list()
assert len(result) == 3
for sample in result:
assert sample.increment == 100
async def test_sort(preset_documents):
q = Sample.find_many(Sample.integer > 1, sort="-integer")
assert q.sort_expressions == [("integer", SortDirection.DESCENDING)]
q = Sample.find_many(Sample.integer > 1, sort="integer")
assert q.sort_expressions == [("integer", SortDirection.ASCENDING)]
q = Sample.find_many(Sample.integer > 1).sort("-integer")
assert q.sort_expressions == [("integer", SortDirection.DESCENDING)]
q = (
Sample.find_many(Sample.integer > 1)
.find_many(Sample.integer < 100)
.sort("-integer")
)
assert q.sort_expressions == [("integer", SortDirection.DESCENDING)]
result = await Sample.find_many(
Sample.integer > 1, sort="-integer"
).to_list()
i_buf = None
for a in result:
if i_buf is None:
i_buf = a.integer
assert i_buf >= a.integer
i_buf = a.integer
result = await Sample.find_many(
Sample.integer > 1, sort="+integer"
).to_list()
i_buf = None
for a in result:
if i_buf is None:
i_buf = a.integer
assert i_buf <= a.integer
i_buf = a.integer
result = await Sample.find_many(
Sample.integer > 1, sort="integer"
).to_list()
i_buf = None
for a in result:
if i_buf is None:
i_buf = a.integer
assert i_buf <= a.integer
i_buf = a.integer
result = await Sample.find_many(
Sample.integer > 1, sort=-Sample.integer
).to_list()
i_buf = None
for a in result:
if i_buf is None:
i_buf = a.integer
assert i_buf >= a.integer
i_buf = a.integer
with pytest.raises(TypeError):
Sample.find_many(Sample.integer > 1, sort=1)
async def test_aggregate(preset_documents):
q = Sample.aggregate(
[{"$group": {"_id": "$string", "total": {"$sum": "$integer"}}}]
)
assert q.get_aggregation_pipeline() == [
{"$group": {"_id": "$string", "total": {"$sum": "$integer"}}}
]
result = await q.to_list()
assert len(result) == 4
assert {"_id": "test_3", "total": 3} in result
assert {"_id": "test_1", "total": 3} in result
assert {"_id": "test_0", "total": 0} in result
assert {"_id": "test_2", "total": 6} in result
async def test_current_date(session):
q = Sample.find_many(Sample.integer == 1).current_date(
{Sample.timestamp: "timestamp"}, session=session)
assert isinstance(q, UpdateQuery)
assert isinstance(q, UpdateMany)
assert q.session == session
assert q.update_query == {"$currentDate": {"timestamp": "timestamp"}}
q = (Sample.find_many(Sample.integer == 1).update(Max(
{Sample.integer: 10})).current_date({Sample.timestamp: "timestamp"}))
assert isinstance(q, UpdateQuery)
assert isinstance(q, UpdateMany)
assert q.update_query == {
"$max": {
"integer": 10
},
"$currentDate": {
"timestamp": "timestamp"
},
}
async def test_inc(session):
q = Sample.find_many(Sample.integer == 1).inc({Sample.integer: 100},
session=session)
assert isinstance(q, UpdateQuery)
assert isinstance(q, UpdateMany)
assert q.session == session
assert q.update_query == {"$inc": {"integer": 100}}
q = (Sample.find_many(Sample.integer == 1).update(Max(
{Sample.integer: 10})).inc({Sample.integer: 100}))
assert isinstance(q, UpdateQuery)
assert isinstance(q, UpdateMany)
assert q.update_query == {
"$max": {
"integer": 10
},
"$inc": {
"integer": 100
},
}
async def test_find_many(preset_documents):
result = (await Sample.find_many(Sample.integer > 1
).find_many(Sample.nested.optional == None
).to_list()) # noqa
assert len(result) == 2
for a in result:
assert a.integer > 1
assert a.nested.optional is None
len_result = 0
async for a in Sample.find_many(Sample.integer > 1).find_many(
Sample.nested.optional == None): # noqa
assert a in result
len_result += 1
assert len_result == len(result)
async def test_delete_many(preset_documents):
count_before = await Sample.count()
count_find = (await
Sample.find_many(Sample.integer > 1
).find_many(Sample.nested.optional == None
).count()) # noqa
delete_result = await Sample.find_many(Sample.integer > 1).find_many(
Sample.nested.optional == None).delete() # noqa
count_deleted = delete_result.deleted_count
count_after = await Sample.count()
assert count_before - count_find == count_after
assert count_after + count_deleted == count_before
assert isinstance(
Sample.find_many(Sample.integer > 1).find_many(
Sample.nested.optional == None).delete_many(),
DeleteMany,
) # noqa
def sample_doc_not_saved(point):
nested = Nested(
integer=0,
option_1=Option1(s="TEST"),
union=Option1(s="TEST"),
optional=None,
)
geo = GeoObject(coordinates=[
point["longitude"],
point["latitude"],
])
return Sample(
timestamp=datetime.utcnow(),
increment=0,
integer=0,
float_num=0,
string="TEST_NOT_SAVED",
nested=nested,
optional=None,
union=Option1(s="TEST"),
geo=geo,
)
async def preset_documents(point):
docs = []
for i in range(10):
timestamp = datetime.utcnow() - timedelta(days=i)
integer_1: int = i // 3
integer_2: int = i // 2
float_num = integer_1 + 0.3
string: str = f"test_{integer_1}"
option_1 = Option1(s="TEST")
option_2 = Option2(f=3.14)
union = option_1 if i % 2 else option_2
optional = option_2 if not i % 3 else None
geo = GeoObject(
coordinates=[
point["longitude"] + i / 10,
point["latitude"] + i / 10,
]
)
nested = Nested(
integer=integer_2,
option_1=option_1,
union=union,
optional=optional,
)
sample = Sample(
timestamp=timestamp,
increment=i,
integer=integer_1,
float_num=float_num,
string=string,
nested=nested,
optional=optional,
union=union,
geo=geo,
)
docs.append(sample)
await Sample.insert_many(documents=docs)
def test_eq():
q = Sample.find_many(Sample.integer == 1)
assert q.get_filter_query() == {"integer": 1}
def test_lt():
q = Sample.find_many(Sample.integer < 1)
assert q.get_filter_query() == {"integer": {"$lt": 1}}
def test_ne():
q = Sample.find_many(Sample.integer != 1)
assert q.get_filter_query() == {"integer": {"$ne": 1}}
def test_nin():
q = Sample.find_many(NotIn(Sample.integer, [1, 2, 3, 4]))
assert dict(q.get_filter_query()) == {"integer": {"$nin": [1, 2, 3, 4]}}
