diff --git a/mongokat/collection.py b/mongokat/collection.py index 7133d45..a4fc45a 100644 --- a/mongokat/collection.py +++ b/mongokat/collection.py @@ -133,6 +133,12 @@ def exists(self, query, **args): def count(self, *args, **kwargs): return self._collection_with_options(kwargs).count(*args, **kwargs) + def count_documents(self, *args, **kwargs): + return self._collection_with_options(kwargs).count_documents(*args, **kwargs) + + def estimated_document_count(self, **kwargs): + return self._collection_with_options(kwargs).estimated_document_count(**kwargs) + def distinct(self, *args, **kwargs): return self._collection_with_options(kwargs).distinct(*args, **kwargs) diff --git a/tests/test_api.py b/tests/test_api.py index 5701ba9..e2aa8aa 100644 --- a/tests/test_api.py +++ b/tests/test_api.py @@ -89,12 +89,16 @@ def test_document_common_methods(Sample): import collections assert Sample.collection.find().count() == 0 + assert Sample.collection.find().count_documents({}) == 0 + assert Sample.collection.find().estimated_document_count() == 0 # Instanciate new_object = Sample({"name": "XXX", "url": "http://example.com"}) # Should not save to DB yet. assert Sample.collection.find().count() == 0 + assert Sample.collection.find().count_documents({}) == 0 + assert Sample.collection.find().estimated_document_count() == 0 # Now save() new_object.save() @@ -106,6 +110,8 @@ def test_document_common_methods(Sample): assert type(new_object["_id"]) == ObjectId assert Sample.collection.find().count() == 1 + assert Sample.collection.find().count_documents({}) == 1 + assert Sample.collection.find().estimated_document_count() == 1 db_object = Sample.collection.find_one() assert type(db_object) == dict assert db_object["name"] == "XXX" @@ -115,6 +121,8 @@ def test_document_common_methods(Sample): assert type(inserted_object) == ObjectId assert Sample.collection.find().count() == 2 + assert Sample.collection.find().count_documents({}) == 2 + assert Sample.collection.find().estimated_document_count() == 2 # Find back with different methods orm_object = Sample.find_by_id(db_object["_id"]) @@ -168,11 +176,15 @@ def test_document_common_methods(Sample): orm_object.save() assert Sample.collection.find().count() == 2 + assert Sample.collection.find().count_documents({}) == 2 + assert Sample.collection.find().estimated_document_count() == 2 # FIXME not anymore as we are requesting _id for each query # orm_object.save(force=True) # assert Sample.collection.find().count() == 3 + # assert Sample.collection.find().count_documents({}) == 3 + # assert Sample.collection.find().estimated_document_count() == 3 orm_object = Sample.find_by_id(db_object["_id"], fields=["url", "_id"]) assert dict(orm_object) == {"url": "http://example.com", "_id": db_object["_id"]} @@ -186,6 +198,9 @@ def test_document_common_methods(Sample): orm_object.save() assert Sample.collection.find().count() == 2 + assert Sample.collection.find().count_documents({}) == 2 + assert Sample.collection.find().estimated_document_count() == 2 + db_object = Sample.collection.find_one({"_id": db_object["_id"]}) assert "name" in db_object @@ -193,6 +208,8 @@ def test_document_common_methods(Sample): # Should not add anything new assert Sample.collection.find().count() == 2 + assert Sample.collection.find().count_documents({}) == 2 + assert Sample.collection.find().estimated_document_count() == 2 db_object = Sample.collection.find_one({"_id": db_object["_id"]}) assert "name" not in db_object