mirror of
https://github.com/microsoft/agent-framework.git
synced 2026-06-16 21:04:09 +08:00
Laveesh Rohra
e3eff65a6b
* Add worker and clients * Clean code and refactor common code * Implement sample * Add sample * Update readmes * Fix tests * Fix tests * Update requirements * Fix typo * Address comments * use response.text
321 lines
13 KiB
Python
321 lines
13 KiB
Python
# Copyright (c) Microsoft. All rights reserved.
|
|
|
|
"""Unit tests for DurableAgentExecutor implementations.
|
|
|
|
Focuses on critical behavioral flows for executor strategies.
|
|
Run with: pytest tests/test_executors.py -v
|
|
"""
|
|
|
|
import time
|
|
from typing import Any
|
|
from unittest.mock import Mock
|
|
|
|
import pytest
|
|
from agent_framework import AgentRunResponse, Role
|
|
from durabletask.entities import EntityInstanceId
|
|
from durabletask.task import Task
|
|
from pydantic import BaseModel
|
|
|
|
from agent_framework_durabletask import DurableAgentThread
|
|
from agent_framework_durabletask._constants import DEFAULT_MAX_POLL_RETRIES, DEFAULT_POLL_INTERVAL_SECONDS
|
|
from agent_framework_durabletask._executors import (
|
|
ClientAgentExecutor,
|
|
DurableAgentTask,
|
|
OrchestrationAgentExecutor,
|
|
)
|
|
from agent_framework_durabletask._models import AgentSessionId, RunRequest
|
|
|
|
|
|
# Fixtures
|
|
@pytest.fixture
|
|
def mock_client() -> Mock:
|
|
"""Provide a mock client for ClientAgentExecutor tests."""
|
|
client = Mock()
|
|
client.signal_entity = Mock()
|
|
client.get_entity = Mock(return_value=None)
|
|
return client
|
|
|
|
|
|
@pytest.fixture
|
|
def mock_entity_task() -> Mock:
|
|
"""Provide a mock entity task."""
|
|
return Mock(spec=Task)
|
|
|
|
|
|
@pytest.fixture
|
|
def mock_orchestration_context(mock_entity_task: Mock) -> Mock:
|
|
"""Provide a mock orchestration context with call_entity configured."""
|
|
context = Mock()
|
|
context.call_entity = Mock(return_value=mock_entity_task)
|
|
return context
|
|
|
|
|
|
@pytest.fixture
|
|
def sample_run_request() -> RunRequest:
|
|
"""Provide a sample RunRequest for tests."""
|
|
return RunRequest(message="test message", correlation_id="test-123")
|
|
|
|
|
|
@pytest.fixture
|
|
def client_executor(mock_client: Mock) -> ClientAgentExecutor:
|
|
"""Provide a ClientAgentExecutor with minimal polling for fast tests."""
|
|
return ClientAgentExecutor(mock_client, max_poll_retries=1, poll_interval_seconds=0.01)
|
|
|
|
|
|
@pytest.fixture
|
|
def orchestration_executor(mock_orchestration_context: Mock) -> OrchestrationAgentExecutor:
|
|
"""Provide an OrchestrationAgentExecutor."""
|
|
return OrchestrationAgentExecutor(mock_orchestration_context)
|
|
|
|
|
|
@pytest.fixture
|
|
def successful_agent_response() -> dict[str, Any]:
|
|
"""Provide a successful agent response dictionary."""
|
|
return {
|
|
"messages": [{"role": "assistant", "contents": [{"type": "text", "text": "Hello!"}]}],
|
|
"created_at": "2025-12-30T10:00:00Z",
|
|
}
|
|
|
|
|
|
class TestExecutorThreadCreation:
|
|
"""Test that executors properly create DurableAgentThread with parameters."""
|
|
|
|
def test_client_executor_creates_durable_thread(self, mock_client: Mock) -> None:
|
|
"""Verify ClientAgentExecutor creates DurableAgentThread instances."""
|
|
executor = ClientAgentExecutor(mock_client)
|
|
|
|
thread = executor.get_new_thread("test_agent")
|
|
|
|
assert isinstance(thread, DurableAgentThread)
|
|
|
|
def test_client_executor_forwards_kwargs_to_thread(self, mock_client: Mock) -> None:
|
|
"""Verify ClientAgentExecutor forwards kwargs to DurableAgentThread creation."""
|
|
executor = ClientAgentExecutor(mock_client)
|
|
|
|
thread = executor.get_new_thread("test_agent", service_thread_id="client-123")
|
|
|
|
assert isinstance(thread, DurableAgentThread)
|
|
assert thread.service_thread_id == "client-123"
|
|
|
|
def test_orchestration_executor_creates_durable_thread(
|
|
self, orchestration_executor: OrchestrationAgentExecutor
|
|
) -> None:
|
|
"""Verify OrchestrationAgentExecutor creates DurableAgentThread instances."""
|
|
thread = orchestration_executor.get_new_thread("test_agent")
|
|
|
|
assert isinstance(thread, DurableAgentThread)
|
|
|
|
def test_orchestration_executor_forwards_kwargs_to_thread(
|
|
self, orchestration_executor: OrchestrationAgentExecutor
|
|
) -> None:
|
|
"""Verify OrchestrationAgentExecutor forwards kwargs to DurableAgentThread creation."""
|
|
thread = orchestration_executor.get_new_thread("test_agent", service_thread_id="orch-456")
|
|
|
|
assert isinstance(thread, DurableAgentThread)
|
|
assert thread.service_thread_id == "orch-456"
|
|
|
|
|
|
class TestClientAgentExecutorRun:
|
|
"""Test that ClientAgentExecutor.run_durable_agent works as implemented."""
|
|
|
|
def test_client_executor_run_returns_response(
|
|
self, client_executor: ClientAgentExecutor, sample_run_request: RunRequest
|
|
) -> None:
|
|
"""Verify ClientAgentExecutor.run_durable_agent returns AgentRunResponse (synchronous)."""
|
|
result = client_executor.run_durable_agent("test_agent", sample_run_request)
|
|
|
|
# Verify it returns an AgentRunResponse (synchronous, not a coroutine)
|
|
assert isinstance(result, AgentRunResponse)
|
|
assert result is not None
|
|
|
|
|
|
class TestClientAgentExecutorPollingConfiguration:
|
|
"""Test polling configuration parameters for ClientAgentExecutor."""
|
|
|
|
def test_executor_uses_default_polling_parameters(self, mock_client: Mock) -> None:
|
|
"""Verify executor initializes with default polling parameters."""
|
|
executor = ClientAgentExecutor(mock_client)
|
|
|
|
assert executor.max_poll_retries == DEFAULT_MAX_POLL_RETRIES
|
|
assert executor.poll_interval_seconds == DEFAULT_POLL_INTERVAL_SECONDS
|
|
|
|
def test_executor_accepts_custom_polling_parameters(self, mock_client: Mock) -> None:
|
|
"""Verify executor accepts and stores custom polling parameters."""
|
|
executor = ClientAgentExecutor(mock_client, max_poll_retries=20, poll_interval_seconds=0.5)
|
|
|
|
assert executor.max_poll_retries == 20
|
|
assert executor.poll_interval_seconds == 0.5
|
|
|
|
def test_executor_respects_custom_max_poll_retries(self, mock_client: Mock, sample_run_request: RunRequest) -> None:
|
|
"""Verify executor respects custom max_poll_retries during polling."""
|
|
# Create executor with only 2 retries
|
|
executor = ClientAgentExecutor(mock_client, max_poll_retries=2, poll_interval_seconds=0.01)
|
|
|
|
# Run the agent
|
|
result = executor.run_durable_agent("test_agent", sample_run_request)
|
|
|
|
# Verify it returns AgentRunResponse (should timeout after 2 attempts)
|
|
assert isinstance(result, AgentRunResponse)
|
|
|
|
# Verify get_entity was called 2 times (max_poll_retries)
|
|
assert mock_client.get_entity.call_count == 2
|
|
|
|
def test_executor_respects_custom_poll_interval(self, mock_client: Mock, sample_run_request: RunRequest) -> None:
|
|
"""Verify executor respects custom poll_interval_seconds during polling."""
|
|
# Create executor with very short interval
|
|
executor = ClientAgentExecutor(mock_client, max_poll_retries=3, poll_interval_seconds=0.01)
|
|
|
|
# Measure time taken
|
|
start = time.time()
|
|
result = executor.run_durable_agent("test_agent", sample_run_request)
|
|
elapsed = time.time() - start
|
|
|
|
# Should take roughly 3 * 0.01 = 0.03 seconds (plus overhead)
|
|
# Be generous with timing to avoid flakiness
|
|
assert elapsed < 0.2 # Should be quick with 0.01 interval
|
|
assert isinstance(result, AgentRunResponse)
|
|
|
|
|
|
class TestOrchestrationAgentExecutorRun:
|
|
"""Test OrchestrationAgentExecutor.run_durable_agent implementation."""
|
|
|
|
def test_orchestration_executor_run_returns_durable_agent_task(
|
|
self, orchestration_executor: OrchestrationAgentExecutor, sample_run_request: RunRequest
|
|
) -> None:
|
|
"""Verify OrchestrationAgentExecutor.run_durable_agent returns DurableAgentTask."""
|
|
result = orchestration_executor.run_durable_agent("test_agent", sample_run_request)
|
|
|
|
assert isinstance(result, DurableAgentTask)
|
|
|
|
def test_orchestration_executor_calls_entity_with_correct_parameters(
|
|
self,
|
|
mock_orchestration_context: Mock,
|
|
orchestration_executor: OrchestrationAgentExecutor,
|
|
sample_run_request: RunRequest,
|
|
) -> None:
|
|
"""Verify call_entity is invoked with correct entity ID and request."""
|
|
orchestration_executor.run_durable_agent("test_agent", sample_run_request)
|
|
|
|
# Verify call_entity was called once
|
|
assert mock_orchestration_context.call_entity.call_count == 1
|
|
|
|
# Get the call arguments
|
|
call_args = mock_orchestration_context.call_entity.call_args
|
|
entity_id_arg = call_args[0][0]
|
|
operation_arg = call_args[0][1]
|
|
request_dict_arg = call_args[0][2]
|
|
|
|
# Verify entity ID
|
|
assert isinstance(entity_id_arg, EntityInstanceId)
|
|
assert entity_id_arg.entity == "dafx-test_agent"
|
|
|
|
# Verify operation name
|
|
assert operation_arg == "run"
|
|
|
|
# Verify request dict
|
|
assert request_dict_arg == sample_run_request.to_dict()
|
|
|
|
def test_orchestration_executor_uses_thread_session_id(
|
|
self,
|
|
mock_orchestration_context: Mock,
|
|
orchestration_executor: OrchestrationAgentExecutor,
|
|
sample_run_request: RunRequest,
|
|
) -> None:
|
|
"""Verify executor uses thread's session ID when provided."""
|
|
# Create thread with specific session ID
|
|
session_id = AgentSessionId(name="test_agent", key="specific-key-123")
|
|
thread = DurableAgentThread.from_session_id(session_id)
|
|
|
|
result = orchestration_executor.run_durable_agent("test_agent", sample_run_request, thread=thread)
|
|
|
|
# Verify call_entity was called with the specific key
|
|
call_args = mock_orchestration_context.call_entity.call_args
|
|
entity_id_arg = call_args[0][0]
|
|
|
|
assert entity_id_arg.key == "specific-key-123"
|
|
assert isinstance(result, DurableAgentTask)
|
|
|
|
|
|
class TestDurableAgentTask:
|
|
"""Test DurableAgentTask completion and response transformation."""
|
|
|
|
def test_durable_agent_task_transforms_successful_result(
|
|
self, mock_entity_task: Mock, successful_agent_response: dict[str, Any]
|
|
) -> None:
|
|
"""Verify DurableAgentTask converts successful entity result to AgentRunResponse."""
|
|
mock_entity_task.is_failed = False
|
|
mock_entity_task.get_result = Mock(return_value=successful_agent_response)
|
|
|
|
task = DurableAgentTask(entity_task=mock_entity_task, response_format=None, correlation_id="test-123")
|
|
|
|
# Simulate child task completion
|
|
task.on_child_completed(mock_entity_task)
|
|
|
|
assert task.is_complete
|
|
result = task.get_result()
|
|
assert isinstance(result, AgentRunResponse)
|
|
assert len(result.messages) == 1
|
|
assert result.messages[0].role == Role.ASSISTANT
|
|
|
|
def test_durable_agent_task_propagates_failure(self, mock_entity_task: Mock) -> None:
|
|
"""Verify DurableAgentTask propagates task failures."""
|
|
mock_entity_task.is_failed = True
|
|
mock_entity_task.get_exception = Mock(return_value=ValueError("Entity error"))
|
|
|
|
task = DurableAgentTask(entity_task=mock_entity_task, response_format=None, correlation_id="test-123")
|
|
|
|
# Simulate child task completion with failure
|
|
task.on_child_completed(mock_entity_task)
|
|
|
|
assert task.is_complete
|
|
assert task.is_failed
|
|
exception = task.get_exception()
|
|
assert isinstance(exception, ValueError)
|
|
assert str(exception) == "Entity error"
|
|
|
|
def test_durable_agent_task_validates_response_format(self, mock_entity_task: Mock) -> None:
|
|
"""Verify DurableAgentTask validates response format when provided."""
|
|
mock_entity_task.is_failed = False
|
|
mock_entity_task.get_result = Mock(
|
|
return_value={
|
|
"messages": [{"role": "assistant", "contents": [{"type": "text", "text": '{"answer": "42"}'}]}],
|
|
"created_at": "2025-12-30T10:00:00Z",
|
|
}
|
|
)
|
|
|
|
class TestResponse(BaseModel):
|
|
answer: str
|
|
|
|
task = DurableAgentTask(entity_task=mock_entity_task, response_format=TestResponse, correlation_id="test-123")
|
|
|
|
# Simulate child task completion
|
|
task.on_child_completed(mock_entity_task)
|
|
|
|
assert task.is_complete
|
|
result = task.get_result()
|
|
assert isinstance(result, AgentRunResponse)
|
|
|
|
def test_durable_agent_task_ignores_duplicate_completion(
|
|
self, mock_entity_task: Mock, successful_agent_response: dict[str, Any]
|
|
) -> None:
|
|
"""Verify DurableAgentTask ignores duplicate completion calls."""
|
|
mock_entity_task.is_failed = False
|
|
mock_entity_task.get_result = Mock(return_value=successful_agent_response)
|
|
|
|
task = DurableAgentTask(entity_task=mock_entity_task, response_format=None, correlation_id="test-123")
|
|
|
|
# Simulate child task completion twice
|
|
task.on_child_completed(mock_entity_task)
|
|
first_result = task.get_result()
|
|
|
|
task.on_child_completed(mock_entity_task)
|
|
second_result = task.get_result()
|
|
|
|
# Should be the same result, get_result should only be called once
|
|
assert first_result is second_result
|
|
assert mock_entity_task.get_result.call_count == 1
|
|
|
|
|
|
if __name__ == "__main__":
|
|
pytest.main([__file__, "-v", "--tb=short"])
|