diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index 7de1d8018..f7a527927 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -7619,25 +7619,34 @@ def test_parallel_interrupts_double( def test_pregel_loop_refcount(): - class State(TypedDict): - messages: Annotated[list, add_messages] + gc.collect() + try: + gc.disable() - graph_builder = StateGraph(State) + class State(TypedDict): + messages: Annotated[list, add_messages] - def chatbot(state: State): - return {"messages": [("ai", "HIYA")]} + graph_builder = StateGraph(State) - graph_builder.add_node("chatbot", chatbot) - graph_builder.set_entry_point("chatbot") - graph_builder.set_finish_point("chatbot") - graph = graph_builder.compile() + def chatbot(state: State): + return {"messages": [("ai", "HIYA")]} - for _ in range(5): - graph.invoke({"messages": [{"role": "user", "content": "hi"}]}) - assert ( - len([obj for obj in gc.get_objects() if isinstance(obj, SyncPregelLoop)]) - == 0 - ) - assert ( - len([obj for obj in gc.get_objects() if isinstance(obj, PregelRunner)]) == 0 - ) + graph_builder.add_node("chatbot", chatbot) + graph_builder.set_entry_point("chatbot") + graph_builder.set_finish_point("chatbot") + graph = graph_builder.compile() + + for _ in range(5): + graph.invoke({"messages": [{"role": "user", "content": "hi"}]}) + assert ( + len( + [obj for obj in gc.get_objects() if isinstance(obj, SyncPregelLoop)] + ) + == 0 + ) + assert ( + len([obj for obj in gc.get_objects() if isinstance(obj, PregelRunner)]) + == 0 + ) + finally: + gc.enable() diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py index dc59abcd9..7158edd6e 100644 --- a/libs/langgraph/tests/test_pregel_async.py +++ b/libs/langgraph/tests/test_pregel_async.py @@ -7844,25 +7844,38 @@ async def test_handles_multiple_interrupts_from_tasks() -> None: async def test_pregel_loop_refcount(): - class State(TypedDict): - messages: Annotated[list, add_messages] + gc.collect() + try: + gc.disable() - graph_builder = StateGraph(State) + class State(TypedDict): + messages: Annotated[list, add_messages] - async def chatbot(state: State): - return {"messages": [("ai", "HIYA")]} + graph_builder = StateGraph(State) - graph_builder.add_node("chatbot", chatbot) - graph_builder.set_entry_point("chatbot") - graph_builder.set_finish_point("chatbot") - graph = graph_builder.compile() + async def chatbot(state: State): + return {"messages": [("ai", "HIYA")]} - for _ in range(5): - await graph.ainvoke({"messages": [{"role": "user", "content": "hi"}]}) - assert ( - len([obj for obj in gc.get_objects() if isinstance(obj, AsyncPregelLoop)]) - == 0 - ) - assert ( - len([obj for obj in gc.get_objects() if isinstance(obj, PregelRunner)]) == 0 - ) + graph_builder.add_node("chatbot", chatbot) + graph_builder.set_entry_point("chatbot") + graph_builder.set_finish_point("chatbot") + graph = graph_builder.compile() + + for _ in range(5): + await graph.ainvoke({"messages": [{"role": "user", "content": "hi"}]}) + assert ( + len( + [ + obj + for obj in gc.get_objects() + if isinstance(obj, AsyncPregelLoop) + ] + ) + == 0 + ) + assert ( + len([obj for obj in gc.get_objects() if isinstance(obj, PregelRunner)]) + == 0 + ) + finally: + gc.enable()