From 510dbd3dd7eb10f801f2bd7379ff4d93760e27d5 Mon Sep 17 00:00:00 2001
From: Christoph Zwerschke <cito@online.de>
Date: Sun, 11 Aug 2024 00:26:37 +0200
Subject: [PATCH] Fix execute_stream_async_iterator

---
 src/graphql/execution/execute.py | 19 +++++++++++--------
 tests/execution/test_stream.py   | 10 ++++++++--
 2 files changed, 19 insertions(+), 10 deletions(-)

diff --git a/src/graphql/execution/execute.py b/src/graphql/execution/execute.py
index 47786357..3cbfda72 100644
--- a/src/graphql/execution/execute.py
+++ b/src/graphql/execution/execute.py
@@ -1683,6 +1683,7 @@ async def execute_stream_async_iterator(
         index = initial_index
         previous_incremental_data_record = parent_context
 
+        done = False
         while True:
             item_path = Path(path, index, None)
             incremental_data_record = (
@@ -1692,7 +1693,7 @@ async def execute_stream_async_iterator(
             )
 
             try:
-                data = await self.execute_stream_async_iterator_item(
+                completed_item = await self.execute_stream_async_iterator_item(
                     async_iterator,
                     field_group,
                     info,
@@ -1701,8 +1702,6 @@ async def execute_stream_async_iterator(
                     path,
                     item_path,
                 )
-            except StopAsyncIteration:
-                break
             except GraphQLError as error:
                 incremental_publisher.add_field_error(incremental_data_record, error)
                 incremental_publisher.filter(path, incremental_data_record)
@@ -1716,12 +1715,16 @@ async def execute_stream_async_iterator(
                     # so we need to remember that this iterator is already canceled
                     self._canceled_iterators.add(async_iterator)
                 break
-            else:
-                incremental_publisher.complete_stream_items_record(
-                    incremental_data_record,
-                    [data],
-                )
+            except StopAsyncIteration:
+                done = True
+
+            incremental_publisher.complete_stream_items_record(
+                incremental_data_record,
+                [completed_item],
+            )
 
+            if done:
+                break
             previous_incremental_data_record = incremental_data_record
             index += 1
 
diff --git a/tests/execution/test_stream.py b/tests/execution/test_stream.py
index a9667185..d4dea0b0 100644
--- a/tests/execution/test_stream.py
+++ b/tests/execution/test_stream.py
@@ -743,6 +743,9 @@ async def friend_list(_info):
                         "path": ["friendList", 2],
                     }
                 ],
+                "hasNext": True,
+            },
+            {
                 "hasNext": False,
             },
         ]
@@ -783,6 +786,9 @@ async def friend_list(_info):
                         "path": ["friendList", 2],
                     }
                 ],
+                "hasNext": True,
+            },
+            {
                 "hasNext": False,
             },
         ]
@@ -856,10 +862,10 @@ async def friend_list(_info):
                             "path": ["friendList", 2],
                         }
                     ],
-                    "hasNext": False,
+                    "hasNext": True,
                 },
             },
-            {"done": True, "value": None},
+            {"done": False, "value": {"hasNext": False}},
             {"done": True, "value": None},
         ]