1
1
import pytest
2
2
3
3
from strands import Agent , tool
4
+ from strands .experimental .hooks import AfterModelInvocationEvent , BeforeModelInvocationEvent
5
+ from strands .hooks import AfterInvocationEvent , AgentInitializedEvent , BeforeInvocationEvent , MessageAddedEvent
4
6
from strands .multiagent .graph import GraphBuilder
5
7
from strands .types .content import ContentBlock
8
+ from tests .fixtures .mock_hook_provider import MockHookProvider
6
9
7
10
8
11
@tool
@@ -18,49 +21,59 @@ def multiply_numbers(x: int, y: int) -> int:
18
21
19
22
20
23
@pytest .fixture
21
- def math_agent ():
24
+ def hook_provider ():
25
+ return MockHookProvider ("all" )
26
+
27
+
28
+ @pytest .fixture
29
+ def math_agent (hook_provider ):
22
30
"""Create an agent specialized in mathematical operations."""
23
31
return Agent (
24
32
model = "us.amazon.nova-pro-v1:0" ,
25
33
system_prompt = "You are a mathematical assistant. Always provide clear, step-by-step calculations." ,
34
+ hooks = [hook_provider ],
26
35
tools = [calculate_sum , multiply_numbers ],
27
36
)
28
37
29
38
30
39
@pytest .fixture
31
- def analysis_agent ():
40
+ def analysis_agent (hook_provider ):
32
41
"""Create an agent specialized in data analysis."""
33
42
return Agent (
34
43
model = "us.amazon.nova-pro-v1:0" ,
44
+ hooks = [hook_provider ],
35
45
system_prompt = "You are a data analysis expert. Provide insights and interpretations of numerical results." ,
36
46
)
37
47
38
48
39
49
@pytest .fixture
40
- def summary_agent ():
50
+ def summary_agent (hook_provider ):
41
51
"""Create an agent specialized in summarization."""
42
52
return Agent (
43
53
model = "us.amazon.nova-lite-v1:0" ,
54
+ hooks = [hook_provider ],
44
55
system_prompt = "You are a summarization expert. Create concise, clear summaries of complex information." ,
45
56
)
46
57
47
58
48
59
@pytest .fixture
49
- def validation_agent ():
60
+ def validation_agent (hook_provider ):
50
61
"""Create an agent specialized in validation."""
51
62
return Agent (
52
63
model = "us.amazon.nova-pro-v1:0" ,
64
+ hooks = [hook_provider ],
53
65
system_prompt = "You are a validation expert. Check results for accuracy and completeness." ,
54
66
)
55
67
56
68
57
69
@pytest .fixture
58
- def image_analysis_agent ():
70
+ def image_analysis_agent (hook_provider ):
59
71
"""Create an agent specialized in image analysis."""
60
72
return Agent (
73
+ hooks = [hook_provider ],
61
74
system_prompt = (
62
75
"You are an image analysis expert. Describe what you see in images and provide detailed analysis."
63
- )
76
+ ),
64
77
)
65
78
66
79
@@ -149,7 +162,7 @@ def proceed_to_second_summary(state):
149
162
150
163
151
164
@pytest .mark .asyncio
152
- async def test_graph_execution_with_image (image_analysis_agent , summary_agent , yellow_img ):
165
+ async def test_graph_execution_with_image (image_analysis_agent , summary_agent , yellow_img , hook_provider ):
153
166
"""Test graph execution with multi-modal image input."""
154
167
builder = GraphBuilder ()
155
168
@@ -186,3 +199,16 @@ async def test_graph_execution_with_image(image_analysis_agent, summary_agent, y
186
199
# Verify both nodes completed
187
200
assert "image_analyzer" in result .results
188
201
assert "summarizer" in result .results
202
+
203
+ expected_hook_events = [
204
+ AgentInitializedEvent ,
205
+ BeforeInvocationEvent ,
206
+ MessageAddedEvent ,
207
+ BeforeModelInvocationEvent ,
208
+ AfterModelInvocationEvent ,
209
+ MessageAddedEvent ,
210
+ AfterInvocationEvent ,
211
+ ]
212
+
213
+ assert hook_provider .extract_for (image_analysis_agent ).event_types_received == expected_hook_events
214
+ assert hook_provider .extract_for (summary_agent ).event_types_received == expected_hook_events
0 commit comments