diff --git a/pyproject.toml b/pyproject.toml
index 3f0ee6c..ae78bc1 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -60,8 +60,12 @@ lint.ignore = [
 [tool.ruff.lint.pydocstyle]
 convention = "google"
 
+[tool.pytest.ini_options]
+asyncio_mode = "auto"
+
 [dependency-groups]
 dev = [
     "langgraph-cli[inmem]>=0.1.71",
     "pytest>=8.3.5",
+    "pytest-asyncio>=1.0.0",
 ]
diff --git a/src/react_agent/graph.py b/src/react_agent/graph.py
index 9f3a1ae..cb2ff49 100644
--- a/src/react_agent/graph.py
+++ b/src/react_agent/graph.py
@@ -65,7 +65,11 @@ async def call_model(state: State) -> Dict[str, List[AIMessage]]:
 
 # Define a new graph
 
-builder = StateGraph(State, input=InputState, config_schema=Configuration)
+builder = StateGraph(
+    State,
+    input_schema=InputState,
+    config_schema=Configuration
+)
 
 # Define the two nodes we will cycle between
 builder.add_node(call_model)
diff --git a/tests/integration_tests/test_graph.py b/tests/integration_tests/test_graph.py
index cd90afb..6ea83bc 100644
--- a/tests/integration_tests/test_graph.py
+++ b/tests/integration_tests/test_graph.py
@@ -1,14 +1,17 @@
 import pytest
 from langsmith import unit
+from langchain_core.messages import HumanMessage
 
 from react_agent import graph
+from react_agent.state import InputState
 
 
-@pytest.mark.asyncio
 @unit
 async def test_react_agent_simple_passthrough() -> None:
+    message = HumanMessage(content="Who is the founder of LangChain?")
+    input_state = InputState(messages=[message])
     res = await graph.ainvoke(
-        {"messages": [("user", "Who is the founder of LangChain?")]},
+        input_state,
         {"configurable": {"system_prompt": "You are a helpful AI assistant."}},
     )