mirror of
https://github.com/microsoft/agent-framework.git
synced 2026-06-16 21:04:09 +08:00
Python: Complete durableagent package (#3058)
* 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
This commit is contained in:
Laveesh Rohra
committed by
GitHub
Unverified
GitHub
parent
a5b36dc379
commit
e3eff65a6b
@@ -1,402 +0,0 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
"""Unit tests for data models (AgentSessionId, RunRequest, AgentResponse)."""
|
||||
|
||||
import azure.durable_functions as df
|
||||
import pytest
|
||||
from agent_framework import Role
|
||||
from agent_framework_durabletask import RunRequest
|
||||
from pydantic import BaseModel
|
||||
|
||||
from agent_framework_azurefunctions._models import AgentSessionId
|
||||
|
||||
|
||||
class ModuleStructuredResponse(BaseModel):
|
||||
value: int
|
||||
|
||||
|
||||
class TestAgentSessionId:
|
||||
"""Test suite for AgentSessionId."""
|
||||
|
||||
def test_init_creates_session_id(self) -> None:
|
||||
"""Test that AgentSessionId initializes correctly."""
|
||||
session_id = AgentSessionId(name="AgentEntity", key="test-key-123")
|
||||
|
||||
assert session_id.name == "AgentEntity"
|
||||
assert session_id.key == "test-key-123"
|
||||
|
||||
def test_with_random_key_generates_guid(self) -> None:
|
||||
"""Test that with_random_key generates a GUID."""
|
||||
session_id = AgentSessionId.with_random_key(name="AgentEntity")
|
||||
|
||||
assert session_id.name == "AgentEntity"
|
||||
assert len(session_id.key) == 32 # UUID hex is 32 chars
|
||||
# Verify it's a valid hex string
|
||||
int(session_id.key, 16)
|
||||
|
||||
def test_with_random_key_unique_keys(self) -> None:
|
||||
"""Test that with_random_key generates unique keys."""
|
||||
session_id1 = AgentSessionId.with_random_key(name="AgentEntity")
|
||||
session_id2 = AgentSessionId.with_random_key(name="AgentEntity")
|
||||
|
||||
assert session_id1.key != session_id2.key
|
||||
|
||||
def test_to_entity_id_conversion(self) -> None:
|
||||
"""Test conversion to EntityId."""
|
||||
session_id = AgentSessionId(name="AgentEntity", key="test-key")
|
||||
entity_id = session_id.to_entity_id()
|
||||
|
||||
assert isinstance(entity_id, df.EntityId)
|
||||
assert entity_id.name == "dafx-AgentEntity"
|
||||
assert entity_id.key == "test-key"
|
||||
|
||||
def test_from_entity_id_conversion(self) -> None:
|
||||
"""Test creation from EntityId."""
|
||||
entity_id = df.EntityId(name="dafx-AgentEntity", key="test-key")
|
||||
session_id = AgentSessionId.from_entity_id(entity_id)
|
||||
|
||||
assert isinstance(session_id, AgentSessionId)
|
||||
assert session_id.name == "AgentEntity"
|
||||
assert session_id.key == "test-key"
|
||||
|
||||
def test_round_trip_entity_id_conversion(self) -> None:
|
||||
"""Test round-trip conversion to and from EntityId."""
|
||||
original = AgentSessionId(name="AgentEntity", key="test-key")
|
||||
entity_id = original.to_entity_id()
|
||||
restored = AgentSessionId.from_entity_id(entity_id)
|
||||
|
||||
assert restored.name == original.name
|
||||
assert restored.key == original.key
|
||||
|
||||
def test_str_representation(self) -> None:
|
||||
"""Test string representation."""
|
||||
session_id = AgentSessionId(name="AgentEntity", key="test-key-123")
|
||||
str_repr = str(session_id)
|
||||
|
||||
assert str_repr == "@AgentEntity@test-key-123"
|
||||
|
||||
def test_repr_representation(self) -> None:
|
||||
"""Test repr representation."""
|
||||
session_id = AgentSessionId(name="AgentEntity", key="test-key")
|
||||
repr_str = repr(session_id)
|
||||
|
||||
assert "AgentSessionId" in repr_str
|
||||
assert "AgentEntity" in repr_str
|
||||
assert "test-key" in repr_str
|
||||
|
||||
def test_parse_valid_session_id(self) -> None:
|
||||
"""Test parsing valid session ID string."""
|
||||
session_id = AgentSessionId.parse("@AgentEntity@test-key-123")
|
||||
|
||||
assert session_id.name == "AgentEntity"
|
||||
assert session_id.key == "test-key-123"
|
||||
|
||||
def test_parse_invalid_format_no_prefix(self) -> None:
|
||||
"""Test parsing invalid format without @ prefix."""
|
||||
with pytest.raises(ValueError) as exc_info:
|
||||
AgentSessionId.parse("AgentEntity@test-key")
|
||||
|
||||
assert "Invalid agent session ID format" in str(exc_info.value)
|
||||
|
||||
def test_parse_invalid_format_single_part(self) -> None:
|
||||
"""Test parsing invalid format with single part."""
|
||||
with pytest.raises(ValueError) as exc_info:
|
||||
AgentSessionId.parse("@AgentEntity")
|
||||
|
||||
assert "Invalid agent session ID format" in str(exc_info.value)
|
||||
|
||||
def test_parse_with_multiple_at_signs_in_key(self) -> None:
|
||||
"""Test parsing with @ signs in the key."""
|
||||
session_id = AgentSessionId.parse("@AgentEntity@key-with@symbols")
|
||||
|
||||
assert session_id.name == "AgentEntity"
|
||||
assert session_id.key == "key-with@symbols"
|
||||
|
||||
def test_parse_round_trip(self) -> None:
|
||||
"""Test round-trip parse and string conversion."""
|
||||
original = AgentSessionId(name="AgentEntity", key="test-key")
|
||||
str_repr = str(original)
|
||||
parsed = AgentSessionId.parse(str_repr)
|
||||
|
||||
assert parsed.name == original.name
|
||||
assert parsed.key == original.key
|
||||
|
||||
def test_to_entity_name_adds_prefix(self) -> None:
|
||||
"""Test that to_entity_name adds the dafx- prefix."""
|
||||
entity_name = AgentSessionId.to_entity_name("TestAgent")
|
||||
assert entity_name == "dafx-TestAgent"
|
||||
|
||||
def test_from_entity_id_strips_prefix(self) -> None:
|
||||
"""Test that from_entity_id strips the dafx- prefix."""
|
||||
entity_id = df.EntityId(name="dafx-TestAgent", key="key123")
|
||||
session_id = AgentSessionId.from_entity_id(entity_id)
|
||||
|
||||
assert session_id.name == "TestAgent"
|
||||
assert session_id.key == "key123"
|
||||
|
||||
def test_from_entity_id_raises_without_prefix(self) -> None:
|
||||
"""Test that from_entity_id raises ValueError when entity name lacks the prefix."""
|
||||
entity_id = df.EntityId(name="TestAgent", key="key123")
|
||||
|
||||
with pytest.raises(ValueError) as exc_info:
|
||||
AgentSessionId.from_entity_id(entity_id)
|
||||
|
||||
assert "not a valid agent session ID" in str(exc_info.value)
|
||||
assert "dafx-" in str(exc_info.value)
|
||||
|
||||
|
||||
class TestRunRequest:
|
||||
"""Test suite for RunRequest."""
|
||||
|
||||
def test_init_with_defaults(self) -> None:
|
||||
"""Test RunRequest initialization with defaults."""
|
||||
request = RunRequest(message="Hello")
|
||||
|
||||
assert request.message == "Hello"
|
||||
assert request.role == Role.USER
|
||||
assert request.response_format is None
|
||||
assert request.enable_tool_calls is True
|
||||
|
||||
def test_init_with_all_fields(self) -> None:
|
||||
"""Test RunRequest initialization with all fields."""
|
||||
schema = ModuleStructuredResponse
|
||||
request = RunRequest(
|
||||
message="Hello",
|
||||
role=Role.SYSTEM,
|
||||
response_format=schema,
|
||||
enable_tool_calls=False,
|
||||
)
|
||||
|
||||
assert request.message == "Hello"
|
||||
assert request.role == Role.SYSTEM
|
||||
assert request.response_format is schema
|
||||
assert request.enable_tool_calls is False
|
||||
|
||||
def test_init_coerces_string_role(self) -> None:
|
||||
"""Ensure string role values are coerced into Role instances."""
|
||||
request = RunRequest(message="Hello", role="system") # type: ignore[arg-type]
|
||||
|
||||
assert request.role == Role.SYSTEM
|
||||
|
||||
def test_to_dict_with_defaults(self) -> None:
|
||||
"""Test to_dict with default values."""
|
||||
request = RunRequest(message="Test message")
|
||||
data = request.to_dict()
|
||||
|
||||
assert data["message"] == "Test message"
|
||||
assert data["enable_tool_calls"] is True
|
||||
assert data["role"] == "user"
|
||||
assert "response_format" not in data or data["response_format"] is None
|
||||
assert "thread_id" not in data
|
||||
|
||||
def test_to_dict_with_all_fields(self) -> None:
|
||||
"""Test to_dict with all fields."""
|
||||
schema = ModuleStructuredResponse
|
||||
request = RunRequest(
|
||||
message="Hello",
|
||||
role=Role.ASSISTANT,
|
||||
response_format=schema,
|
||||
enable_tool_calls=False,
|
||||
)
|
||||
data = request.to_dict()
|
||||
|
||||
assert data["message"] == "Hello"
|
||||
assert data["role"] == "assistant"
|
||||
assert data["response_format"]["__response_schema_type__"] == "pydantic_model"
|
||||
assert data["response_format"]["module"] == schema.__module__
|
||||
assert data["response_format"]["qualname"] == schema.__qualname__
|
||||
assert data["enable_tool_calls"] is False
|
||||
assert "thread_id" not in data
|
||||
|
||||
def test_from_dict_with_defaults(self) -> None:
|
||||
"""Test from_dict with minimal data."""
|
||||
data = {"message": "Hello"}
|
||||
request = RunRequest.from_dict(data)
|
||||
|
||||
assert request.message == "Hello"
|
||||
assert request.role == Role.USER
|
||||
assert request.enable_tool_calls is True
|
||||
|
||||
def test_from_dict_ignores_thread_id_field(self) -> None:
|
||||
"""Ensure legacy thread_id input does not break RunRequest parsing."""
|
||||
request = RunRequest.from_dict({"message": "Hello", "thread_id": "ignored"})
|
||||
|
||||
assert request.message == "Hello"
|
||||
|
||||
def test_from_dict_with_all_fields(self) -> None:
|
||||
"""Test from_dict with all fields."""
|
||||
data = {
|
||||
"message": "Test",
|
||||
"role": "system",
|
||||
"response_format": {
|
||||
"__response_schema_type__": "pydantic_model",
|
||||
"module": ModuleStructuredResponse.__module__,
|
||||
"qualname": ModuleStructuredResponse.__qualname__,
|
||||
},
|
||||
"enable_tool_calls": False,
|
||||
}
|
||||
request = RunRequest.from_dict(data)
|
||||
|
||||
assert request.message == "Test"
|
||||
assert request.role == Role.SYSTEM
|
||||
assert request.response_format is ModuleStructuredResponse
|
||||
assert request.enable_tool_calls is False
|
||||
|
||||
def test_from_dict_with_unknown_role_preserves_value(self) -> None:
|
||||
"""Test from_dict keeps custom roles intact."""
|
||||
data = {"message": "Test", "role": "reviewer"}
|
||||
request = RunRequest.from_dict(data)
|
||||
|
||||
assert request.role.value == "reviewer"
|
||||
assert request.role != Role.USER
|
||||
|
||||
def test_from_dict_empty_message(self) -> None:
|
||||
"""Test from_dict with empty message."""
|
||||
request = RunRequest.from_dict({})
|
||||
|
||||
assert request.message == ""
|
||||
assert request.role == Role.USER
|
||||
|
||||
def test_round_trip_dict_conversion(self) -> None:
|
||||
"""Test round-trip to_dict and from_dict."""
|
||||
original = RunRequest(
|
||||
message="Test message",
|
||||
role=Role.SYSTEM,
|
||||
response_format=ModuleStructuredResponse,
|
||||
enable_tool_calls=False,
|
||||
)
|
||||
|
||||
data = original.to_dict()
|
||||
restored = RunRequest.from_dict(data)
|
||||
|
||||
assert restored.message == original.message
|
||||
assert restored.role == original.role
|
||||
assert restored.response_format is ModuleStructuredResponse
|
||||
assert restored.enable_tool_calls == original.enable_tool_calls
|
||||
|
||||
def test_round_trip_with_pydantic_response_format(self) -> None:
|
||||
"""Ensure Pydantic response formats serialize and deserialize properly."""
|
||||
original = RunRequest(
|
||||
message="Structured",
|
||||
response_format=ModuleStructuredResponse,
|
||||
)
|
||||
|
||||
data = original.to_dict()
|
||||
|
||||
assert data["response_format"]["__response_schema_type__"] == "pydantic_model"
|
||||
assert data["response_format"]["module"] == ModuleStructuredResponse.__module__
|
||||
assert data["response_format"]["qualname"] == ModuleStructuredResponse.__qualname__
|
||||
|
||||
restored = RunRequest.from_dict(data)
|
||||
assert restored.response_format is ModuleStructuredResponse
|
||||
|
||||
def test_init_with_correlationId(self) -> None:
|
||||
"""Test RunRequest initialization with correlationId."""
|
||||
request = RunRequest(message="Test message", correlation_id="corr-123")
|
||||
|
||||
assert request.message == "Test message"
|
||||
assert request.correlation_id == "corr-123"
|
||||
|
||||
def test_to_dict_with_correlationId(self) -> None:
|
||||
"""Test to_dict includes correlationId."""
|
||||
request = RunRequest(message="Test", correlation_id="corr-456")
|
||||
data = request.to_dict()
|
||||
|
||||
assert data["message"] == "Test"
|
||||
assert data["correlationId"] == "corr-456"
|
||||
|
||||
def test_from_dict_with_correlationId(self) -> None:
|
||||
"""Test from_dict with correlationId."""
|
||||
data = {"message": "Test", "correlationId": "corr-789"}
|
||||
request = RunRequest.from_dict(data)
|
||||
|
||||
assert request.message == "Test"
|
||||
assert request.correlation_id == "corr-789"
|
||||
|
||||
def test_round_trip_with_correlationId(self) -> None:
|
||||
"""Test round-trip to_dict and from_dict with correlationId."""
|
||||
original = RunRequest(
|
||||
message="Test message",
|
||||
role=Role.SYSTEM,
|
||||
correlation_id="corr-123",
|
||||
)
|
||||
|
||||
data = original.to_dict()
|
||||
restored = RunRequest.from_dict(data)
|
||||
|
||||
assert restored.message == original.message
|
||||
assert restored.role == original.role
|
||||
assert restored.correlation_id == original.correlation_id
|
||||
|
||||
def test_init_with_orchestration_id(self) -> None:
|
||||
"""Test RunRequest initialization with orchestration_id."""
|
||||
request = RunRequest(
|
||||
message="Test message",
|
||||
orchestration_id="orch-123",
|
||||
)
|
||||
|
||||
assert request.message == "Test message"
|
||||
assert request.orchestration_id == "orch-123"
|
||||
|
||||
def test_to_dict_with_orchestration_id(self) -> None:
|
||||
"""Test to_dict includes orchestrationId."""
|
||||
request = RunRequest(
|
||||
message="Test",
|
||||
orchestration_id="orch-456",
|
||||
)
|
||||
data = request.to_dict()
|
||||
|
||||
assert data["message"] == "Test"
|
||||
assert data["orchestrationId"] == "orch-456"
|
||||
|
||||
def test_to_dict_excludes_orchestration_id_when_none(self) -> None:
|
||||
"""Test to_dict excludes orchestrationId when not set."""
|
||||
request = RunRequest(
|
||||
message="Test",
|
||||
)
|
||||
data = request.to_dict()
|
||||
|
||||
assert "orchestrationId" not in data
|
||||
|
||||
def test_from_dict_with_orchestration_id(self) -> None:
|
||||
"""Test from_dict with orchestrationId."""
|
||||
data = {
|
||||
"message": "Test",
|
||||
"orchestrationId": "orch-789",
|
||||
}
|
||||
request = RunRequest.from_dict(data)
|
||||
|
||||
assert request.message == "Test"
|
||||
assert request.orchestration_id == "orch-789"
|
||||
|
||||
def test_round_trip_with_orchestration_id(self) -> None:
|
||||
"""Test round-trip to_dict and from_dict with orchestration_id."""
|
||||
original = RunRequest(
|
||||
message="Test message",
|
||||
role=Role.SYSTEM,
|
||||
correlation_id="corr-123",
|
||||
orchestration_id="orch-123",
|
||||
)
|
||||
|
||||
data = original.to_dict()
|
||||
restored = RunRequest.from_dict(data)
|
||||
|
||||
assert restored.message == original.message
|
||||
assert restored.role == original.role
|
||||
assert restored.correlation_id == original.correlation_id
|
||||
assert restored.orchestration_id == original.orchestration_id
|
||||
|
||||
|
||||
class TestModelIntegration:
|
||||
"""Test suite for integration between models."""
|
||||
|
||||
def test_run_request_with_session_id_string(self) -> None:
|
||||
"""AgentSessionId string can still be used by callers, but is not stored on RunRequest."""
|
||||
session_id = AgentSessionId.with_random_key("AgentEntity")
|
||||
session_id_str = str(session_id)
|
||||
|
||||
assert session_id_str.startswith("@AgentEntity@")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
pytest.main([__file__, "-v", "--tb=short"])
|
||||
@@ -6,11 +6,11 @@ from typing import Any
|
||||
from unittest.mock import Mock
|
||||
|
||||
import pytest
|
||||
from agent_framework import AgentRunResponse, AgentThread, ChatMessage
|
||||
from agent_framework import AgentRunResponse, ChatMessage
|
||||
from agent_framework_durabletask import DurableAIAgent
|
||||
from azure.durable_functions.models.Task import TaskBase, TaskState
|
||||
|
||||
from agent_framework_azurefunctions import AgentFunctionApp, DurableAIAgent
|
||||
from agent_framework_azurefunctions._models import AgentSessionId, DurableAgentThread
|
||||
from agent_framework_azurefunctions import AgentFunctionApp
|
||||
from agent_framework_azurefunctions._orchestration import AgentTask
|
||||
|
||||
|
||||
@@ -38,46 +38,96 @@ def _create_entity_task(task_id: int = 1) -> TaskBase:
|
||||
return _FakeTask(task_id)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_context():
|
||||
"""Create a mock orchestration context with UUID support."""
|
||||
context = Mock()
|
||||
context.instance_id = "test-instance"
|
||||
context.current_utc_datetime = Mock()
|
||||
return context
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_context_with_uuid() -> tuple[Mock, str]:
|
||||
"""Create a mock context with a single UUID."""
|
||||
from uuid import UUID
|
||||
|
||||
context = Mock()
|
||||
context.instance_id = "test-instance"
|
||||
context.current_utc_datetime = Mock()
|
||||
test_uuid = UUID("aaaaaaaa-aaaa-aaaa-aaaa-aaaaaaaaaaaa")
|
||||
context.new_uuid = Mock(return_value=test_uuid)
|
||||
return context, test_uuid.hex
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_context_with_multiple_uuids() -> tuple[Mock, list[str]]:
|
||||
"""Create a mock context with multiple UUIDs via side_effect."""
|
||||
from uuid import UUID
|
||||
|
||||
context = Mock()
|
||||
context.instance_id = "test-instance"
|
||||
context.current_utc_datetime = Mock()
|
||||
uuids = [
|
||||
UUID("aaaaaaaa-aaaa-aaaa-aaaa-aaaaaaaaaaaa"),
|
||||
UUID("bbbbbbbb-bbbb-bbbb-bbbb-bbbbbbbbbbbb"),
|
||||
UUID("cccccccc-cccc-cccc-cccc-cccccccccccc"),
|
||||
]
|
||||
context.new_uuid = Mock(side_effect=uuids)
|
||||
# Return the hex versions for assertion checking
|
||||
hex_uuids = [uuid.hex for uuid in uuids]
|
||||
return context, hex_uuids
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def executor_with_uuid() -> tuple[Any, Mock, str]:
|
||||
"""Create an executor with a mocked generate_unique_id method."""
|
||||
from agent_framework_azurefunctions._orchestration import AzureFunctionsAgentExecutor
|
||||
|
||||
context = Mock()
|
||||
context.instance_id = "test-instance"
|
||||
context.current_utc_datetime = Mock()
|
||||
|
||||
executor = AzureFunctionsAgentExecutor(context)
|
||||
test_uuid_hex = "aaaaaaaa-aaaa-aaaa-aaaa-aaaaaaaaaaaa"
|
||||
executor.generate_unique_id = Mock(return_value=test_uuid_hex)
|
||||
|
||||
return executor, context, test_uuid_hex
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def executor_with_multiple_uuids() -> tuple[Any, Mock, list[str]]:
|
||||
"""Create an executor with multiple mocked UUIDs."""
|
||||
from agent_framework_azurefunctions._orchestration import AzureFunctionsAgentExecutor
|
||||
|
||||
context = Mock()
|
||||
context.instance_id = "test-instance"
|
||||
context.current_utc_datetime = Mock()
|
||||
|
||||
executor = AzureFunctionsAgentExecutor(context)
|
||||
uuid_hexes = [
|
||||
"aaaaaaaa-aaaa-aaaa-aaaa-aaaaaaaaaaaa",
|
||||
"bbbbbbbb-bbbb-bbbb-bbbb-bbbbbbbbbbbb",
|
||||
"cccccccc-cccc-cccc-cccc-cccccccccccc",
|
||||
"dddddddd-dddd-dddd-dddd-dddddddddddd",
|
||||
"eeeeeeee-eeee-eeee-eeee-eeeeeeeeeeee",
|
||||
]
|
||||
executor.generate_unique_id = Mock(side_effect=uuid_hexes)
|
||||
|
||||
return executor, context, uuid_hexes
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def executor_with_context(mock_context_with_uuid: tuple[Mock, str]) -> tuple[Any, Mock]:
|
||||
"""Create an executor with a mocked context."""
|
||||
from agent_framework_azurefunctions._orchestration import AzureFunctionsAgentExecutor
|
||||
|
||||
context, _ = mock_context_with_uuid
|
||||
return AzureFunctionsAgentExecutor(context), context
|
||||
|
||||
|
||||
class TestAgentResponseHelpers:
|
||||
"""Tests for helper utilities that prepare AgentRunResponse values."""
|
||||
|
||||
@staticmethod
|
||||
def _create_agent_task() -> AgentTask:
|
||||
entity_task = _create_entity_task()
|
||||
return AgentTask(entity_task, None, "correlation-id")
|
||||
|
||||
def test_load_agent_response_from_instance(self) -> None:
|
||||
task = self._create_agent_task()
|
||||
response = AgentRunResponse(messages=[ChatMessage(role="assistant", text='{"foo": "bar"}')])
|
||||
|
||||
loaded = task._load_agent_response(response)
|
||||
|
||||
assert loaded is response
|
||||
assert loaded.value is None
|
||||
|
||||
def test_load_agent_response_from_serialized(self) -> None:
|
||||
task = self._create_agent_task()
|
||||
serialized = AgentRunResponse(messages=[ChatMessage(role="assistant", text="structured")]).to_dict()
|
||||
serialized["value"] = {"answer": 42}
|
||||
|
||||
loaded = task._load_agent_response(serialized)
|
||||
|
||||
assert loaded is not None
|
||||
assert loaded.value == {"answer": 42}
|
||||
loaded_dict = loaded.to_dict()
|
||||
assert loaded_dict["type"] == "agent_run_response"
|
||||
|
||||
def test_load_agent_response_rejects_none(self) -> None:
|
||||
task = self._create_agent_task()
|
||||
|
||||
with pytest.raises(ValueError):
|
||||
task._load_agent_response(None)
|
||||
|
||||
def test_load_agent_response_rejects_unsupported_type(self) -> None:
|
||||
task = self._create_agent_task()
|
||||
|
||||
with pytest.raises(TypeError, match="Unsupported type"):
|
||||
task._load_agent_response(["invalid", "list"]) # type: ignore[arg-type]
|
||||
"""Tests for response handling through public AgentTask API."""
|
||||
|
||||
def test_try_set_value_success(self) -> None:
|
||||
"""Test try_set_value correctly processes successful task completion."""
|
||||
@@ -144,335 +194,10 @@ class TestAgentResponseHelpers:
|
||||
assert isinstance(task.result.value, TestSchema)
|
||||
assert task.result.value.answer == "42"
|
||||
|
||||
def test_ensure_response_format_parses_value(self) -> None:
|
||||
"""Test _ensure_response_format correctly parses response value."""
|
||||
from pydantic import BaseModel
|
||||
|
||||
class SampleSchema(BaseModel):
|
||||
name: str
|
||||
|
||||
task = self._create_agent_task()
|
||||
response = AgentRunResponse(messages=[ChatMessage(role="assistant", text='{"name": "test"}')])
|
||||
|
||||
# Value should be None initially
|
||||
assert response.value is None
|
||||
|
||||
# Parse the value
|
||||
task._ensure_response_format(SampleSchema, "test-correlation", response)
|
||||
|
||||
# Value should now be parsed
|
||||
assert isinstance(response.value, SampleSchema)
|
||||
assert response.value.name == "test"
|
||||
|
||||
def test_ensure_response_format_skips_if_already_parsed(self) -> None:
|
||||
"""Test _ensure_response_format does not re-parse if value already matches format."""
|
||||
from pydantic import BaseModel
|
||||
|
||||
class SampleSchema(BaseModel):
|
||||
name: str
|
||||
|
||||
task = self._create_agent_task()
|
||||
existing_value = SampleSchema(name="existing")
|
||||
response = AgentRunResponse(
|
||||
messages=[ChatMessage(role="assistant", text='{"name": "new"}')],
|
||||
value=existing_value,
|
||||
)
|
||||
|
||||
# Call _ensure_response_format
|
||||
task._ensure_response_format(SampleSchema, "test-correlation", response)
|
||||
|
||||
# Value should remain unchanged (not re-parsed)
|
||||
assert response.value is existing_value
|
||||
assert response.value.name == "existing"
|
||||
|
||||
|
||||
class TestDurableAIAgent:
|
||||
"""Test suite for DurableAIAgent wrapper."""
|
||||
|
||||
def test_init(self) -> None:
|
||||
"""Test DurableAIAgent initialization."""
|
||||
mock_context = Mock()
|
||||
mock_context.instance_id = "test-instance-123"
|
||||
|
||||
agent = DurableAIAgent(mock_context, "TestAgent")
|
||||
|
||||
assert agent.context == mock_context
|
||||
assert agent.agent_name == "TestAgent"
|
||||
|
||||
def test_implements_agent_protocol(self) -> None:
|
||||
"""Test that DurableAIAgent implements AgentProtocol."""
|
||||
from agent_framework import AgentProtocol
|
||||
|
||||
mock_context = Mock()
|
||||
agent = DurableAIAgent(mock_context, "TestAgent")
|
||||
|
||||
# Check that agent satisfies AgentProtocol
|
||||
assert isinstance(agent, AgentProtocol)
|
||||
|
||||
def test_has_agent_protocol_properties(self) -> None:
|
||||
"""Test that DurableAIAgent has AgentProtocol properties."""
|
||||
mock_context = Mock()
|
||||
agent = DurableAIAgent(mock_context, "TestAgent")
|
||||
|
||||
# AgentProtocol properties
|
||||
assert hasattr(agent, "id")
|
||||
assert hasattr(agent, "name")
|
||||
assert hasattr(agent, "description")
|
||||
assert hasattr(agent, "display_name")
|
||||
|
||||
# Verify values
|
||||
assert agent.name == "TestAgent"
|
||||
assert agent.description == "Durable agent proxy for TestAgent"
|
||||
assert agent.display_name == "TestAgent"
|
||||
assert agent.id is not None # Auto-generated UUID
|
||||
|
||||
def test_get_new_thread(self) -> None:
|
||||
"""Test creating a new agent thread."""
|
||||
mock_context = Mock()
|
||||
mock_context.instance_id = "test-instance-456"
|
||||
mock_context.new_uuid = Mock(return_value="test-guid-456")
|
||||
|
||||
agent = DurableAIAgent(mock_context, "WriterAgent")
|
||||
thread = agent.get_new_thread()
|
||||
|
||||
assert isinstance(thread, DurableAgentThread)
|
||||
assert thread.session_id is not None
|
||||
session_id = thread.session_id
|
||||
assert isinstance(session_id, AgentSessionId)
|
||||
assert session_id.name == "WriterAgent"
|
||||
assert session_id.key == "test-guid-456"
|
||||
mock_context.new_uuid.assert_called_once()
|
||||
|
||||
def test_get_new_thread_deterministic(self) -> None:
|
||||
"""Test that get_new_thread creates deterministic session IDs."""
|
||||
|
||||
mock_context = Mock()
|
||||
mock_context.instance_id = "test-instance-789"
|
||||
mock_context.new_uuid = Mock(side_effect=["session-guid-1", "session-guid-2"])
|
||||
|
||||
agent = DurableAIAgent(mock_context, "EditorAgent")
|
||||
|
||||
# Create multiple threads - they should have unique session IDs
|
||||
thread1 = agent.get_new_thread()
|
||||
thread2 = agent.get_new_thread()
|
||||
|
||||
assert isinstance(thread1, DurableAgentThread)
|
||||
assert isinstance(thread2, DurableAgentThread)
|
||||
|
||||
session_id1 = thread1.session_id
|
||||
session_id2 = thread2.session_id
|
||||
assert session_id1 is not None and session_id2 is not None
|
||||
assert isinstance(session_id1, AgentSessionId)
|
||||
assert isinstance(session_id2, AgentSessionId)
|
||||
assert session_id1.name == "EditorAgent"
|
||||
assert session_id2.name == "EditorAgent"
|
||||
assert session_id1.key == "session-guid-1"
|
||||
assert session_id2.key == "session-guid-2"
|
||||
assert mock_context.new_uuid.call_count == 2
|
||||
|
||||
def test_run_creates_entity_call(self) -> None:
|
||||
"""Test that run() creates proper entity call and returns a Task."""
|
||||
mock_context = Mock()
|
||||
mock_context.instance_id = "test-instance-001"
|
||||
mock_context.new_uuid = Mock(side_effect=["thread-guid", "correlation-guid"])
|
||||
|
||||
entity_task = _create_entity_task()
|
||||
mock_context.call_entity = Mock(return_value=entity_task)
|
||||
|
||||
agent = DurableAIAgent(mock_context, "TestAgent")
|
||||
|
||||
# Create thread
|
||||
thread = agent.get_new_thread()
|
||||
|
||||
# Call run() - returns AgentTask directly
|
||||
task = agent.run(messages="Test message", thread=thread, enable_tool_calls=True)
|
||||
|
||||
assert isinstance(task, AgentTask)
|
||||
assert task.children[0] == entity_task
|
||||
|
||||
# Verify call_entity was called with correct parameters
|
||||
assert mock_context.call_entity.called
|
||||
call_args = mock_context.call_entity.call_args
|
||||
entity_id, operation, request = call_args[0]
|
||||
|
||||
assert operation == "run"
|
||||
assert request["message"] == "Test message"
|
||||
assert request["enable_tool_calls"] is True
|
||||
assert "correlationId" in request
|
||||
assert request["correlationId"] == "correlation-guid"
|
||||
assert "thread_id" not in request
|
||||
# Verify orchestration ID is set from context.instance_id
|
||||
assert "orchestrationId" in request
|
||||
assert request["orchestrationId"] == "test-instance-001"
|
||||
|
||||
def test_run_sets_orchestration_id(self) -> None:
|
||||
"""Test that run() sets the orchestration_id from context.instance_id."""
|
||||
mock_context = Mock()
|
||||
mock_context.instance_id = "my-orchestration-123"
|
||||
mock_context.new_uuid = Mock(side_effect=["thread-guid", "correlation-guid"])
|
||||
|
||||
entity_task = _create_entity_task()
|
||||
mock_context.call_entity = Mock(return_value=entity_task)
|
||||
|
||||
agent = DurableAIAgent(mock_context, "TestAgent")
|
||||
thread = agent.get_new_thread()
|
||||
|
||||
agent.run(messages="Test", thread=thread)
|
||||
|
||||
call_args = mock_context.call_entity.call_args
|
||||
request = call_args[0][2]
|
||||
|
||||
assert request["orchestrationId"] == "my-orchestration-123"
|
||||
|
||||
def test_run_without_thread(self) -> None:
|
||||
"""Test that run() works without explicit thread (creates unique session key)."""
|
||||
mock_context = Mock()
|
||||
mock_context.instance_id = "test-instance-002"
|
||||
mock_context.new_uuid = Mock(side_effect=["auto-generated-guid", "correlation-guid"])
|
||||
|
||||
entity_task = _create_entity_task()
|
||||
mock_context.call_entity = Mock(return_value=entity_task)
|
||||
|
||||
agent = DurableAIAgent(mock_context, "TestAgent")
|
||||
|
||||
# Call without thread
|
||||
task = agent.run(messages="Test message")
|
||||
|
||||
assert isinstance(task, AgentTask)
|
||||
assert task.children[0] == entity_task
|
||||
|
||||
# Verify the entity ID uses the auto-generated GUID with dafx- prefix
|
||||
call_args = mock_context.call_entity.call_args
|
||||
entity_id = call_args[0][0]
|
||||
assert entity_id.name == "dafx-TestAgent"
|
||||
assert entity_id.key == "auto-generated-guid"
|
||||
# Should be called twice: once for session_key, once for correlationId
|
||||
assert mock_context.new_uuid.call_count == 2
|
||||
|
||||
def test_run_with_response_format(self) -> None:
|
||||
"""Test that run() passes response format correctly."""
|
||||
mock_context = Mock()
|
||||
mock_context.instance_id = "test-instance-003"
|
||||
|
||||
entity_task = _create_entity_task()
|
||||
mock_context.call_entity = Mock(return_value=entity_task)
|
||||
|
||||
agent = DurableAIAgent(mock_context, "TestAgent")
|
||||
|
||||
from pydantic import BaseModel
|
||||
|
||||
class SampleSchema(BaseModel):
|
||||
key: str
|
||||
|
||||
# Create thread and call
|
||||
thread = agent.get_new_thread()
|
||||
|
||||
task = agent.run(messages="Test message", thread=thread, response_format=SampleSchema)
|
||||
|
||||
assert isinstance(task, AgentTask)
|
||||
assert task.children[0] == entity_task
|
||||
|
||||
# Verify schema was passed in the call_entity arguments
|
||||
call_args = mock_context.call_entity.call_args
|
||||
input_data = call_args[0][2] # Third argument is input_data
|
||||
assert "response_format" in input_data
|
||||
assert input_data["response_format"]["__response_schema_type__"] == "pydantic_model"
|
||||
assert input_data["response_format"]["module"] == SampleSchema.__module__
|
||||
assert input_data["response_format"]["qualname"] == SampleSchema.__qualname__
|
||||
|
||||
def test_messages_to_string(self) -> None:
|
||||
"""Test converting ChatMessage list to string."""
|
||||
from agent_framework import ChatMessage
|
||||
|
||||
mock_context = Mock()
|
||||
agent = DurableAIAgent(mock_context, "TestAgent")
|
||||
|
||||
messages = [
|
||||
ChatMessage(role="user", text="Hello"),
|
||||
ChatMessage(role="assistant", text="Hi there"),
|
||||
ChatMessage(role="user", text="How are you?"),
|
||||
]
|
||||
|
||||
result = agent._messages_to_string(messages)
|
||||
|
||||
assert result == "Hello\nHi there\nHow are you?"
|
||||
|
||||
def test_run_with_chat_message(self) -> None:
|
||||
"""Test that run() handles ChatMessage input."""
|
||||
from agent_framework import ChatMessage
|
||||
|
||||
mock_context = Mock()
|
||||
mock_context.new_uuid = Mock(side_effect=["thread-guid", "correlation-guid"])
|
||||
entity_task = _create_entity_task()
|
||||
mock_context.call_entity = Mock(return_value=entity_task)
|
||||
|
||||
agent = DurableAIAgent(mock_context, "TestAgent")
|
||||
thread = agent.get_new_thread()
|
||||
|
||||
# Call with ChatMessage
|
||||
msg = ChatMessage(role="user", text="Hello")
|
||||
task = agent.run(messages=msg, thread=thread)
|
||||
|
||||
assert isinstance(task, AgentTask)
|
||||
assert task.children[0] == entity_task
|
||||
|
||||
# Verify message was converted to string
|
||||
call_args = mock_context.call_entity.call_args
|
||||
request = call_args[0][2]
|
||||
assert request["message"] == "Hello"
|
||||
|
||||
def test_run_stream_raises_not_implemented(self) -> None:
|
||||
"""Test that run_stream() method raises NotImplementedError."""
|
||||
mock_context = Mock()
|
||||
agent = DurableAIAgent(mock_context, "TestAgent")
|
||||
|
||||
with pytest.raises(NotImplementedError) as exc_info:
|
||||
agent.run_stream("Test message")
|
||||
|
||||
error_msg = str(exc_info.value)
|
||||
assert "Streaming is not supported" in error_msg
|
||||
|
||||
def test_entity_id_format(self) -> None:
|
||||
"""Test that EntityId is created with correct format (name, key)."""
|
||||
from azure.durable_functions import EntityId
|
||||
|
||||
mock_context = Mock()
|
||||
mock_context.new_uuid = Mock(return_value="test-guid-789")
|
||||
mock_context.call_entity = Mock(return_value=_create_entity_task())
|
||||
|
||||
agent = DurableAIAgent(mock_context, "WriterAgent")
|
||||
thread = agent.get_new_thread()
|
||||
|
||||
# Call run() to trigger entity ID creation
|
||||
agent.run("Test", thread=thread)
|
||||
|
||||
# Verify call_entity was called with correct EntityId
|
||||
call_args = mock_context.call_entity.call_args
|
||||
entity_id = call_args[0][0]
|
||||
|
||||
# EntityId should be EntityId(name="dafx-WriterAgent", key="test-guid-789")
|
||||
# Which formats as "@dafx-writeragent@test-guid-789"
|
||||
assert isinstance(entity_id, EntityId)
|
||||
assert entity_id.name == "dafx-WriterAgent"
|
||||
assert entity_id.key == "test-guid-789"
|
||||
assert str(entity_id) == "@dafx-writeragent@test-guid-789"
|
||||
|
||||
|
||||
class TestAgentFunctionAppGetAgent:
|
||||
"""Test suite for AgentFunctionApp.get_agent."""
|
||||
|
||||
def test_get_agent_method(self) -> None:
|
||||
"""Test get_agent method creates DurableAIAgent for registered agent."""
|
||||
app = _app_with_registered_agents("MyAgent")
|
||||
mock_context = Mock()
|
||||
mock_context.instance_id = "test-instance-100"
|
||||
|
||||
agent = app.get_agent(mock_context, "MyAgent")
|
||||
|
||||
assert isinstance(agent, DurableAIAgent)
|
||||
assert agent.agent_name == "MyAgent"
|
||||
assert agent.context == mock_context
|
||||
|
||||
def test_get_agent_raises_for_unregistered_agent(self) -> None:
|
||||
"""Test get_agent raises ValueError when agent is not registered."""
|
||||
app = _app_with_registered_agents("KnownAgent")
|
||||
@@ -484,15 +209,9 @@ class TestAgentFunctionAppGetAgent:
|
||||
class TestOrchestrationIntegration:
|
||||
"""Integration tests for orchestration scenarios."""
|
||||
|
||||
def test_sequential_agent_calls_simulation(self) -> None:
|
||||
def test_sequential_agent_calls_simulation(self, executor_with_multiple_uuids: tuple[Any, Mock, list[str]]) -> None:
|
||||
"""Simulate sequential agent calls in an orchestration."""
|
||||
mock_context = Mock()
|
||||
mock_context.instance_id = "test-orchestration-001"
|
||||
# new_uuid will be called 3 times:
|
||||
# 1. thread creation
|
||||
# 2. correlationId for first call
|
||||
# 3. correlationId for second call
|
||||
mock_context.new_uuid = Mock(side_effect=["deterministic-guid-001", "corr-1", "corr-2"])
|
||||
executor, context, uuid_hexes = executor_with_multiple_uuids
|
||||
|
||||
# Track entity calls
|
||||
entity_calls: list[dict[str, Any]] = []
|
||||
@@ -501,10 +220,10 @@ class TestOrchestrationIntegration:
|
||||
entity_calls.append({"entity_id": str(entity_id), "operation": operation, "input": input_data})
|
||||
return _create_entity_task()
|
||||
|
||||
mock_context.call_entity = Mock(side_effect=mock_call_entity_side_effect)
|
||||
context.call_entity = Mock(side_effect=mock_call_entity_side_effect)
|
||||
|
||||
app = _app_with_registered_agents("WriterAgent")
|
||||
agent = app.get_agent(mock_context, "WriterAgent")
|
||||
# Create agent directly with executor (not via app.get_agent)
|
||||
agent = DurableAIAgent(executor, "WriterAgent")
|
||||
|
||||
# Create thread
|
||||
thread = agent.get_new_thread()
|
||||
@@ -520,18 +239,15 @@ class TestOrchestrationIntegration:
|
||||
# Verify both calls used the same entity (same session key)
|
||||
assert len(entity_calls) == 2
|
||||
assert entity_calls[0]["entity_id"] == entity_calls[1]["entity_id"]
|
||||
# EntityId format is @dafx-writeragent@deterministic-guid-001
|
||||
assert entity_calls[0]["entity_id"] == "@dafx-writeragent@deterministic-guid-001"
|
||||
# new_uuid called 3 times: thread + 2 correlation IDs
|
||||
assert mock_context.new_uuid.call_count == 3
|
||||
# EntityId format is @dafx-writeragent@<uuid_hex>
|
||||
expected_entity_id = f"@dafx-writeragent@{uuid_hexes[0]}"
|
||||
assert entity_calls[0]["entity_id"] == expected_entity_id
|
||||
# generate_unique_id called 3 times: thread + 2 correlation IDs
|
||||
assert executor.generate_unique_id.call_count == 3
|
||||
|
||||
def test_multiple_agents_in_orchestration(self) -> None:
|
||||
def test_multiple_agents_in_orchestration(self, executor_with_multiple_uuids: tuple[Any, Mock, list[str]]) -> None:
|
||||
"""Test using multiple different agents in one orchestration."""
|
||||
mock_context = Mock()
|
||||
mock_context.instance_id = "test-orchestration-002"
|
||||
# Mock new_uuid to return different GUIDs for each call
|
||||
# Order: writer thread, editor thread, writer correlation, editor correlation
|
||||
mock_context.new_uuid = Mock(side_effect=["writer-guid-001", "editor-guid-002", "writer-corr", "editor-corr"])
|
||||
executor, context, uuid_hexes = executor_with_multiple_uuids
|
||||
|
||||
entity_calls: list[str] = []
|
||||
|
||||
@@ -539,11 +255,11 @@ class TestOrchestrationIntegration:
|
||||
entity_calls.append(str(entity_id))
|
||||
return _create_entity_task()
|
||||
|
||||
mock_context.call_entity = Mock(side_effect=mock_call_entity_side_effect)
|
||||
context.call_entity = Mock(side_effect=mock_call_entity_side_effect)
|
||||
|
||||
app = _app_with_registered_agents("WriterAgent", "EditorAgent")
|
||||
writer = app.get_agent(mock_context, "WriterAgent")
|
||||
editor = app.get_agent(mock_context, "EditorAgent")
|
||||
# Create agents directly with executor (not via app.get_agent)
|
||||
writer = DurableAIAgent(executor, "WriterAgent")
|
||||
editor = DurableAIAgent(executor, "EditorAgent")
|
||||
|
||||
writer_thread = writer.get_new_thread()
|
||||
editor_thread = editor.get_new_thread()
|
||||
@@ -557,62 +273,11 @@ class TestOrchestrationIntegration:
|
||||
|
||||
# Verify different entity IDs were used
|
||||
assert len(entity_calls) == 2
|
||||
# EntityId format is @dafx-agentname@guid (lowercased agent name with dafx- prefix)
|
||||
assert entity_calls[0] == "@dafx-writeragent@writer-guid-001"
|
||||
assert entity_calls[1] == "@dafx-editoragent@editor-guid-002"
|
||||
|
||||
|
||||
class TestAgentThreadSerialization:
|
||||
"""Test that AgentThread can be serialized for orchestration state."""
|
||||
|
||||
async def test_agent_thread_serialize(self) -> None:
|
||||
"""Test that AgentThread can be serialized."""
|
||||
thread = AgentThread()
|
||||
|
||||
# Serialize
|
||||
serialized = await thread.serialize()
|
||||
|
||||
assert isinstance(serialized, dict)
|
||||
assert "service_thread_id" in serialized
|
||||
|
||||
async def test_agent_thread_deserialize(self) -> None:
|
||||
"""Test that AgentThread can be deserialized."""
|
||||
thread = AgentThread()
|
||||
serialized = await thread.serialize()
|
||||
|
||||
# Deserialize
|
||||
restored = await AgentThread.deserialize(serialized)
|
||||
|
||||
assert isinstance(restored, AgentThread)
|
||||
assert restored.service_thread_id == thread.service_thread_id
|
||||
|
||||
async def test_durable_agent_thread_serialization(self) -> None:
|
||||
"""Test that DurableAgentThread persists session metadata during serialization."""
|
||||
mock_context = Mock()
|
||||
mock_context.instance_id = "test-instance-999"
|
||||
mock_context.new_uuid = Mock(return_value="test-guid-999")
|
||||
|
||||
agent = DurableAIAgent(mock_context, "TestAgent")
|
||||
thread = agent.get_new_thread()
|
||||
|
||||
assert isinstance(thread, DurableAgentThread)
|
||||
# Verify custom attribute and property exist
|
||||
assert thread.session_id is not None
|
||||
session_id = thread.session_id
|
||||
assert isinstance(session_id, AgentSessionId)
|
||||
assert session_id.name == "TestAgent"
|
||||
assert session_id.key == "test-guid-999"
|
||||
|
||||
# Standard serialization should still work
|
||||
serialized = await thread.serialize()
|
||||
assert isinstance(serialized, dict)
|
||||
assert serialized.get("durable_session_id") == str(session_id)
|
||||
|
||||
# After deserialization, we'd need to restore the custom attribute
|
||||
# This would be handled by the orchestration framework
|
||||
restored = await DurableAgentThread.deserialize(serialized)
|
||||
assert isinstance(restored, DurableAgentThread)
|
||||
assert restored.session_id == session_id
|
||||
# EntityId format is @dafx-agentname@uuid_hex (lowercased agent name with dafx- prefix)
|
||||
expected_writer_id = f"@dafx-writeragent@{uuid_hexes[0]}"
|
||||
expected_editor_id = f"@dafx-editoragent@{uuid_hexes[1]}"
|
||||
assert entity_calls[0] == expected_writer_id
|
||||
assert entity_calls[1] == expected_editor_id
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
|
||||
Reference in New Issue
Block a user