Warn invalid state

tests/test_state.py

@@ -0,0 +1,48 @@
+from typing import Annotated as Annotated2
+from typing import Any
+
+import pytest
+from pydantic.v1 import BaseModel
+from typing_extensions import Annotated, TypedDict
+
+from langgraph.graph.state import _warn_invalid_state_schema
+
+
+class State(BaseModel):
+    foo: str
+    bar: int
+
+
+class State2(TypedDict):
+    foo: str
+    bar: int
+
+
+@pytest.mark.parametrize(
+    "schema",
+    [
+        {"foo": "bar"},
+        ["hi", lambda x, y: x + y],
+        State(foo="bar", bar=1),
+        State2(foo="bar", bar=1),
+    ],
+)
+def test_warns_invalid_schema(schema: Any):
+    with pytest.warns(UserWarning):
+        _warn_invalid_state_schema(schema)
+
+
+@pytest.mark.parametrize(
+    "schema",
+    [
+        Annotated[dict, lambda x, y: y],
+        Annotated2[list, lambda x, y: y],
+        dict,
+        State,
+        State2,
+    ],
+)
+def test_doesnt_warn_valid_schema(schema: Any):
+    # Assert the function does not raise a warning
+    with pytest.warns(None):
+        _warn_invalid_state_schema(schema)