Add tests

tests/test_case_when.py

@@ -2,7 +2,7 @@
 from tortoise import connections
 from tortoise.contrib import test
 from tortoise.expressions import Case, F, Q, When
-from tortoise.functions import Coalesce
+from tortoise.functions import Coalesce, Count
 
 
 class TestCaseWhen(test.TestCase):
@@ -131,3 +131,52 @@ async def test_case_when_in_where(self):
         else:
             expected_sql = "SELECT \"intnum\" \"intnum\" FROM \"intfields\" WHERE CASE WHEN \"intnum\">=8 THEN 'big' WHEN \"intnum\"<=2 THEN 'small' ELSE 'middle' END IN ('big','small')"
         self.assertEqual(sql, expected_sql)
+
+    async def test_annotation_in_when_annotation(self):
+        sql = (
+            IntFields.all()
+            .annotate(intnum_plus_1=F("intnum") + 1)
+            .annotate(bigger_than_10=Case(When(Q(intnum_plus_1__gte=10), then=True), default=False))
+            .values("id", "intnum", "intnum_plus_1", "bigger_than_10")
+            .sql()
+        )
+
+        dialect = self.db.schema_generator.DIALECT
+        if dialect == "mysql":
+            expected_sql = "SELECT `id` `id`,`intnum` `intnum`,`intnum`+1 `intnum_plus_1`,CASE WHEN `intnum`+1>=10 THEN true ELSE false END `bigger_than_10` FROM `intfields`"
+        else:
+            expected_sql = 'SELECT "id" "id","intnum" "intnum","intnum"+1 "intnum_plus_1",CASE WHEN "intnum"+1>=10 THEN true ELSE false END "bigger_than_10" FROM "intfields"'
+        self.assertEqual(sql, expected_sql)
+
+    async def test_func_annotation_in_when_annotation(self):
+        sql = (
+            IntFields.all()
+            .annotate(intnum_col=Coalesce("intnum", 0))
+            .annotate(is_zero=Case(When(Q(intnum_col=0), then=True), default=False))
+            .values("id", "intnum_col", "is_zero")
+            .sql()
+        )
+
+        dialect = self.db.schema_generator.DIALECT
+        if dialect == "mysql":
+            expected_sql = "SELECT `id` `id`,COALESCE(`intnum`,0) `intnum_col`,CASE WHEN COALESCE(`intnum`,0)=0 THEN true ELSE false END `is_zero` FROM `intfields`"
+        else:
+            expected_sql = 'SELECT "id" "id",COALESCE("intnum",0) "intnum_col",CASE WHEN COALESCE("intnum",0)=0 THEN true ELSE false END "is_zero" FROM "intfields"'
+        self.assertEqual(sql, expected_sql)
+
+    async def test_case_when_in_group_by(self):
+        sql = (
+            IntFields.all()
+            .annotate(is_zero=Case(When(Q(intnum=0), then=True), default=False))
+            .annotate(count=Count("id"))
+            .group_by("is_zero")
+            .values("is_zero", "count")
+            .sql()
+        )
+
+        dialect = self.db.schema_generator.DIALECT
+        if dialect == "mysql":
+            expected_sql = "SELECT CASE WHEN `intnum`=0 THEN true ELSE false END `is_zero`,COUNT(`id`) `count` FROM `intfields` GROUP BY `is_zero`"
+        else:
+            expected_sql = 'SELECT CASE WHEN "intnum"=0 THEN true ELSE false END "is_zero",COUNT("id") "count" FROM "intfields" GROUP BY "is_zero"'
+        self.assertEqual(sql, expected_sql)

tests/test_q.py

@@ -1,9 +1,10 @@
+import operator
 from unittest import TestCase as _TestCase
 
 from tests.testmodels import CharFields, IntFields
 from tortoise.contrib.test import TestCase
 from tortoise.exceptions import OperationalError
-from tortoise.expressions import Q, ResolveContext
+from tortoise.expressions import F, Q, ResolveContext
 
 
 class TestQ(_TestCase):
@@ -231,3 +232,21 @@ def test_q_with_blank_or3(self):
         q = Q() | Q(id__gt=5)
         r = q.resolve(self.char_fields_context)
         self.assertEqual(r.where_criterion.get_sql(), '"id">5')
+
+    def test_annotations_resolved(self):
+        q = Q(id__gt=5) | Q(annotated__lt=5)
+        r = q.resolve(
+            ResolveContext(
+                model=IntFields,
+                table=IntFields._meta.basequery,
+                annotations={"annotated": F("annotated")},
+                custom_filters={
+                    "annotated__lt": {
+                        "field": "annotated",
+                        "source_field": "annotated",
+                        "operator": operator.lt,
+                    }
+                },
+            )
+        )
+        self.assertEqual(r.where_criterion.get_sql(), '"id">5 OR "annotated"<5')