diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py
index 40d6d5b64..98f676d52 100644
--- a/libs/langgraph/tests/test_pregel.py
+++ b/libs/langgraph/tests/test_pregel.py
@@ -2901,7 +2901,7 @@ def test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic_inp
     request: pytest.FixtureRequest,
     checkpointer_name: str,
 ) -> None:
-    from pydantic import BaseModel, ConfigDict
+    from pydantic import BaseModel
 
     checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}")
 
@@ -2917,10 +2917,10 @@ def test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic_inp
     class InnerObject(BaseModel):
         yo: int
 
-    class State(BaseModel):
-        model_config = ConfigDict(arbitrary_types_allowed=True)
-
+    class QueryModel(BaseModel):
         query: str
+
+    class State(QueryModel):
         inner: InnerObject
         answer: Optional[str] = None
         docs: Annotated[list[str], sorted_add]
@@ -2930,8 +2930,7 @@ def test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic_inp
         answer: Optional[str] = None
         docs: Optional[list[str]] = None
 
-    class Input(BaseModel):
-        query: str
+    class Input(QueryModel):
         inner: InnerObject
 
     class Output(BaseModel):