Skip to content

Commit 954c492

Browse files
authored
refactor: Consolidate agent state unit tests (#334)
1 parent f20a405 commit 954c492

File tree

4 files changed

+34
-36
lines changed

4 files changed

+34
-36
lines changed

tests/strands/agent/test_agent_state.py

Lines changed: 34 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,11 @@
22

33
import pytest
44

5+
from strands import Agent, tool
56
from strands.agent.state import AgentState
7+
from strands.types.content import Messages
8+
9+
from ...fixtures.mocked_model_provider import MockedModelProvider
610

711

812
def test_set_and_get():
@@ -109,3 +113,33 @@ def test_initial_state():
109113
assert state.get("key1") == "value1"
110114
assert state.get("key2") == "value2"
111115
assert state.get() == initial
116+
117+
118+
def test_agent_state_update_from_tool():
119+
@tool
120+
def update_state(agent: Agent):
121+
agent.state.set("hello", "world")
122+
agent.state.set("foo", "baz")
123+
124+
agent_messages: Messages = [
125+
{
126+
"role": "assistant",
127+
"content": [{"toolUse": {"name": "update_state", "toolUseId": "123", "input": {}}}],
128+
},
129+
{"role": "assistant", "content": [{"text": "I invoked a tool!"}]},
130+
]
131+
mocked_model_provider = MockedModelProvider(agent_messages)
132+
133+
agent = Agent(
134+
model=mocked_model_provider,
135+
tools=[update_state],
136+
state={"foo": "bar"},
137+
)
138+
139+
assert agent.state.get("hello") is None
140+
assert agent.state.get("foo") == "bar"
141+
142+
agent("Invoke Mocked!")
143+
144+
assert agent.state.get("hello") == "world"
145+
assert agent.state.get("foo") == "baz"

tests/strands/mocked_model_provider/__init__.py

Whitespace-only changes.

tests/strands/mocked_model_provider/test_agent_state_updates.py

Lines changed: 0 additions & 36 deletions
This file was deleted.

0 commit comments

Comments
 (0)