Add MongoTestCaseMixin.assertAggregateQuery()

Author: timgraham

django_mongodb_backend/test.py

@@ -0,0 +1,21 @@
+"""Not a public API."""
+
+from bson import SON, ObjectId
+
+
+class MongoTestCaseMixin:
+    maxDiff = None
+
+    def assertAggregateQuery(self, query, expected_collection, expected_pipeline):
+        """
+        Assert that the logged query is equal to:
+            db.{expected_collection}.aggregate({expected_pipeline})
+        """
+        prefix, pipeline = query.split("(", 1)
+        _, collection, operator = prefix.split(".")
+        self.assertEqual(operator, "aggregate")
+        self.assertEqual(collection, expected_collection)
+        self.assertEqual(
+            eval(pipeline[:-1], {"SON": SON, "ObjectId": ObjectId}, {}),  # noqa: S307
+            expected_pipeline,
+        )

tests/lookup_/tests.py

@@ -1,5 +1,7 @@
 from django.test import TestCase
 
+from django_mongodb_backend.test import MongoTestCaseMixin
+
 from .models import Book, Number
 
 
@@ -17,16 +19,23 @@ def test_lte(self):
         self.assertQuerySetEqual(Number.objects.filter(num__lte=3), self.objs[:4])
 
 
-class RegexTests(TestCase):
+class RegexTests(MongoTestCaseMixin, TestCase):
     def test_mql(self):
         # $regexMatch must not cast the input to string, otherwise MongoDB
         # can't use the field's indexes.
         with self.assertNumQueries(1) as ctx:
             list(Book.objects.filter(title__regex="Moby Dick"))
         query = ctx.captured_queries[0]["sql"]
-        self.assertEqual(
+        self.assertAggregateQuery(
             query,
-            "db.lookup__book.aggregate(["
-            "{'$match': {'$expr': {'$regexMatch': {'input': '$title', "
-            "'regex': 'Moby Dick', 'options': ''}}}}])",
+            "lookup__book",
+            [
+                {
+                    "$match": {
+                        "$expr": {
+                            "$regexMatch": {"input": "$title", "regex": "Moby Dick", "options": ""}
+                        }
+                    }
+                }
+            ],
         )