fix: fix nested MPTT connections when trying to access totalCount

strawberry-graphql · Jun 12, 2024 · 98d9c9c · 98d9c9c
1 parent e318ad2
commit 98d9c9c
Show file tree

Hide file tree

Showing 2 changed files with 86 additions and 3 deletions.
diff --git a/strawberry_django/relay.py b/strawberry_django/relay.py
@@ -102,7 +102,7 @@ def resolve_connection(
 
         if isinstance(nodes, models.QuerySet) and is_optimized_by_prefetching(nodes):
             try:
-                return cls.resolve_connection_from_cache(
+                conn = cls.resolve_connection_from_cache(
                     nodes,
                     info=info,
                     before=before,
@@ -120,6 +120,10 @@ def resolve_connection(
                     RuntimeWarning,
                     stacklevel=2,
                 )
+            else:
+                conn = cast(Self, conn)
+                conn.nodes = nodes
+                return conn
 
         conn = super().resolve_connection(
             nodes,

diff --git a/tests/relay/mptt/test_nested_children.py b/tests/relay/mptt/test_nested_children.py
@@ -1,5 +1,6 @@
 import pytest
 import strawberry
+from strawberry.relay.utils import to_base64
 
 import strawberry_django
 from strawberry_django.optimizer import DjangoOptimizerExtension
@@ -19,10 +20,11 @@ class Query:
 @pytest.mark.django_db(transaction=True)
 def test_nested_children_total_count():
     parent = MPTTAuthor.objects.create(name="Parent")
-    MPTTAuthor.objects.create(name="Child", parent=parent)
+    child1 = MPTTAuthor.objects.create(name="Child1", parent=parent)
+    child2 = MPTTAuthor.objects.create(name="Child2", parent=parent)
     query = """\
     query {
-      authors(last: 1) {
+      authors(first: 1) {
         totalCount
         edges {
           node {
@@ -45,3 +47,80 @@ def test_nested_children_total_count():
 
     result = schema.execute_sync(query)
     assert not result.errors
+    assert result.data == {
+        "authors": {
+            "totalCount": 3,
+            "edges": [
+                {
+                    "node": {
+                        "id": to_base64("MPTTAuthorType", parent.pk),
+                        "name": "Parent",
+                        "children": {
+                            "totalCount": 2,
+                            "edges": [
+                                {
+                                    "node": {
+                                        "id": to_base64("MPTTAuthorType", child1.pk),
+                                        "name": "Child1",
+                                    }
+                                },
+                                {
+                                    "node": {
+                                        "id": to_base64("MPTTAuthorType", child2.pk),
+                                        "name": "Child2",
+                                    }
+                                },
+                            ],
+                        },
+                    }
+                }
+            ],
+        }
+    }
+
+
+@pytest.mark.django_db(transaction=True)
+def test_nested_children_total_count_no_children():
+    parent = MPTTAuthor.objects.create(name="Parent")
+    query = """\
+    query {
+      authors {
+        totalCount
+        edges {
+          node {
+            id
+            name
+            children {
+              totalCount
+              edges {
+                node {
+                  id
+                  name
+                }
+              }
+            }
+          }
+        }
+      }
+    }
+    """
+
+    result = schema.execute_sync(query)
+    assert not result.errors
+    assert result.data == {
+        "authors": {
+            "totalCount": 1,
+            "edges": [
+                {
+                    "node": {
+                        "id": to_base64("MPTTAuthorType", parent.pk),
+                        "name": "Parent",
+                        "children": {
+                            "totalCount": 0,
+                            "edges": [],
+                        },
+                    }
+                }
+            ],
+        }
+    }