Update tests

This commit is contained in:
Nuno Campos
2024-12-04 14:41:30 -08:00
parent fb01d65dc0
commit 5c7a6689af
4 changed files with 44 additions and 16 deletions
-2
View File
@@ -72,8 +72,6 @@ CONFIG_KEY_CHECKPOINT_ID = sys.intern("checkpoint_id")
CONFIG_KEY_CHECKPOINT_NS = sys.intern("checkpoint_ns")
# holds the current checkpoint_ns, "" for root graph
CONFIG_KEY_NODE_FINISHED = sys.intern("__pregel_node_finished")
# callback to be called when a node is finished
CONFIG_KEY_RESUME_VALUE = sys.intern("__pregel_resume_value")
# holds the value that "answers" an interrupt() call
CONFIG_KEY_WRITES = sys.intern("__pregel_writes")
# read-only list of existing task writes
+18
View File
@@ -35,3 +35,21 @@ class AnyDict(dict):
return False
else:
return True
class AnyList(list):
def __init__(self, *args, **kwargs) -> None:
super().__init__(*args, **kwargs)
def __eq__(self, other: object) -> bool:
if not self and isinstance(other, list):
return True
if not isinstance(other, list) or len(self) != len(other):
return False
for i, v in enumerate(self):
if v == other[i]:
continue
else:
return False
else:
return True
+13 -7
View File
@@ -15,7 +15,7 @@ from langgraph.graph.state import StateGraph
from langgraph.pregel import Pregel
from langgraph.scheduler.kafka import serde
from langgraph.scheduler.kafka.types import MessageToOrchestrator, Topics
from tests.any import AnyDict
from tests.any import AnyDict, AnyList
from tests.drain import drain_topics_async
from tests.messages import _AnyIdAIMessage, _AnyIdHumanMessage
@@ -196,7 +196,8 @@ async def test_subgraph_w_interrupt(
"__pregel_resuming": False,
"__pregel_store": None,
"__pregel_task_id": history[0].tasks[0].id,
"__pregel_resume_value": None,
"__pregel_scratchpad": {},
"__pregel_writes": AnyList(),
"checkpoint_id": None,
"checkpoint_map": {
"": history[0].config["configurable"]["checkpoint_id"]
@@ -261,7 +262,8 @@ async def test_subgraph_w_interrupt(
"__pregel_resuming": False,
"__pregel_store": None,
"__pregel_task_id": history[0].tasks[0].id,
"__pregel_resume_value": None,
"__pregel_scratchpad": {},
"__pregel_writes": AnyList(),
"checkpoint_id": c.config["configurable"]["checkpoint_id"],
"checkpoint_map": {
"": history[0].config["configurable"]["checkpoint_id"]
@@ -356,7 +358,8 @@ async def test_subgraph_w_interrupt(
"__pregel_resuming": False,
"__pregel_store": None,
"__pregel_task_id": history[0].tasks[0].id,
"__pregel_resume_value": None,
"__pregel_scratchpad": {},
"__pregel_writes": AnyList(),
"checkpoint_id": c.config["configurable"]["checkpoint_id"],
"checkpoint_map": {
"": history[0].config["configurable"]["checkpoint_id"]
@@ -461,7 +464,8 @@ async def test_subgraph_w_interrupt(
"__pregel_resuming": True,
"__pregel_store": None,
"__pregel_task_id": history[1].tasks[0].id,
"__pregel_resume_value": None,
"__pregel_scratchpad": {},
"__pregel_writes": AnyList(),
"checkpoint_id": None,
"checkpoint_map": {
"": history[1].config["configurable"]["checkpoint_id"]
@@ -521,7 +525,8 @@ async def test_subgraph_w_interrupt(
"__pregel_resuming": True,
"__pregel_store": None,
"__pregel_task_id": history[1].tasks[0].id,
"__pregel_resume_value": None,
"__pregel_scratchpad": {},
"__pregel_writes": AnyList(),
"checkpoint_id": c.config["configurable"]["checkpoint_id"],
"checkpoint_map": {
"": history[1].config["configurable"]["checkpoint_id"]
@@ -637,7 +642,8 @@ async def test_subgraph_w_interrupt(
"__pregel_resuming": True,
"__pregel_store": None,
"__pregel_task_id": history[1].tasks[0].id,
"__pregel_resume_value": None,
"__pregel_scratchpad": {},
"__pregel_writes": AnyList(),
"checkpoint_id": c.config["configurable"]["checkpoint_id"],
"checkpoint_map": {
"": history[1].config["configurable"]["checkpoint_id"]
@@ -15,7 +15,7 @@ from langgraph.pregel import Pregel
from langgraph.scheduler.kafka import serde
from langgraph.scheduler.kafka.default_sync import DefaultProducer
from langgraph.scheduler.kafka.types import MessageToOrchestrator, Topics
from tests.any import AnyDict
from tests.any import AnyDict, AnyList
from tests.drain import drain_topics
from tests.messages import _AnyIdAIMessage, _AnyIdHumanMessage
@@ -195,7 +195,8 @@ def test_subgraph_w_interrupt(
"__pregel_resuming": False,
"__pregel_store": None,
"__pregel_task_id": history[0].tasks[0].id,
"__pregel_resume_value": None,
"__pregel_scratchpad": {},
"__pregel_writes": AnyList(),
"checkpoint_id": None,
"checkpoint_map": {
"": history[0].config["configurable"]["checkpoint_id"]
@@ -260,7 +261,8 @@ def test_subgraph_w_interrupt(
"__pregel_dedupe_tasks": True,
"__pregel_resuming": False,
"__pregel_task_id": history[0].tasks[0].id,
"__pregel_resume_value": None,
"__pregel_scratchpad": {},
"__pregel_writes": AnyList(),
"checkpoint_id": c.config["configurable"]["checkpoint_id"],
"checkpoint_map": {
"": history[0].config["configurable"]["checkpoint_id"]
@@ -355,7 +357,8 @@ def test_subgraph_w_interrupt(
"__pregel_store": None,
"__pregel_resuming": False,
"__pregel_task_id": history[0].tasks[0].id,
"__pregel_resume_value": None,
"__pregel_scratchpad": {},
"__pregel_writes": AnyList(),
"checkpoint_id": c.config["configurable"]["checkpoint_id"],
"checkpoint_map": {
"": history[0].config["configurable"]["checkpoint_id"]
@@ -459,7 +462,8 @@ def test_subgraph_w_interrupt(
"__pregel_store": None,
"__pregel_resuming": True,
"__pregel_task_id": history[1].tasks[0].id,
"__pregel_resume_value": None,
"__pregel_scratchpad": {},
"__pregel_writes": AnyList(),
"checkpoint_id": None,
"checkpoint_map": {
"": history[1].config["configurable"]["checkpoint_id"]
@@ -519,7 +523,8 @@ def test_subgraph_w_interrupt(
"__pregel_store": None,
"__pregel_resuming": True,
"__pregel_task_id": history[1].tasks[0].id,
"__pregel_resume_value": None,
"__pregel_scratchpad": {},
"__pregel_writes": AnyList(),
"checkpoint_id": c.config["configurable"]["checkpoint_id"],
"checkpoint_map": {
"": history[1].config["configurable"]["checkpoint_id"]
@@ -635,7 +640,8 @@ def test_subgraph_w_interrupt(
"__pregel_resuming": True,
"__pregel_store": None,
"__pregel_task_id": history[1].tasks[0].id,
"__pregel_resume_value": None,
"__pregel_scratchpad": {},
"__pregel_writes": AnyList(),
"checkpoint_id": c.config["configurable"]["checkpoint_id"],
"checkpoint_map": {
"": history[1].config["configurable"]["checkpoint_id"]