1
1
import math
2
2
import time
3
- from unittest .mock import MagicMock , Mock
3
+ from unittest .mock import MagicMock , Mock , patch
4
4
5
5
import pytest
6
6
@@ -93,6 +93,22 @@ def mock_swarm(mock_agents):
93
93
return swarm
94
94
95
95
96
+ @pytest .fixture
97
+ def mock_strands_tracer ():
98
+ with patch ("strands.multiagent.swarm.get_tracer" ) as mock_get_tracer :
99
+ mock_tracer_instance = MagicMock ()
100
+ mock_span = MagicMock ()
101
+ mock_tracer_instance .start_multiagent_span .return_value = mock_span
102
+ mock_get_tracer .return_value = mock_tracer_instance
103
+ yield mock_tracer_instance
104
+
105
+
106
+ @pytest .fixture
107
+ def mock_use_span ():
108
+ with patch ("strands.multiagent.swarm.trace_api.use_span" ) as mock_use_span :
109
+ yield mock_use_span
110
+
111
+
96
112
def test_swarm_structure_and_nodes (mock_swarm , mock_agents ):
97
113
"""Test swarm structure and SwarmNode properties."""
98
114
# Test swarm structure
@@ -214,7 +230,7 @@ def test_swarm_state_should_continue(mock_swarm):
214
230
215
231
216
232
@pytest .mark .asyncio
217
- async def test_swarm_execution_async (mock_swarm , mock_agents ):
233
+ async def test_swarm_execution_async (mock_strands_tracer , mock_use_span , mock_swarm , mock_agents ):
218
234
"""Test asynchronous swarm execution."""
219
235
# Execute swarm
220
236
task = [ContentBlock (text = "Analyze this task" ), ContentBlock (text = "Additional context" )]
@@ -237,8 +253,11 @@ async def test_swarm_execution_async(mock_swarm, mock_agents):
237
253
assert hasattr (result , "node_history" )
238
254
assert len (result .node_history ) == 1
239
255
256
+ mock_strands_tracer .start_multiagent_span .assert_called ()
257
+ mock_use_span .assert_called_once ()
240
258
241
- def test_swarm_synchronous_execution (mock_agents ):
259
+
260
+ def test_swarm_synchronous_execution (mock_strands_tracer , mock_use_span , mock_agents ):
242
261
"""Test synchronous swarm execution using __call__ method."""
243
262
agents = list (mock_agents .values ())
244
263
swarm = Swarm (
@@ -279,6 +298,9 @@ def test_swarm_synchronous_execution(mock_agents):
279
298
for node in swarm .nodes .values ():
280
299
node .executor .tool_registry .process_tools .assert_called ()
281
300
301
+ mock_strands_tracer .start_multiagent_span .assert_called ()
302
+ mock_use_span .assert_called_once ()
303
+
282
304
283
305
def test_swarm_builder_validation (mock_agents ):
284
306
"""Test swarm builder validation and error handling."""
@@ -405,7 +427,7 @@ def test_swarm_tool_creation_and_execution():
405
427
assert completion_result ["status" ] == "success"
406
428
407
429
408
- def test_swarm_failure_handling ():
430
+ def test_swarm_failure_handling (mock_strands_tracer , mock_use_span ):
409
431
"""Test swarm execution with agent failures."""
410
432
# Test execution with agent failures
411
433
failing_agent = create_mock_agent ("failing_agent" )
@@ -416,6 +438,8 @@ def test_swarm_failure_handling():
416
438
# The swarm catches exceptions internally and sets status to FAILED
417
439
result = failing_swarm ("Test failure handling" )
418
440
assert result .status == Status .FAILED
441
+ mock_strands_tracer .start_multiagent_span .assert_called ()
442
+ mock_use_span .assert_called_once ()
419
443
420
444
421
445
def test_swarm_metrics_handling ():
0 commit comments