Fix unit test.

This commit is contained in:
Andrew Nguonly
2024-10-11 18:53:45 -07:00
parent 2f819a6a9b
commit a277b86fcb
2 changed files with 16 additions and 14 deletions
+14 -12
View File
@@ -400,13 +400,14 @@ class RemoteGraph(PregelProtocol, Runnable):
interrupt_after=interrupt_after, # type: ignore
stream_subgraphs=subgraphs,
):
if chunk.event == INTERRUPT:
raise GraphInterrupt()
if chunk.event == "updates":
if INTERRUPT in chunk.data:
raise GraphInterrupt()
# Don't emit 'updates' events if the original list of stream modes
# didn't include it.
if chunk.event == "updates" and not include_updates:
continue
# Don't emit 'updates' events if the original list of stream
# modes didn't include it.
if not include_updates:
continue
yield chunk
@@ -434,13 +435,14 @@ class RemoteGraph(PregelProtocol, Runnable):
interrupt_after=interrupt_after, # type: ignore
stream_subgraphs=subgraphs,
):
if chunk.event == INTERRUPT:
raise GraphInterrupt()
if chunk.event == "updates":
if INTERRUPT in chunk.data:
raise GraphInterrupt()
# Don't emit 'updates' events if the original list of stream modes
# didn't include it.
if chunk.event == "updates" and not include_updates:
continue
# Don't emit 'updates' events if the original list of stream
# modes didn't include it.
if not include_updates:
continue
yield chunk
+2 -2
View File
@@ -479,7 +479,7 @@ def test_stream():
StreamPart(event="values", data={"chunk": "data2"}),
StreamPart(event="values", data={"chunk": "data3"}),
StreamPart(event="updates", data={"chunk": "data4"}),
StreamPart(event="__interrupt__", data={}),
StreamPart(event="updates", data={"__interrupt__": ()}),
]
# call method / assertions
@@ -527,7 +527,7 @@ async def test_astream():
StreamPart(event="values", data={"chunk": "data2"}),
StreamPart(event="values", data={"chunk": "data3"}),
StreamPart(event="updates", data={"chunk": "data4"}),
StreamPart(event="__interrupt__", data={}),
StreamPart(event="updates", data={"__interrupt__": ()}),
]
mock_async_client.runs.stream.return_value = async_iter