Fix stream record test

Author: Cito

src/graphql/execution/incremental_publisher.py

@@ -9,6 +9,7 @@
     AsyncGenerator,
     AsyncIterator,
     Awaitable,
+    Collection,
     NamedTuple,
     Sequence,
     Union,
@@ -328,7 +329,7 @@ def __ne__(self, other: object) -> bool:
 class InitialResult(NamedTuple):
     """The state of the initial result"""
 
-    children: set[IncrementalDataRecord]
+    children: dict[IncrementalDataRecord, None]
     is_completed: bool
 
 
@@ -356,16 +357,19 @@ class IncrementalPublisher:
 
     Each Incremental Data record also contains similar metadata, i.e. these records also
     contain similar ``is_completed`` and ``children`` properties.
+
+    Note: Instead of sets we use dicts (with values set to None) which preserve order
+    and thereby achieve more deterministic results.
     """
 
     _initial_result: InitialResult
-    _released: set[IncrementalDataRecord]
-    _pending: set[IncrementalDataRecord]
+    _released: dict[IncrementalDataRecord, None]
+    _pending: dict[IncrementalDataRecord, None]
 
     def __init__(self) -> None:
-        self._initial_result = InitialResult(set(), False)
-        self._released = set()
-        self._pending = set()
+        self._initial_result = InitialResult({}, False)
+        self._released = {}
+        self._pending = {}
         self._resolve = Event()
         self._tasks: set[Awaitable] = set()
 
@@ -382,13 +386,10 @@ async def subscribe(
             while not is_done:
                 released = self._released
                 for item in released:
-                    pending.discard(item)
-                self._released = (
-                    set()
-                )  # TODO, solve differently? use clear() and local variable?
+                    del pending[item]
+                self._released = {}
 
                 result = self._get_incremental_result(released)
-                released.clear()
 
                 if not self.has_next():
                     is_done = True
@@ -420,7 +421,8 @@ def prepare_new_deferred_fragment_record(
     ) -> DeferredFragmentRecord:
         deferred_fragment_record = DeferredFragmentRecord(label, path, parent_context)
 
-        (parent_context or self._initial_result).children.add(deferred_fragment_record)
+        context = parent_context or self._initial_result
+        context.children[deferred_fragment_record] = None
         return deferred_fragment_record
 
     def prepare_new_stream_items_record(
@@ -434,7 +436,8 @@ def prepare_new_stream_items_record(
             label, path, parent_context, async_iterator
         )
 
-        (parent_context or self._initial_result).children.add(stream_items_record)
+        context = parent_context or self._initial_result
+        context.children[stream_items_record] = None
         return stream_items_record
 
     def complete_deferred_fragment_record(
@@ -484,7 +487,7 @@ def filter(
 
             self._delete(child)
             parent = child.parent_context or self._initial_result
-            parent.children.discard(child)
+            del parent.children[child]
 
             if isinstance(child, StreamItemsRecord):
                 async_iterator = child.async_iterator
@@ -504,26 +507,26 @@ def _reset(self) -> None:
         self._resolve.clear()
 
     def _introduce(self, item: IncrementalDataRecord) -> None:
-        self._pending.add(item)
+        self._pending[item] = None
 
     def _release(self, item: IncrementalDataRecord) -> None:
         if item in self._pending:
-            self._pending.remove(item)
-            self._released.add(item)
+            del self._pending[item]
+            self._released[item] = None
             self._trigger()
 
     def _push(self, item: IncrementalDataRecord) -> None:
-        self._released.add(item)
-        self._pending.add(item)
+        self._released[item] = None
+        self._pending[item] = None
         self._trigger()
 
     def _delete(self, item: IncrementalDataRecord) -> None:
-        self._released.discard(item)
-        self._pending.discard(item)
+        del self._released[item]
+        del self._pending[item]
         self._trigger()
 
     def _get_incremental_result(
-        self, completed_records: set[IncrementalDataRecord]
+        self, completed_records: Collection[IncrementalDataRecord]
     ) -> SubsequentIncrementalExecutionResult | None:
         incremental_results: list[IncrementalResult] = []
         encountered_completed_async_iterator = False
@@ -574,13 +577,13 @@ def _publish(self, incremental_data_record: IncrementalDataRecord) -> None:
 
     def _get_descendants(
         self,
-        children: set[IncrementalDataRecord],
-        descendants: set[IncrementalDataRecord] | None = None,
-    ) -> set[IncrementalDataRecord]:
+        children: dict[IncrementalDataRecord, None],
+        descendants: dict[IncrementalDataRecord, None] | None = None,
+    ) -> dict[IncrementalDataRecord, None]:
         if descendants is None:
-            descendants = set()
+            descendants = {}
         for child in children:
-            descendants.add(child)
+            descendants[child] = None
             self._get_descendants(child.children, descendants)
         return descendants
 
@@ -604,7 +607,7 @@ class DeferredFragmentRecord:
     path: list[str | int]
     data: dict[str, Any] | None
     parent_context: IncrementalDataRecord | None
-    children: set[IncrementalDataRecord]
+    children: dict[IncrementalDataRecord, None]
     is_completed: bool
 
     def __init__(
@@ -617,7 +620,7 @@ def __init__(
         self.path = path.as_list() if path else []
         self.parent_context = parent_context
         self.errors = []
-        self.children = set()
+        self.children = {}
         self.is_completed = False
         self.data = None
 
@@ -641,7 +644,7 @@ class StreamItemsRecord:
     path: list[str | int]
     items: list[str] | None
     parent_context: IncrementalDataRecord | None
-    children: set[IncrementalDataRecord]
+    children: dict[IncrementalDataRecord, None]
     async_iterator: AsyncIterator[Any] | None
     is_completed_async_iterator: bool
     is_completed: bool
@@ -658,7 +661,7 @@ def __init__(
         self.parent_context = parent_context
         self.async_iterator = async_iterator
         self.errors = []
-        self.children = set()
+        self.children = {}
         self.is_completed = False
         self.items = None
 

tests/execution/test_defer.py

@@ -1,7 +1,6 @@
 from __future__ import annotations
 
 from asyncio import sleep
-from operator import itemgetter
 from typing import Any, AsyncGenerator, NamedTuple
 
 import pytest
@@ -515,8 +514,6 @@ async def can_defer_a_fragment_within_an_already_deferred_fragment():
             """
         )
         result = await complete(document)
-        # incremental results can sometimes arrive in different order which is ok
-        result[1]["incremental"].sort(key=itemgetter("label"))
 
         assert result == [
             {"data": {"hero": {}}, "hasNext": True},

tests/execution/test_stream.py

@@ -6,7 +6,6 @@
 import pytest
 from graphql.error import GraphQLError
 from graphql.execution import (
-    ExecutionContext,
     ExecutionResult,
     ExperimentalIncrementalExecutionResults,
     IncrementalStreamResult,
@@ -173,13 +172,9 @@ def can_format_and_print_incremental_stream_result():
         )
 
     def can_print_stream_record():
-        context = ExecutionContext.build(schema, parse("{ hero { id } }"))
-        assert isinstance(context, ExecutionContext)
-        record = StreamItemsRecord(None, None, None, None, context)
+        record = StreamItemsRecord(None, None, None, None)
         assert str(record) == "StreamItemsRecord(path=[])"
-        record = StreamItemsRecord(
-            "foo", Path(None, "bar", "Bar"), None, record, context
-        )
+        record = StreamItemsRecord("foo", Path(None, "bar", "Bar"), record, None)
         assert (
             str(record) == "StreamItemsRecord("
             "path=['bar'], label='foo', parent_context)"