agent/agent-framework
1
0
Fork 0
mirror of https://github.com/microsoft/agent-framework.git synced 2026-06-16 21:04:09 +08:00
Code Issues Packages Projects Releases Wiki Activity
Files
aa88195dcde3dc858f356cf22b87fbe2acb90c5f
agent-framework/python/packages/purview/tests/test_chat_middleware.py
T
Eduard van Valkenburg 838a7fd61d Python: [BREAKING] Types API Review improvements (#3647) 
* Replace Role and FinishReason classes with NewType + Literal

- Remove EnumLike metaclass from _types.py
- Replace Role class with NewType('Role', str) + RoleLiteral
- Replace FinishReason class with NewType('FinishReason', str) + FinishReasonLiteral
- Update all usages across codebase to use string literals
- Remove .value access patterns (direct string comparison now works)
- Add backward compatibility for legacy dict serialization format
- Update tests to reflect new string-based types

Addresses #3591, #3615

* Simplify ChatResponse and AgentResponse type hints (#3592)

- Remove overloads from ChatResponse.__init__
- Remove text parameter from ChatResponse.__init__
- Remove | dict[str, Any] from finish_reason and usage_details params
- Remove **kwargs from AgentResponse.__init__
- Both now accept ChatMessage | Sequence[ChatMessage] | None for messages
- Update docstrings and examples to reflect changes
- Fix tests that were using removed kwargs
- Fix Role type hint usage in ag-ui utils

* Remove text parameter from ChatResponseUpdate and AgentResponseUpdate (#3597)

- Remove text parameter from ChatResponseUpdate.__init__
- Remove text parameter from AgentResponseUpdate.__init__
- Remove **kwargs from both update classes
- Simplify contents parameter type to Sequence[Content] | None
- Update all usages to use contents=[Content.from_text(...)] pattern
- Fix imports in test files
- Update docstrings and examples

* Rename from_chat_response_updates to from_updates (#3593)

- ChatResponse.from_chat_response_updates → ChatResponse.from_updates
- ChatResponse.from_chat_response_generator → ChatResponse.from_update_generator
- AgentResponse.from_agent_run_response_updates → AgentResponse.from_updates

* Remove try_parse_value method from ChatResponse and AgentResponse (#3595)

- Remove try_parse_value method from ChatResponse
- Remove try_parse_value method from AgentResponse
- Remove try_parse_value calls from from_updates and from_update_generator methods
- Update samples to use try/except with response.value instead
- Update tests to use response.value pattern
- Users should now use response.value with try/except for safe parsing

* Add agent_id to AgentResponse and clarify author_name documentation (#3596)

- Add agent_id parameter to AgentResponse class
- Document that author_name is on ChatMessage objects, not responses
- Update ChatResponse docstring with author_name note
- Update AgentResponse docstring with author_name note

* Simplify ChatMessage.__init__ signature (#3618)

- Make contents a positional argument accepting Sequence[Content | str]
- Auto-convert strings in contents to TextContent
- Remove overloads, keep text kwarg for backward compatibility with serialization
- Update _parse_content_list to handle string items
- Update all usages across codebase to use new format: ChatMessage("role", ["text"])

* Allow Content as input on run and get_response

- Update prepare_messages and normalize_messages to accept Content
- Update type signatures in _agents.py and _clients.py
- Add tests for Content input handling

* Fix ChatMessage usage across packages and samples

Update all remaining ChatMessage(role=..., text=...) to use new
ChatMessage('role', ['text']) signature.

* Fix Role string usage and response format parsing

- Fix redis provider: remove .value access on string literals
- Fix durabletask ensure_response_format: set _response_format before accessing .value

* Fix ollama .value and ai_model_id issues, handle None in content list

- Fix ollama _chat_client: remove .value on string literals
- Fix ollama _chat_client: rename ai_model_id to model_id
- Fix _parse_content_list: skip None values gracefully

* Fix A2AAgent type signature to include Content

* Fix Role/FinishReason NewType dict annotations and improve test coverage to 95%

* Fix mypy errors for Role/FinishReason NewType usage

* Fix Role.TOOL and Role.ASSISTANT usage in _orchestrator_helpers.py

* Fix Role NewType usage in durabletask _models.py
2026-02-04 10:13:23 +00:00

351 lines
15 KiB
Python
Raw Blame History

# Copyright (c) Microsoft. All rights reserved.
"""Tests for Purview chat middleware."""
from dataclasses import dataclass
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
from agent_framework import ChatContext, ChatMessage
from azure.core.credentials import AccessToken
from agent_framework_purview import PurviewChatPolicyMiddleware, PurviewSettings
@dataclass
class DummyChatClient:
name: str = "dummy"
class TestPurviewChatPolicyMiddleware:
@pytest.fixture
def mock_credential(self) -> AsyncMock:
credential = AsyncMock()
credential.get_token = AsyncMock(return_value=AccessToken("fake-token", 9999999999))
return credential
@pytest.fixture
def settings(self) -> PurviewSettings:
return PurviewSettings(app_name="Test App", tenant_id="test-tenant")
@pytest.fixture
def middleware(self, mock_credential: AsyncMock, settings: PurviewSettings) -> PurviewChatPolicyMiddleware:
return PurviewChatPolicyMiddleware(mock_credential, settings)
@pytest.fixture
def chat_context(self) -> ChatContext:
chat_client = DummyChatClient()
chat_options = MagicMock()
chat_options.model = "test-model"
return ChatContext(chat_client=chat_client, messages=[ChatMessage("user", ["Hello"])], options=chat_options)
async def test_initialization(self, middleware: PurviewChatPolicyMiddleware) -> None:
assert middleware._client is not None
assert middleware._processor is not None
async def test_allows_clean_prompt(
self, middleware: PurviewChatPolicyMiddleware, chat_context: ChatContext
) -> None:
with patch.object(middleware._processor, "process_messages", return_value=(False, "user-123")) as mock_proc:
next_called = False
async def mock_next(ctx: ChatContext) -> None:
nonlocal next_called
next_called = True
class Result:
def __init__(self):
self.messages = [ChatMessage("assistant", ["Hi there"])]
ctx.result = Result()
await middleware.process(chat_context, mock_next)
assert next_called
assert mock_proc.call_count == 2
assert chat_context.result.messages[0].role == "assistant"
async def test_blocks_prompt(self, middleware: PurviewChatPolicyMiddleware, chat_context: ChatContext) -> None:
with patch.object(middleware._processor, "process_messages", return_value=(True, "user-123")):
async def mock_next(ctx: ChatContext) -> None: # should not run
raise AssertionError("next should not be called when prompt blocked")
await middleware.process(chat_context, mock_next)
assert chat_context.terminate
assert chat_context.result
assert hasattr(chat_context.result, "messages")
msg = chat_context.result.messages[0]
assert msg.role in ("system", "system")
assert "blocked" in msg.text.lower()
async def test_blocks_response(self, middleware: PurviewChatPolicyMiddleware, chat_context: ChatContext) -> None:
call_state = {"count": 0}
async def side_effect(messages, activity, user_id=None):
call_state["count"] += 1
should_block = call_state["count"] == 2
return (should_block, "user-123")
with patch.object(middleware._processor, "process_messages", side_effect=side_effect):
async def mock_next(ctx: ChatContext) -> None:
class Result:
def __init__(self):
self.messages = [ChatMessage("assistant", ["Sensitive output"])] # pragma: no cover
ctx.result = Result()
await middleware.process(chat_context, mock_next)
assert call_state["count"] == 2
msgs = getattr(chat_context.result, "messages", None) or chat_context.result
first_msg = msgs[0]
assert first_msg.role in ("system", "system")
assert "blocked" in first_msg.text.lower()
async def test_streaming_skips_post_check(self, middleware: PurviewChatPolicyMiddleware) -> None:
chat_client = DummyChatClient()
chat_options = MagicMock()
chat_options.model = "test-model"
streaming_context = ChatContext(
chat_client=chat_client,
messages=[ChatMessage("user", ["Hello"])],
options=chat_options,
is_streaming=True,
)
with patch.object(middleware._processor, "process_messages", return_value=(False, "user-123")) as mock_proc:
async def mock_next(ctx: ChatContext) -> None:
ctx.result = MagicMock()
await middleware.process(streaming_context, mock_next)
assert mock_proc.call_count == 1
async def test_chat_middleware_handles_post_check_exception(
self, middleware: PurviewChatPolicyMiddleware, chat_context: ChatContext
) -> None:
"""Test that exceptions in post-check are logged but don't affect result when ignore_exceptions=True."""
# Set ignore_exceptions to True to test exception suppression
middleware._settings.ignore_exceptions = True
call_count = 0
async def mock_process_messages(*args, **kwargs):
nonlocal call_count
call_count += 1
if call_count == 1:
return (False, "user-123") # Pre-check succeeds
raise Exception("Post-check error") # Post-check fails
with patch.object(middleware._processor, "process_messages", side_effect=mock_process_messages):
async def mock_next(ctx: ChatContext) -> None:
result = MagicMock()
result.messages = [ChatMessage("assistant", ["Response"])]
ctx.result = result
await middleware.process(chat_context, mock_next)
# Should have been called twice (pre and post)
assert call_count == 2
# Result should still be set
assert chat_context.result is not None
async def test_chat_middleware_uses_consistent_user_id(
self, middleware: PurviewChatPolicyMiddleware, chat_context: ChatContext
) -> None:
"""Test that the same user_id from pre-check is used in post-check."""
captured_user_ids = []
async def mock_process_messages(messages, activity, user_id=None):
captured_user_ids.append(user_id)
return (False, "resolved-user-123")
with patch.object(middleware._processor, "process_messages", side_effect=mock_process_messages):
async def mock_next(ctx: ChatContext) -> None:
result = MagicMock()
result.messages = [ChatMessage("assistant", ["Response"])]
ctx.result = result
await middleware.process(chat_context, mock_next)
# Should have been called twice
assert len(captured_user_ids) == 2
# First call should have None (no user_id provided yet)
assert captured_user_ids[0] is None
# Second call should have the resolved user_id from first call
assert captured_user_ids[1] == "resolved-user-123"
async def test_chat_middleware_handles_payment_required_pre_check(self, mock_credential: AsyncMock) -> None:
"""Test that 402 in pre-check is handled based on settings."""
from agent_framework_purview._exceptions import PurviewPaymentRequiredError
# Test with ignore_payment_required=False
settings = PurviewSettings(app_name="Test App", ignore_payment_required=False)
middleware = PurviewChatPolicyMiddleware(mock_credential, settings)
chat_client = DummyChatClient()
chat_options = MagicMock()
chat_options.model = "test-model"
context = ChatContext(chat_client=chat_client, messages=[ChatMessage("user", ["Hello"])], options=chat_options)
async def mock_process_messages(*args, **kwargs):
raise PurviewPaymentRequiredError("Payment required")
with patch.object(middleware._processor, "process_messages", side_effect=mock_process_messages):
async def mock_next(ctx: ChatContext) -> None:
raise AssertionError("next should not be called")
# Should raise the exception
with pytest.raises(PurviewPaymentRequiredError):
await middleware.process(context, mock_next)
async def test_chat_middleware_handles_payment_required_post_check(self, mock_credential: AsyncMock) -> None:
"""Test that 402 in post-check is raised when ignore_payment_required=False."""
from agent_framework_purview._exceptions import PurviewPaymentRequiredError
settings = PurviewSettings(app_name="Test App", ignore_payment_required=False)
middleware = PurviewChatPolicyMiddleware(mock_credential, settings)
chat_client = DummyChatClient()
chat_options = MagicMock()
chat_options.model = "test-model"
context = ChatContext(chat_client=chat_client, messages=[ChatMessage("user", ["Hello"])], options=chat_options)
call_count = 0
async def side_effect(*args, **kwargs):
nonlocal call_count
call_count += 1
if call_count == 1:
return (False, "user-123")
raise PurviewPaymentRequiredError("Payment required")
with patch.object(middleware._processor, "process_messages", side_effect=side_effect):
async def mock_next(ctx: ChatContext) -> None:
result = MagicMock()
result.messages = [ChatMessage("assistant", ["OK"])]
ctx.result = result
with pytest.raises(PurviewPaymentRequiredError):
await middleware.process(context, mock_next)
async def test_chat_middleware_ignores_payment_required_when_configured(self, mock_credential: AsyncMock) -> None:
"""Test that 402 is ignored when ignore_payment_required=True."""
from agent_framework_purview._exceptions import PurviewPaymentRequiredError
settings = PurviewSettings(app_name="Test App", ignore_payment_required=True)
middleware = PurviewChatPolicyMiddleware(mock_credential, settings)
chat_client = DummyChatClient()
chat_options = MagicMock()
chat_options.model = "test-model"
context = ChatContext(chat_client=chat_client, messages=[ChatMessage("user", ["Hello"])], options=chat_options)
async def mock_process_messages(*args, **kwargs):
raise PurviewPaymentRequiredError("Payment required")
with patch.object(middleware._processor, "process_messages", side_effect=mock_process_messages):
async def mock_next(ctx: ChatContext) -> None:
result = MagicMock()
result.messages = [ChatMessage("assistant", ["Response"])]
context.result = result
# Should not raise, just log
await middleware.process(context, mock_next)
# Next should have been called
assert context.result is not None
async def test_chat_middleware_handles_result_without_messages_attribute(
self, middleware: PurviewChatPolicyMiddleware, chat_context: ChatContext
) -> None:
"""Test middleware handles result that doesn't have messages attribute."""
with patch.object(middleware._processor, "process_messages", return_value=(False, "user-123")):
async def mock_next(ctx: ChatContext) -> None:
# Set result to something without messages attribute
ctx.result = "Some string result"
await middleware.process(chat_context, mock_next)
# Should not crash, result should be unchanged
assert chat_context.result == "Some string result"
async def test_chat_middleware_with_ignore_exceptions(self, mock_credential: AsyncMock) -> None:
"""Test that middleware respects ignore_exceptions setting."""
settings = PurviewSettings(app_name="Test App", ignore_exceptions=True)
middleware = PurviewChatPolicyMiddleware(mock_credential, settings)
chat_client = DummyChatClient()
chat_options = MagicMock()
chat_options.model = "test-model"
context = ChatContext(chat_client=chat_client, messages=[ChatMessage("user", ["Hello"])], options=chat_options)
async def mock_process_messages(*args, **kwargs):
raise ValueError("Some error")
with patch.object(middleware._processor, "process_messages", side_effect=mock_process_messages):
async def mock_next(ctx: ChatContext) -> None:
result = MagicMock()
result.messages = [ChatMessage("assistant", ["Response"])]
context.result = result
# Should not raise, just log
await middleware.process(context, mock_next)
# Next should have been called
assert context.result is not None
async def test_chat_middleware_raises_on_pre_check_exception_when_ignore_exceptions_false(
self, mock_credential: AsyncMock
) -> None:
"""Test that exceptions are propagated by default when ignore_exceptions=False."""
settings = PurviewSettings(app_name="Test App", ignore_exceptions=False)
middleware = PurviewChatPolicyMiddleware(mock_credential, settings)
chat_client = DummyChatClient()
chat_options = MagicMock()
chat_options.model = "test-model"
context = ChatContext(chat_client=chat_client, messages=[ChatMessage("user", ["Hello"])], options=chat_options)
with patch.object(middleware._processor, "process_messages", side_effect=ValueError("boom")):
async def mock_next(_: ChatContext) -> None:
raise AssertionError("next should not be called")
with pytest.raises(ValueError, match="boom"):
await middleware.process(context, mock_next)
async def test_chat_middleware_raises_on_post_check_exception_when_ignore_exceptions_false(
self, mock_credential: AsyncMock
) -> None:
"""Test that post-check exceptions are propagated by default."""
settings = PurviewSettings(app_name="Test App", ignore_exceptions=False)
middleware = PurviewChatPolicyMiddleware(mock_credential, settings)
chat_client = DummyChatClient()
chat_options = MagicMock()
chat_options.model = "test-model"
context = ChatContext(chat_client=chat_client, messages=[ChatMessage("user", ["Hello"])], options=chat_options)
call_count = 0
async def side_effect(*args, **kwargs):
nonlocal call_count
call_count += 1
if call_count == 1:
return (False, "user-123")
raise ValueError("post")
with patch.object(middleware._processor, "process_messages", side_effect=side_effect):
async def mock_next(ctx: ChatContext) -> None:
result = MagicMock()
result.messages = [ChatMessage("assistant", ["OK"])]
ctx.result = result
with pytest.raises(ValueError, match="post"):
await middleware.process(context, mock_next)
Reference in New Issue View Git Blame Copy Permalink