@@ -327,18 +327,22 @@ async def test_consume_all_continues_on_queue_empty_if_not_really_closed(
327327def test_agent_task_callback_sets_exception (event_consumer : EventConsumer ):
328328 """Test that agent_task_callback sets _exception if the task had one."""
329329 mock_task = MagicMock (spec = asyncio .Task )
330+ mock_task .cancelled .return_value = False
331+ mock_task .done .return_value = True
330332 sample_exception = ValueError ('Task failed' )
331333 mock_task .exception .return_value = sample_exception
332334
333335 event_consumer .agent_task_callback (mock_task )
334336
335337 assert event_consumer ._exception == sample_exception
336- # mock_task.exception.assert_called_once() # Removing this, as exception() might be called internally by the check
338+ mock_task .exception .assert_called_once ()
337339
338340
339341def test_agent_task_callback_no_exception (event_consumer : EventConsumer ):
340342 """Test that agent_task_callback does nothing if the task has no exception."""
341343 mock_task = MagicMock (spec = asyncio .Task )
344+ mock_task .cancelled .return_value = False
345+ mock_task .done .return_value = True
342346 mock_task .exception .return_value = None # No exception
343347
344348 event_consumer .agent_task_callback (mock_task )
@@ -347,6 +351,34 @@ def test_agent_task_callback_no_exception(event_consumer: EventConsumer):
347351 mock_task .exception .assert_called_once ()
348352
349353
354+ def test_agent_task_callback_cancelled_task (event_consumer : EventConsumer ):
355+ """Test that agent_task_callback does nothing if the task has no exception."""
356+ mock_task = MagicMock (spec = asyncio .Task )
357+ mock_task .cancelled .return_value = True
358+ mock_task .done .return_value = True
359+ sample_exception = ValueError ('Task still running' )
360+ mock_task .exception .return_value = sample_exception
361+
362+ event_consumer .agent_task_callback (mock_task )
363+
364+ assert event_consumer ._exception is None # Should remain None
365+ mock_task .exception .assert_not_called ()
366+
367+
368+ def test_agent_task_callback_not_done_task (event_consumer : EventConsumer ):
369+ """Test that agent_task_callback does nothing if the task has no exception."""
370+ mock_task = MagicMock (spec = asyncio .Task )
371+ mock_task .cancelled .return_value = False
372+ mock_task .done .return_value = False
373+ sample_exception = ValueError ('Task is cancelled' )
374+ mock_task .exception .return_value = sample_exception
375+
376+ event_consumer .agent_task_callback (mock_task )
377+
378+ assert event_consumer ._exception is None # Should remain None
379+ mock_task .exception .assert_not_called ()
380+
381+
350382@pytest .mark .asyncio
351383async def test_consume_all_handles_validation_error (
352384 event_consumer : EventConsumer , mock_event_queue : AsyncMock
0 commit comments