Fixed #36210, Refs #36181 -- Allowed Subquery usage in further lookups against composite pks.

Follow-up to 8561100425876bde3be4b2a22324655f74ff9609.

Co-authored-by: Simon Charette <charette.s@gmail.com>

1 files changed, 41 insertions, 4 deletions

diff --git a/tests/composite_pk/test_filter.py b/tests/composite_pk/test_filter.py
index 902bd832ff..d7ecfbec11 100644
--- a/tests/composite_pk/test_filter.py
+++ b/tests/composite_pk/test_filter.py
@@ -1,6 +1,6 @@
 from unittest.mock import patch
 
-from django.db import connection
+from django.db import NotSupportedError, connection
 from django.db.models import (
     Case,
     F,
@@ -14,7 +14,7 @@ from django.db.models import (
 )
 from django.db.models.functions import Cast
 from django.db.models.lookups import Exact
-from django.test import TestCase, skipUnlessDBFeature
+from django.test import TestCase, skipIfDBFeature, skipUnlessDBFeature
 
 from .models import Comment, Tenant, User
 
@@ -492,6 +492,39 @@ class CompositePKFilterTests(TestCase):
                 queryset = Comment.objects.filter(**{f"id{lookup}": subquery})
                 self.assertEqual(queryset.count(), expected_count)
 
+    def test_outer_ref_pk_filter_on_pk_exact(self):
+        subquery = Subquery(User.objects.filter(pk=OuterRef("pk")).values("pk")[:1])
+        qs = Comment.objects.filter(pk=subquery)
+        self.assertEqual(qs.count(), 2)
+
+    @skipUnlessDBFeature("supports_tuple_comparison_against_subquery")
+    def test_outer_ref_pk_filter_on_pk_comparison(self):
+        subquery = Subquery(User.objects.filter(pk=OuterRef("pk")).values("pk")[:1])
+        tests = [
+            ("gt", 0),
+            ("gte", 2),
+            ("lt", 0),
+            ("lte", 2),
+        ]
+        for lookup, expected_count in tests:
+            with self.subTest(f"pk__{lookup}"):
+                qs = Comment.objects.filter(**{f"pk__{lookup}": subquery})
+                self.assertEqual(qs.count(), expected_count)
+
+    @skipIfDBFeature("supports_tuple_comparison_against_subquery")
+    def test_outer_ref_pk_filter_on_pk_comparison_unsupported(self):
+        subquery = Subquery(User.objects.filter(pk=OuterRef("pk")).values("pk")[:1])
+        tests = ["gt", "gte", "lt", "lte"]
+        for lookup in tests:
+            with self.subTest(f"pk__{lookup}"):
+                qs = Comment.objects.filter(**{f"pk__{lookup}": subquery})
+                with self.assertRaisesMessage(
+                    NotSupportedError,
+                    f'"{lookup}" cannot be used to target composite fields '
+                    "through subqueries on this backend",
+                ):
+                    qs.count()
+
     def test_unsupported_rhs(self):
         pk = Exact(F("tenant_id"), 1)
         msg = (
@@ -561,7 +594,11 @@ class CompositePKFilterTests(TestCase):
 @skipUnlessDBFeature("supports_tuple_lookups")
 class CompositePKFilterTupleLookupFallbackTests(CompositePKFilterTests):
     def setUp(self):
-        feature_patch = patch.object(
+        feature_patch_1 = patch.object(
             connection.features, "supports_tuple_lookups", False
         )
-        self.enterContext(feature_patch)
+        feature_patch_2 = patch.object(
+            connection.features, "supports_tuple_comparison_against_subquery", False
+        )
+        self.enterContext(feature_patch_1)
+        self.enterContext(feature_patch_2)