Add tests for polymorphic foreign key optimization

Author: diesieben07

tests/polymorphism/test_optimizer.py

@@ -193,7 +193,7 @@ def test_polymorphic_relation():
     """
 
     # Company, ContentType, base table, two subtables = 5 queries
-    with assert_num_queries(1):
+    with assert_num_queries(5):
         result = schema.execute_sync(query)
     assert not result.errors
     assert result.data == {
@@ -218,7 +218,6 @@ def test_polymorphic_relation():
     }
 
 
-
 @pytest.mark.django_db(transaction=True)
 def test_polymorphic_nested_list():
     company = Company.objects.create(name="Company")

tests/polymorphism_custom/models.py

@@ -1,7 +1,16 @@
 from django.db import models
 
 
+class Company(models.Model):
+    name = models.CharField(max_length=100)
+    main_project = models.ForeignKey('CustomPolyProject', null=True, blank=True, on_delete=models.CASCADE)
+
+    class Meta:
+        ordering = ('name',)
+
+
 class CustomPolyProject(models.Model):
+    company = models.ForeignKey(Company, null=True, blank=True, on_delete=models.CASCADE, related_name='projects')
     topic = models.CharField(max_length=30)
     artist = models.CharField(max_length=30, blank=True)
     supervisor = models.CharField(max_length=30, blank=True)

tests/polymorphism_custom/schema.py

@@ -8,7 +8,7 @@
 from strawberry_django.pagination import OffsetPaginated
 from strawberry_django.relay import ListConnectionWithTotalCount
 
-from .models import CustomPolyProject
+from .models import CustomPolyProject, Company
 
 
 @strawberry_django.interface(CustomPolyProject)
@@ -37,8 +37,16 @@ class ResearchProjectType(ProjectType):
     supervisor: strawberry.auto
 
 
+@strawberry_django.type(Company)
+class CompanyType:
+    name: strawberry.auto
+    main_project: ProjectType | None
+    projects: list[ProjectType]
+
+
 @strawberry.type
 class Query:
+    companies: list[CompanyType] = strawberry_django.field()
     projects: list[ProjectType] = strawberry_django.field()
     projects_paginated: list[ProjectType] = strawberry_django.field(pagination=True)
     projects_offset_paginated: OffsetPaginated[ProjectType] = (

tests/polymorphism_custom/test_optimizer.py

@@ -4,7 +4,7 @@
 
 from tests.utils import assert_num_queries
 
-from .models import CustomPolyProject
+from .models import CustomPolyProject, Company
 from .schema import schema
 
 
@@ -209,3 +209,102 @@ def test_polymorphic_interface_connection():
             ],
         }
     }
+
+
+@pytest.mark.django_db(transaction=True)
+def test_polymorphic_relation():
+    ap = CustomPolyProject.objects.create(topic="Art", artist="Artist")
+    art_company = Company.objects.create(name="ArtCompany", main_project=ap)
+
+    rp = CustomPolyProject.objects.create(topic="Research", supervisor="Supervisor")
+    research_company = Company.objects.create(name="ResearchCompany", main_project=rp)
+
+    query = """\
+    query {
+      companies {
+        name
+        mainProject {
+            __typename
+            topic
+            ... on ArtProjectType {
+              artist
+            }
+            ... on ResearchProjectType {
+              supervisor
+            }
+          }
+      }
+    }
+    """
+
+    with assert_num_queries(2):
+        result = schema.execute_sync(query)
+    assert not result.errors
+    assert result.data == {
+        "companies": [
+            {
+                "name": art_company.name,
+                "mainProject": {
+                    "__typename": "ArtProjectType",
+                    "topic": ap.topic,
+                    "artist": ap.artist,
+                },
+            },
+            {
+                "name": research_company.name,
+                "mainProject": {
+                    "__typename": "ResearchProjectType",
+                    "topic": rp.topic,
+                    "supervisor": rp.supervisor,
+                },
+            }
+        ]
+    }
+
+
+@pytest.mark.django_db(transaction=True)
+def test_polymorphic_nested_list():
+    company = Company.objects.create(name="Company")
+    ap = CustomPolyProject.objects.create(company=company, topic="Art", artist="Artist")
+    rp = CustomPolyProject.objects.create(company=company, topic="Research", supervisor="Supervisor")
+
+    query = """\
+    query {
+      companies {
+        name
+        projects {
+            __typename
+            topic
+            ... on ArtProjectType {
+              artist
+            }
+            ... on ResearchProjectType {
+              supervisor
+            }
+          }
+      }
+    }
+    """
+
+    with assert_num_queries(2):
+        result = schema.execute_sync(query)
+    assert not result.errors
+    assert result.data == {
+        "companies": [
+            {
+                "name": "Company",
+                "projects": [
+                    {
+                        "__typename": "ArtProjectType",
+                        "topic": ap.topic,
+                        "artist": ap.artist,
+                    },
+                    {
+                        "__typename": "ResearchProjectType",
+                        "topic": rp.topic,
+                        "supervisor": rp.supervisor,
+                    },
+                ],
+            }
+        ]
+    }