mirror of
https://github.com/microsoft/agent-framework.git
synced 2026-06-16 21:04:09 +08:00
3dc59c83b5
* WIP * big update to new ResponseStream model * fixed tests and typing * fixed tests and typing * fixed tools typevar import * fix * mypy fix * mypy fixes and some cleanup * fix missing quoted names * and client * fix imports agui * fix anthropic override * fix agui * fix ag ui * fix import * fix anthropic types * fix mypy * refactoring * updated typing * fix 3.11 * fixes * redid layering of chat clients and agents * redid layering of chat clients and agents * Fix lint, type, and test issues after rebase - Add @overload decorators to AgentProtocol.run() for type compatibility - Add missing docstring params (middleware, function_invocation_configuration) - Fix TODO format (TD002) by adding author tags - Fix broken observability tests from upstream: - Replace non-existent use_instrumentation with direct instantiation - Replace non-existent use_agent_instrumentation with AgentTelemetryLayer mixin - Fix get_streaming_response to use get_response(stream=True) - Add AgentInitializationError import - Update streaming exception tests to match actual behavior * Fix AgentExecutionException import error in test_agents.py - Replace non-existent AgentExecutionException with AgentRunException * Fix test import and asyncio deprecation issues - Add 'tests' to pythonpath in ag-ui pyproject.toml for utils_test_ag_ui import - Replace deprecated asyncio.get_event_loop().run_until_complete with asyncio.run * Fix azure-ai test failures - Update _prepare_options patching to use correct class path - Fix test_to_azure_ai_agent_tools_web_search_missing_connection to clear env vars * Convert ag-ui utils_test_ag_ui.py to conftest.py - Move test utilities to conftest.py for proper pytest discovery - Update all test imports to use conftest instead of utils_test_ag_ui - Remove old utils_test_ag_ui.py file - Revert pythonpath change in pyproject.toml * fix: use relative imports for ag-ui test utilities * fix agui * Rename Bare*Client to Raw*Client and BaseChatClient - Renamed BareChatClient to BaseChatClient (abstract base class) - Renamed BareOpenAIChatClient to RawOpenAIChatClient - Renamed BareOpenAIResponsesClient to RawOpenAIResponsesClient - Renamed BareAzureAIClient to RawAzureAIClient - Added warning docstrings to Raw* classes about layer ordering - Updated README in samples/getting_started/agents/custom with layer docs - Added test for span ordering with function calling * Fix layer ordering: FunctionInvocationLayer before ChatTelemetryLayer This ensures each inner LLM call gets its own telemetry span, resulting in the correct span sequence: chat -> execute_tool -> chat Updated all production clients and test mocks to use correct ordering: - ChatMiddlewareLayer (first) - FunctionInvocationLayer (second) - ChatTelemetryLayer (third) - BaseChatClient/Raw...Client (fourth) * Remove run_stream usage * Fix conversation_id propagation * Python: Add BaseAgent implementation for Claude Agent SDK (#3509) * Added ClaudeAgent implementation * Updated streaming logic * Small updates * Small update * Fixes * Small fix * Naming improvements * Updated imports * Addressed comments * Updated package versions * Update Claude agent connector layering * fix test and plugin * Store function middleware in invocation layer * Fix telemetry streaming and ag-ui tests * Remove legacy ag-ui tests folder * updates * Remove terminate flag from FunctionInvocationContext, use MiddlewareTermination instead - Remove terminate attribute from FunctionInvocationContext - Add result attribute to MiddlewareTermination to carry function results - FunctionMiddlewarePipeline.execute() now lets MiddlewareTermination propagate - _auto_invoke_function captures context.result in exception before re-raising - _try_execute_function_calls catches MiddlewareTermination and sets should_terminate - Fix handoff middleware to append to chat_client.function_middleware directly - Update tests to use raise MiddlewareTermination instead of context.terminate - Add middleware flow documentation in samples/concepts/tools/README.md - Fix ag-ui to use FunctionMiddlewarePipeline instead of removed create_function_middleware_pipeline * fix: remove references to removed terminate flag in purview tests, add type ignore * fix: move _test_utils.py from package to test folder * fix: call get_final_response() to trigger context provider notification in streaming test * fix: correct broken links in tools README * docs: clarify default middleware behavior in summary table * fix: ensure inner stream result hooks are called when using map()/from_awaitable() * Fix mypy type errors * Address PR review comments on observability.py - Remove TODO comment about unconsumed streams, add explanatory note instead - Remove redundant _close_span cleanup hook (already called in _finalize_stream) - Clarify behavior: cleanup hooks run after stream iteration, if stream is not consumed the span remains open until garbage collected * Remove gen_ai.client.operation.duration from span attributes Duration is a metrics-only attribute per OpenTelemetry semantic conventions. It should be recorded to the histogram but not set as a span attribute. * Remove duration from _get_response_attributes, pass directly to _capture_response Duration is a metrics-only attribute. It's now passed directly to _capture_response instead of being included in the attributes dict that gets set on the span. * Remove redundant _close_span cleanup hook in AgentTelemetryLayer _finalize_stream already calls _close_span() in its finally block, so adding it as a separate cleanup hook is redundant. * Use weakref.finalize to close span when stream is garbage collected If a user creates a streaming response but never consumes it, the cleanup hooks won't run. Now we register a weak reference finalizer that will close the span when the stream object is garbage collected, ensuring spans don't leak in this scenario. * Fix _get_finalizers_from_stream to use _result_hooks attribute Renamed function to _get_result_hooks_from_stream and fixed it to look for the _result_hooks attribute which is the correct name in ResponseStream class. * Add missing asyncio import in test_request_info_mixin.py * Fix leftover merge conflict marker in image_generation sample * Update integration tests * Fix integration tests: increase max_iterations from 1 to 2 Tests with tool_choice options require at least 2 iterations: 1. First iteration to get function call and execute the tool 2. Second iteration to get the final text response With max_iterations=1, streaming tests would return early with only the function call/result but no final text content. * Fix duplicate function call error in conversation-based APIs When using conversation_id (for Responses/Assistants APIs), the server already has the function call message from the previous response. We should only send the new function result message, not all messages including the function call which would cause a duplicate ID error. Fix: When conversation_id is set, only send the last message (the tool result) instead of all response.messages. * Add regression test for conversation_id propagation between tool iterations Port test from PR #3664 with updates for new streaming API pattern. Tests that conversation_id is properly updated in options dict during function invocation loop iterations. * Fix tool_choice=required to return after tool execution When tool_choice is 'required', the user's intent is to force exactly one tool call. After the tool executes, return immediately with the function call and result - don't continue to call the model again. This fixes integration tests that were failing with empty text responses because with tool_choice=required, the model would keep returning function calls instead of text. Also adds regression tests for: - conversation_id propagation between tool iterations (from PR #3664) - tool_choice=required returns after tool execution * Document tool_choice behavior in tools README - Add table explaining tool_choice values (auto, none, required) - Explain why tool_choice=required returns immediately after tool execution - Add code example showing the difference between required and auto - Update flow diagram to show the early return path for tool_choice=required * Fix tool_choice=None behavior - don't default to 'auto' Remove the hardcoded default of 'auto' for tool_choice in ChatAgent init. When tool_choice is not specified (None), it will now not be sent to the API, allowing the API's default behavior to be used. Users who want tool_choice='auto' can still explicitly set it either in default_options or at runtime. Fixes #3585 * Fix tool_choice=none should not remove tools In OpenAI Assistants client, tools were not being sent when tool_choice='none'. This was incorrect - tool_choice='none' means the model won't call tools, but tools should still be available in the request (they may be used later in the conversation). Fixes #3585 * Add test for tool_choice=none preserving tools Adds a regression test to ensure that when tool_choice='none' is set but tools are provided, the tools are still sent to the API. This verifies the fix for #3585. * Fix tool_choice=none should not remove tools in all clients Apply the same fix to OpenAI Responses client and Azure AI client: - OpenAI Responses: Remove else block that popped tool_choice/parallel_tool_calls - Azure AI: Remove tool_choice != 'none' check when adding tools When tool_choice='none', the model won't call tools, but tools should still be sent to the API so they're available for future turns. Also update README to clarify tool_choice=required supports multiple tools. Fixes #3585 * Keep tool_choice even when tools is None Move tool_choice processing outside of the 'if tools' block in OpenAI Responses client so tool_choice is sent to the API even when no tools are provided. * Update test to match new parallel_tool_calls behavior Changed test_prepare_options_removes_parallel_tool_calls_when_no_tools to test_prepare_options_preserves_parallel_tool_calls_when_no_tools to reflect that parallel_tool_calls is now preserved even when no tools are present, consistent with the tool_choice behavior. * Fix ChatMessage API and Role enum usage after rebase - Update ChatMessage instantiation to use keyword args (role=, text=, contents=) - Fix Role enum comparisons to use .value for string comparison - Add created_at to AgentResponse in error handling - Fix AgentResponse.from_updates -> from_agent_run_response_updates - Fix DurableAgentStateMessage.from_chat_message to convert Role enum to string - Add Role import where needed * Fix additional ChatMessage API and method name changes - Fix ChatMessage usage in workflow files (use text= instead of contents= for strings) - Fix AgentResponse.from_updates -> from_agent_run_response_updates in workflow files - Fix test files for ChatMessage and Role enum usage * Fix remaining ChatMessage API usage in test files * Fix more ChatMessage and Role API changes in source and test files - Fix ChatMessage in _magentic.py replan method - Fix Role enum comparison in test assertions - Fix remaining test files with old ChatMessage syntax * Fix ChatMessage and Role API changes across packages - Add Role import where missing - Fix ChatMessage signature: positional args to keyword args (role=, text=, contents=) - Fix Role enum comparisons: .role.value instead of .role string - Fix FinishReason enum usage in ag-ui event converters - Rename AgentResponse.from_updates to from_agent_run_response_updates in ag-ui Fixes API compatibility after Types API Review improvements merge * Fix ChatMessage and Role API changes in github_copilot tests * Fix ChatMessage and Role API changes in redis and github_copilot packages - Fix redis provider: Role enum comparison using .value - Fix redis tests: ChatMessage signature and Role comparisons - Fix github_copilot tests: ChatMessage signature and Role comparisons - Update docstring examples in redis chat message store * Fix ChatMessage and Role API changes in devui package - Fix executor: ChatMessage signature change - Fix conversations: Role enum to string conversion in two places - Fix tests: ChatMessage signatures and Role comparisons * Fix ChatMessage and Role API changes in a2a and lab packages - Fix a2a tests: Role comparisons and ChatMessage signatures - Fix lab tau2 source: Role enum comparison in flip_messages, log_messages, sliding_window - Fix lab tau2 tests: ChatMessage signatures and Role comparisons * Remove duplicate test files from ag-ui/tests (tests are in ag_ui_tests) * Fix ChatMessage and Role API changes across packages After rebasing on upstream/main which merged PR #3647 (Types API Review improvements), fix all packages to use the new API: - ChatMessage: Use keyword args (role=, text=, contents=) instead of positional args - Role: Compare using .value attribute since it's now an enum Packages fixed: - ag-ui: Fixed Role value extraction bugs in _message_adapters.py - anthropic: Fixed ChatMessage and Role comparisons in tests - azure-ai: Fixed Role comparison in _client.py - azure-ai-search: Fixed ChatMessage and Role in source/tests - bedrock: Fixed ChatMessage signatures in tests - chatkit: Fixed ChatMessage and Role in source/tests - copilotstudio: Fixed ChatMessage and Role in tests - declarative: Fixed ChatMessage in _executors_agents.py - mem0: Fixed ChatMessage and Role in source/tests - purview: Fixed ChatMessage in source/tests * Fix mypy errors for ChatMessage and Role API changes - durabletask: Use str() fallback in role value extraction - core: Fix ChatMessage in _orchestrator_helpers.py to use keyword args - core: Add type ignore for _conversation_state.py contents deserialization - ag-ui: Fix type ignore comments (call-overload instead of arg-type) - azure-ai-search: Fix get_role_value type hint to accept Any - lab: Move get_role_value to module level with Any type hint * Improve CI test timeout configuration - Increase job timeout from 10 to 15 minutes - Reduce per-test timeout to 60s (was 900s/300s) - Add --timeout_method thread for better timeout handling - Add --timeout-verbose to see which tests are slow - Reduce retries from 3 to 2 and delay from 10s to 5s This ensures individual test timeouts are shorter than the job timeout, providing better visibility when tests hang. With 60s timeout and 2 retries, worst case per test is ~180s. * Fix ChatMessage API usage in docstrings and source - Fix ChatMessage positional args in docstrings: _serialization.py, _threads.py, _middleware.py - Fix ChatMessage in tau2 runner.py - Fix role comparison in _orchestrator_helpers.py to use .value - Fix role comparison in _group_chat.py docstring example - Fix role assertions in test_durable_entities.py to use .value * Revert tool_choice/parallel_tool_calls changes - must be removed when no tools OpenAI API requires tool_choice and parallel_tool_calls to only be present when tools are specified. Restored the logic that removes these options when there are no tools. - Restored check in _chat_client.py to remove tool_choice and parallel_tool_calls when no tools present - Restored same logic in _responses_client.py - Reverted test to expect the correct behavior * fixed issue in tests * fix: resolve merge conflict markers in ag-ui tests * fix: restructure ag-ui tests and fix Role/FinishReason to use string types * fix: streaming function invocation and middleware termination - Refactor streaming function invocation to use get_final_response() on inner streams - Fix MiddlewareTermination to accept result parameter for passing results - Fix _AutoHandoffMiddleware to use MiddlewareTermination instead of context.terminate - Fix AgentMiddlewareLayer.run() to properly forward function/chat middleware - Remove duplicate middleware registration in AgentMiddlewareLayer.__init__ - Fix exception handling in _auto_invoke_function to properly capture termination - Fix mypy errors in core package - Update tests to use stream=True parameter for unified run API * fix all tests command * Refactor integration tests to use pytest fixtures - Merge testutils.py into conftest.py for azurefunctions integration tests - Merge dt_testutils.py into conftest.py for durabletask integration tests - Convert all integration tests to use fixtures instead of direct imports (fixes ModuleNotFoundError with --import-mode=importlib) - Add sample_helper fixture for azurefunctions tests - Add agent_client_factory and orchestration_helper fixtures for durabletask - Integration tests now skip with descriptive messages when services unavailable - Restructure devui tests into tests/devui/ with proper conftest.py - Add test organization guidelines to CODING_STANDARD.md - Remove __init__.py from test directories per pytest best practices * Fix pytest_collection_modifyitems to only skip integration tests The hook was skipping all tests in the test session, not just integration tests. Now it only skips items in the integration_tests directory. * Fix mem0 tests failing on Python 3.13 Use patch.object on the imported module instead of @patch with string path to ensure the mock takes effect regardless of import timing. * fix mem0 * another attempt for mem0 * fix for mem0 * fix mem0 * Increase worker initialization wait time in durabletask tests Increase from 2 to 8 seconds to allow time for: - Python startup and module imports - Azure OpenAI client creation - Agent registration with DTS worker - Worker connection to DTS This helps prevent test failures in CI where the first tests may run before the worker is fully ready to process requests. * Fix streaming test to use ResponseStream with finalizer The _consume_stream method now expects a ResponseStream that can provide a final AgentResponse via get_final_response(). Update the test to use ResponseStream with AgentResponse.from_updates as the finalizer. * Fix MockToolCallingAgent to use new ResponseStream API and update samples * small updates to run_stream to run * fix sub workflow * temp fix for az func test --------- Co-authored-by: Dmytro Struk <13853051+dmytrostruk@users.noreply.github.com>
704 lines
30 KiB
Python
704 lines
30 KiB
Python
# Copyright (c) Microsoft. All rights reserved.
|
|
|
|
"""Tests for Purview processor."""
|
|
|
|
from unittest.mock import AsyncMock, MagicMock, patch
|
|
|
|
import pytest
|
|
from agent_framework import ChatMessage
|
|
|
|
from agent_framework_purview import PurviewAppLocation, PurviewLocationType, PurviewSettings
|
|
from agent_framework_purview._models import (
|
|
Activity,
|
|
DlpAction,
|
|
DlpActionInfo,
|
|
ProcessContentResponse,
|
|
RestrictionAction,
|
|
)
|
|
from agent_framework_purview._processor import ScopedContentProcessor, _is_valid_guid
|
|
|
|
|
|
class TestGuidValidation:
|
|
"""Test GUID validation helper."""
|
|
|
|
def test_valid_guid(self) -> None:
|
|
"""Test _is_valid_guid with valid GUIDs."""
|
|
assert _is_valid_guid("12345678-1234-1234-1234-123456789012")
|
|
assert _is_valid_guid("a1b2c3d4-e5f6-4a5b-8c9d-0e1f2a3b4c5d")
|
|
|
|
def test_invalid_guid(self) -> None:
|
|
"""Test _is_valid_guid with invalid GUIDs."""
|
|
assert not _is_valid_guid("not-a-guid")
|
|
assert not _is_valid_guid("")
|
|
assert not _is_valid_guid(None)
|
|
|
|
|
|
class TestScopedContentProcessor:
|
|
"""Test ScopedContentProcessor functionality."""
|
|
|
|
@pytest.fixture
|
|
def mock_client(self) -> AsyncMock:
|
|
"""Create a mock Purview client."""
|
|
client = AsyncMock()
|
|
client.get_user_info_from_token = AsyncMock(
|
|
return_value={
|
|
"tenant_id": "12345678-1234-1234-1234-123456789012",
|
|
"user_id": "12345678-1234-1234-1234-123456789012",
|
|
"client_id": "12345678-1234-1234-1234-123456789012",
|
|
}
|
|
)
|
|
return client
|
|
|
|
@pytest.fixture
|
|
def settings_with_defaults(self) -> PurviewSettings:
|
|
"""Create settings with default values."""
|
|
app_location = PurviewAppLocation(
|
|
location_type=PurviewLocationType.APPLICATION, location_value="12345678-1234-1234-1234-123456789012"
|
|
)
|
|
return PurviewSettings(
|
|
app_name="Test App",
|
|
tenant_id="12345678-1234-1234-1234-123456789012",
|
|
purview_app_location=app_location,
|
|
)
|
|
|
|
@pytest.fixture
|
|
def settings_without_defaults(self) -> PurviewSettings:
|
|
"""Create settings without default values (requiring token info)."""
|
|
return PurviewSettings(app_name="Test App")
|
|
|
|
@pytest.fixture
|
|
def processor(self, mock_client: AsyncMock, settings_with_defaults: PurviewSettings) -> ScopedContentProcessor:
|
|
"""Create a ScopedContentProcessor with mock client."""
|
|
return ScopedContentProcessor(mock_client, settings_with_defaults)
|
|
|
|
async def test_processor_initialization(
|
|
self, mock_client: AsyncMock, settings_with_defaults: PurviewSettings
|
|
) -> None:
|
|
"""Test ScopedContentProcessor initialization."""
|
|
processor = ScopedContentProcessor(mock_client, settings_with_defaults)
|
|
|
|
assert processor._client == mock_client
|
|
assert processor._settings == settings_with_defaults
|
|
|
|
async def test_process_messages_with_defaults(self, processor: ScopedContentProcessor) -> None:
|
|
"""Test process_messages with settings that have defaults."""
|
|
messages = [
|
|
ChatMessage(role="user", text="Hello"),
|
|
ChatMessage(role="assistant", text="Hi there"),
|
|
]
|
|
|
|
with patch.object(processor, "_map_messages", return_value=([], None)) as mock_map:
|
|
should_block, user_id = await processor.process_messages(messages, Activity.UPLOAD_TEXT)
|
|
|
|
assert should_block is False
|
|
assert user_id is None
|
|
mock_map.assert_called_once_with(messages, Activity.UPLOAD_TEXT, None)
|
|
|
|
async def test_process_messages_blocks_content(
|
|
self, processor: ScopedContentProcessor, process_content_request_factory
|
|
) -> None:
|
|
"""Test process_messages returns True when content should be blocked."""
|
|
messages = [ChatMessage(role="user", text="Sensitive content")]
|
|
|
|
mock_request = process_content_request_factory("Sensitive content")
|
|
|
|
mock_response = ProcessContentResponse(**{
|
|
"policyActions": [DlpActionInfo(action=DlpAction.BLOCK_ACCESS, restrictionAction=RestrictionAction.BLOCK)]
|
|
})
|
|
|
|
with (
|
|
patch.object(processor, "_map_messages", return_value=([mock_request], "user-123")),
|
|
patch.object(processor, "_process_with_scopes", return_value=mock_response),
|
|
):
|
|
should_block, user_id = await processor.process_messages(messages, Activity.UPLOAD_TEXT)
|
|
|
|
assert should_block is True
|
|
assert user_id == "user-123"
|
|
|
|
async def test_map_messages_creates_requests(
|
|
self, processor: ScopedContentProcessor, mock_client: AsyncMock
|
|
) -> None:
|
|
"""Test _map_messages creates ProcessContentRequest objects."""
|
|
messages = [
|
|
ChatMessage(
|
|
role="user",
|
|
text="Test message",
|
|
message_id="msg-123",
|
|
author_name="12345678-1234-1234-1234-123456789012",
|
|
),
|
|
]
|
|
|
|
requests, user_id = await processor._map_messages(messages, Activity.UPLOAD_TEXT)
|
|
|
|
assert len(requests) == 1
|
|
assert requests[0].user_id == "12345678-1234-1234-1234-123456789012"
|
|
assert requests[0].tenant_id == "12345678-1234-1234-1234-123456789012"
|
|
assert user_id == "12345678-1234-1234-1234-123456789012"
|
|
|
|
async def test_map_messages_without_defaults_gets_token_info(self, mock_client: AsyncMock) -> None:
|
|
"""Test _map_messages gets token info when settings lack some defaults."""
|
|
settings = PurviewSettings(app_name="Test App", tenant_id="12345678-1234-1234-1234-123456789012")
|
|
processor = ScopedContentProcessor(mock_client, settings)
|
|
messages = [ChatMessage(role="user", text="Test", message_id="msg-123")]
|
|
|
|
requests, user_id = await processor._map_messages(messages, Activity.UPLOAD_TEXT)
|
|
|
|
mock_client.get_user_info_from_token.assert_called_once()
|
|
assert len(requests) == 1
|
|
assert user_id is not None
|
|
|
|
async def test_map_messages_raises_on_missing_tenant_id(self, mock_client: AsyncMock) -> None:
|
|
"""Test _map_messages raises ValueError when tenant_id cannot be determined."""
|
|
settings = PurviewSettings(app_name="Test App") # No tenant_id
|
|
processor = ScopedContentProcessor(mock_client, settings)
|
|
|
|
mock_client.get_user_info_from_token = AsyncMock(
|
|
return_value={"user_id": "test-user", "client_id": "test-client"}
|
|
)
|
|
|
|
messages = [ChatMessage(role="user", text="Test", message_id="msg-123")]
|
|
|
|
with pytest.raises(ValueError, match="Tenant id required"):
|
|
await processor._map_messages(messages, Activity.UPLOAD_TEXT)
|
|
|
|
async def test_check_applicable_scopes_no_scopes(
|
|
self, processor: ScopedContentProcessor, process_content_request_factory
|
|
) -> None:
|
|
"""Test _check_applicable_scopes when no scopes are returned."""
|
|
from agent_framework_purview._models import ProtectionScopesResponse
|
|
|
|
request = process_content_request_factory()
|
|
response = ProtectionScopesResponse(**{"value": None})
|
|
|
|
should_process, actions, execution_mode = processor._check_applicable_scopes(request, response)
|
|
|
|
assert should_process is False
|
|
assert actions == []
|
|
|
|
async def test_check_applicable_scopes_with_block_action(
|
|
self, processor: ScopedContentProcessor, process_content_request_factory
|
|
) -> None:
|
|
"""Test _check_applicable_scopes identifies block actions."""
|
|
from agent_framework_purview._models import (
|
|
PolicyLocation,
|
|
PolicyScope,
|
|
ProtectionScopeActivities,
|
|
ProtectionScopesResponse,
|
|
)
|
|
|
|
request = process_content_request_factory()
|
|
|
|
block_action = DlpActionInfo(action=DlpAction.BLOCK_ACCESS, restrictionAction=RestrictionAction.BLOCK)
|
|
scope_location = PolicyLocation(**{
|
|
"@odata.type": "microsoft.graph.policyLocationApplication",
|
|
"value": "app-id",
|
|
})
|
|
scope = PolicyScope(**{
|
|
"policyActions": [block_action],
|
|
"activities": ProtectionScopeActivities.UPLOAD_TEXT,
|
|
"locations": [scope_location],
|
|
})
|
|
response = ProtectionScopesResponse(**{"value": [scope]})
|
|
|
|
should_process, actions, execution_mode = processor._check_applicable_scopes(request, response)
|
|
|
|
assert should_process is True
|
|
assert len(actions) == 1
|
|
assert actions[0].action == DlpAction.BLOCK_ACCESS
|
|
|
|
async def test_combine_policy_actions(self, processor: ScopedContentProcessor) -> None:
|
|
"""Test _combine_policy_actions merges action lists."""
|
|
action1 = DlpActionInfo(action=DlpAction.BLOCK_ACCESS, restrictionAction=RestrictionAction.BLOCK)
|
|
action2 = DlpActionInfo(action=DlpAction.OTHER, restrictionAction=RestrictionAction.OTHER)
|
|
|
|
combined = processor._combine_policy_actions([action1], [action2])
|
|
|
|
assert len(combined) == 2
|
|
assert action1 in combined
|
|
assert action2 in combined
|
|
|
|
async def test_process_with_scopes_calls_client_methods(
|
|
self, processor: ScopedContentProcessor, mock_client: AsyncMock, process_content_request_factory
|
|
) -> None:
|
|
"""Test _process_with_scopes calls get_protection_scopes when scopes response is empty."""
|
|
from agent_framework_purview._models import (
|
|
ContentActivitiesResponse,
|
|
ProtectionScopesResponse,
|
|
)
|
|
|
|
request = process_content_request_factory()
|
|
|
|
mock_client.get_protection_scopes = AsyncMock(return_value=ProtectionScopesResponse(**{"value": []}))
|
|
mock_client.process_content = AsyncMock(
|
|
return_value=ProcessContentResponse(**{"id": "response-123", "protectionScopeState": "notModified"})
|
|
)
|
|
mock_client.send_content_activities = AsyncMock(return_value=ContentActivitiesResponse(**{"error": None}))
|
|
|
|
response = await processor._process_with_scopes(request)
|
|
|
|
mock_client.get_protection_scopes.assert_called_once()
|
|
# When no scopes apply, process_content is not called (activities are sent in background)
|
|
mock_client.process_content.assert_not_called()
|
|
# The response should have id=204 (No Content) when no scopes apply
|
|
assert response.id == "204"
|
|
|
|
async def test_process_with_scopes_ignores_unexpected_cached_value_type(
|
|
self, processor: ScopedContentProcessor, mock_client: AsyncMock, process_content_request_factory
|
|
) -> None:
|
|
"""Test that a corrupted cache entry does not crash processing."""
|
|
from agent_framework_purview._models import (
|
|
ExecutionMode,
|
|
PolicyLocation,
|
|
PolicyScope,
|
|
ProcessContentResponse,
|
|
ProtectionScopeActivities,
|
|
ProtectionScopesResponse,
|
|
)
|
|
|
|
request = process_content_request_factory()
|
|
|
|
# Return a valid, inline scope so we stay on the normal (non-background) path.
|
|
scope_location = PolicyLocation(**{
|
|
"@odata.type": "microsoft.graph.policyLocationApplication",
|
|
"value": "app-id",
|
|
})
|
|
scope = PolicyScope(**{
|
|
"activities": ProtectionScopeActivities.UPLOAD_TEXT,
|
|
"locations": [scope_location],
|
|
"execution_mode": ExecutionMode.EVALUATE_INLINE,
|
|
})
|
|
mock_client.get_protection_scopes = AsyncMock(return_value=ProtectionScopesResponse(**{"value": [scope]}))
|
|
mock_client.process_content = AsyncMock(
|
|
return_value=ProcessContentResponse(**{"id": "ok", "protectionScopeState": "notModified"})
|
|
)
|
|
|
|
# First cache read is the tenant payment key (None). Second is the scopes cache (corrupt value).
|
|
processor._cache.get = AsyncMock(side_effect=[None, "corrupt-value"]) # type: ignore[method-assign]
|
|
processor._cache.set = AsyncMock() # type: ignore[method-assign]
|
|
|
|
response = await processor._process_with_scopes(request)
|
|
|
|
assert response.id == "ok"
|
|
mock_client.get_protection_scopes.assert_called_once()
|
|
mock_client.process_content.assert_called_once()
|
|
|
|
async def test_process_with_scopes_uses_tenant_payment_exception_cache(
|
|
self, processor: ScopedContentProcessor, mock_client: AsyncMock, process_content_request_factory
|
|
) -> None:
|
|
"""Test that a cached 402 exception short-circuits all subsequent requests for the tenant."""
|
|
from agent_framework_purview._exceptions import PurviewPaymentRequiredError
|
|
|
|
request = process_content_request_factory()
|
|
|
|
processor._cache.get = AsyncMock(return_value=PurviewPaymentRequiredError("Payment required")) # type: ignore[method-assign]
|
|
|
|
with pytest.raises(PurviewPaymentRequiredError):
|
|
await processor._process_with_scopes(request)
|
|
|
|
mock_client.get_protection_scopes.assert_not_called()
|
|
|
|
async def test_process_content_background_retries_on_modified_state(
|
|
self, processor: ScopedContentProcessor, mock_client: AsyncMock, process_content_request_factory
|
|
) -> None:
|
|
"""Test offline background processing invalidates cache and retries when scope state changes."""
|
|
from agent_framework_purview._models import ProcessContentResponse
|
|
|
|
request = process_content_request_factory()
|
|
request.scope_identifier = "etag-1"
|
|
|
|
mock_client.process_content = AsyncMock(
|
|
side_effect=[
|
|
ProcessContentResponse(**{"id": "r1", "protectionScopeState": "modified"}),
|
|
ProcessContentResponse(**{"id": "r2", "protectionScopeState": "notModified"}),
|
|
]
|
|
)
|
|
processor._cache.remove = AsyncMock() # type: ignore[method-assign]
|
|
|
|
await processor._process_content_background(request, cache_key="purview:protection_scopes:abc")
|
|
|
|
processor._cache.remove.assert_called_once_with("purview:protection_scopes:abc")
|
|
assert mock_client.process_content.call_count == 2
|
|
|
|
async def test_map_messages_with_user_id_in_additional_properties(self, mock_client: AsyncMock) -> None:
|
|
"""Test user_id extraction from message additional_properties."""
|
|
settings = PurviewSettings(
|
|
app_name="Test App",
|
|
tenant_id="12345678-1234-1234-1234-123456789012",
|
|
purview_app_location=PurviewAppLocation(
|
|
location_type=PurviewLocationType.APPLICATION, location_value="app-id"
|
|
),
|
|
)
|
|
processor = ScopedContentProcessor(mock_client, settings)
|
|
|
|
messages = [
|
|
ChatMessage(
|
|
role="user",
|
|
text="Test message",
|
|
additional_properties={"user_id": "22345678-1234-1234-1234-123456789012"},
|
|
),
|
|
]
|
|
|
|
requests, user_id = await processor._map_messages(messages, Activity.UPLOAD_TEXT)
|
|
|
|
assert len(requests) == 1
|
|
assert user_id == "22345678-1234-1234-1234-123456789012"
|
|
assert requests[0].user_id == "22345678-1234-1234-1234-123456789012"
|
|
|
|
async def test_map_messages_with_provided_user_id_fallback(self, mock_client: AsyncMock) -> None:
|
|
"""Test using provided_user_id when no other source is available."""
|
|
settings = PurviewSettings(
|
|
app_name="Test App",
|
|
tenant_id="12345678-1234-1234-1234-123456789012",
|
|
purview_app_location=PurviewAppLocation(
|
|
location_type=PurviewLocationType.APPLICATION, location_value="app-id"
|
|
),
|
|
)
|
|
processor = ScopedContentProcessor(mock_client, settings)
|
|
|
|
messages = [ChatMessage(role="user", text="Test message")]
|
|
|
|
requests, user_id = await processor._map_messages(
|
|
messages, Activity.UPLOAD_TEXT, provided_user_id="32345678-1234-1234-1234-123456789012"
|
|
)
|
|
|
|
assert len(requests) == 1
|
|
assert user_id == "32345678-1234-1234-1234-123456789012"
|
|
assert requests[0].user_id == "32345678-1234-1234-1234-123456789012"
|
|
|
|
async def test_map_messages_returns_empty_when_no_user_id(self, mock_client: AsyncMock) -> None:
|
|
"""Test that empty results are returned when user_id cannot be resolved."""
|
|
settings = PurviewSettings(
|
|
app_name="Test App",
|
|
tenant_id="12345678-1234-1234-1234-123456789012",
|
|
purview_app_location=PurviewAppLocation(
|
|
location_type=PurviewLocationType.APPLICATION, location_value="app-id"
|
|
),
|
|
)
|
|
processor = ScopedContentProcessor(mock_client, settings)
|
|
|
|
messages = [ChatMessage(role="user", text="Test message")]
|
|
|
|
requests, user_id = await processor._map_messages(messages, Activity.UPLOAD_TEXT)
|
|
|
|
assert len(requests) == 0
|
|
assert user_id is None
|
|
|
|
async def test_process_content_sends_activities_when_not_applicable(
|
|
self, mock_client: AsyncMock, process_content_request_factory
|
|
) -> None:
|
|
"""Test that response is returned when scopes don't apply (activities sent in background)."""
|
|
settings = PurviewSettings(
|
|
app_name="Test App",
|
|
tenant_id="12345678-1234-1234-1234-123456789012",
|
|
purview_app_location=PurviewAppLocation(
|
|
location_type=PurviewLocationType.APPLICATION, location_value="app-id"
|
|
),
|
|
)
|
|
processor = ScopedContentProcessor(mock_client, settings)
|
|
|
|
pc_request = process_content_request_factory()
|
|
|
|
# Mock get_protection_scopes to return no applicable scopes
|
|
mock_ps_response = MagicMock()
|
|
mock_ps_response.scopes = []
|
|
mock_client.get_protection_scopes.return_value = mock_ps_response
|
|
|
|
# Mock send_content_activities to return success (called in background)
|
|
mock_ca_response = MagicMock()
|
|
mock_ca_response.error = None
|
|
mock_client.send_content_activities.return_value = mock_ca_response
|
|
|
|
response = await processor._process_with_scopes(pc_request)
|
|
|
|
mock_client.get_protection_scopes.assert_called_once()
|
|
mock_client.process_content.assert_not_called()
|
|
# Response should have id=204 when no scopes apply
|
|
assert response.id == "204"
|
|
|
|
async def test_process_content_handles_activities_error(
|
|
self, mock_client: AsyncMock, process_content_request_factory
|
|
) -> None:
|
|
"""Test that errors in background activities don't affect the response."""
|
|
settings = PurviewSettings(
|
|
app_name="Test App",
|
|
tenant_id="12345678-1234-1234-1234-123456789012",
|
|
purview_app_location=PurviewAppLocation(
|
|
location_type=PurviewLocationType.APPLICATION, location_value="app-id"
|
|
),
|
|
)
|
|
processor = ScopedContentProcessor(mock_client, settings)
|
|
|
|
pc_request = process_content_request_factory()
|
|
|
|
# Mock get_protection_scopes to return no applicable scopes
|
|
mock_ps_response = MagicMock()
|
|
mock_ps_response.scopes = []
|
|
mock_client.get_protection_scopes.return_value = mock_ps_response
|
|
|
|
# Mock send_content_activities to return error (called in background task)
|
|
mock_ca_response = MagicMock()
|
|
mock_ca_response.error = "Test error message"
|
|
mock_client.send_content_activities.return_value = mock_ca_response
|
|
|
|
response = await processor._process_with_scopes(pc_request)
|
|
|
|
# Since activities are sent in background, errors don't affect the response
|
|
# Response should have id=204 when no scopes apply
|
|
assert response.id == "204"
|
|
|
|
|
|
class TestUserIdResolution:
|
|
"""Test user ID resolution from various sources."""
|
|
|
|
@pytest.fixture
|
|
def mock_client(self) -> AsyncMock:
|
|
"""Create a mock Purview client."""
|
|
client = AsyncMock()
|
|
client.get_user_info_from_token = AsyncMock(
|
|
return_value={
|
|
"tenant_id": "12345678-1234-1234-1234-123456789012",
|
|
"user_id": "11111111-1111-1111-1111-111111111111",
|
|
"client_id": "12345678-1234-1234-1234-123456789012",
|
|
}
|
|
)
|
|
return client
|
|
|
|
@pytest.fixture
|
|
def settings(self) -> PurviewSettings:
|
|
"""Create settings."""
|
|
return PurviewSettings(
|
|
app_name="Test App",
|
|
tenant_id="12345678-1234-1234-1234-123456789012",
|
|
purview_app_location=PurviewAppLocation(
|
|
location_type=PurviewLocationType.APPLICATION, location_value="app-id"
|
|
),
|
|
)
|
|
|
|
async def test_user_id_from_token_when_no_other_source(self, mock_client: AsyncMock) -> None:
|
|
"""Test user_id is extracted from token when no other source available."""
|
|
settings = PurviewSettings(app_name="Test App") # No tenant_id or app_location
|
|
processor = ScopedContentProcessor(mock_client, settings)
|
|
|
|
messages = [ChatMessage(role="user", text="Test")]
|
|
|
|
requests, user_id = await processor._map_messages(messages, Activity.UPLOAD_TEXT)
|
|
|
|
mock_client.get_user_info_from_token.assert_called_once()
|
|
assert user_id == "11111111-1111-1111-1111-111111111111"
|
|
|
|
async def test_user_id_from_additional_properties_takes_priority(
|
|
self, mock_client: AsyncMock, settings: PurviewSettings
|
|
) -> None:
|
|
"""Test user_id from additional_properties takes priority over token."""
|
|
processor = ScopedContentProcessor(mock_client, settings)
|
|
|
|
messages = [
|
|
ChatMessage(
|
|
role="user",
|
|
text="Test",
|
|
additional_properties={"user_id": "22222222-2222-2222-2222-222222222222"},
|
|
)
|
|
]
|
|
|
|
requests, user_id = await processor._map_messages(messages, Activity.UPLOAD_TEXT)
|
|
|
|
# Token info should not be called since we have user_id in message
|
|
mock_client.get_user_info_from_token.assert_not_called()
|
|
assert user_id == "22222222-2222-2222-2222-222222222222"
|
|
|
|
async def test_user_id_from_author_name_as_fallback(
|
|
self, mock_client: AsyncMock, settings: PurviewSettings
|
|
) -> None:
|
|
"""Test user_id is extracted from author_name when it's a valid GUID."""
|
|
processor = ScopedContentProcessor(mock_client, settings)
|
|
|
|
messages = [
|
|
ChatMessage(
|
|
role="user",
|
|
text="Test",
|
|
author_name="33333333-3333-3333-3333-333333333333",
|
|
)
|
|
]
|
|
|
|
requests, user_id = await processor._map_messages(messages, Activity.UPLOAD_TEXT)
|
|
|
|
assert user_id == "33333333-3333-3333-3333-333333333333"
|
|
|
|
async def test_author_name_ignored_if_not_valid_guid(
|
|
self, mock_client: AsyncMock, settings: PurviewSettings
|
|
) -> None:
|
|
"""Test author_name is ignored if it's not a valid GUID."""
|
|
processor = ScopedContentProcessor(mock_client, settings)
|
|
|
|
messages = [
|
|
ChatMessage(
|
|
role="user",
|
|
text="Test",
|
|
author_name="John Doe", # Not a GUID
|
|
)
|
|
]
|
|
|
|
requests, user_id = await processor._map_messages(messages, Activity.UPLOAD_TEXT)
|
|
|
|
# Should return empty since author_name is not a valid GUID
|
|
assert user_id is None
|
|
assert len(requests) == 0
|
|
|
|
async def test_provided_user_id_used_as_last_resort(
|
|
self, mock_client: AsyncMock, settings: PurviewSettings
|
|
) -> None:
|
|
"""Test provided_user_id parameter is used as last resort."""
|
|
processor = ScopedContentProcessor(mock_client, settings)
|
|
|
|
messages = [ChatMessage(role="user", text="Test")]
|
|
|
|
requests, user_id = await processor._map_messages(
|
|
messages, Activity.UPLOAD_TEXT, provided_user_id="44444444-4444-4444-4444-444444444444"
|
|
)
|
|
|
|
assert user_id == "44444444-4444-4444-4444-444444444444"
|
|
|
|
async def test_invalid_provided_user_id_ignored(self, mock_client: AsyncMock, settings: PurviewSettings) -> None:
|
|
"""Test invalid provided_user_id is ignored."""
|
|
processor = ScopedContentProcessor(mock_client, settings)
|
|
|
|
messages = [ChatMessage(role="user", text="Test")]
|
|
|
|
requests, user_id = await processor._map_messages(messages, Activity.UPLOAD_TEXT, provided_user_id="not-a-guid")
|
|
|
|
assert user_id is None
|
|
assert len(requests) == 0
|
|
|
|
async def test_multiple_messages_same_user_id(self, mock_client: AsyncMock, settings: PurviewSettings) -> None:
|
|
"""Test that all messages use the same resolved user_id."""
|
|
processor = ScopedContentProcessor(mock_client, settings)
|
|
|
|
messages = [
|
|
ChatMessage(
|
|
role="user", text="First", additional_properties={"user_id": "55555555-5555-5555-5555-555555555555"}
|
|
),
|
|
ChatMessage(role="assistant", text="Response"),
|
|
ChatMessage(role="user", text="Second"),
|
|
]
|
|
|
|
requests, user_id = await processor._map_messages(messages, Activity.UPLOAD_TEXT)
|
|
|
|
assert user_id == "55555555-5555-5555-5555-555555555555"
|
|
# All requests should have the same user_id
|
|
assert all(req.user_id == "55555555-5555-5555-5555-555555555555" for req in requests)
|
|
|
|
async def test_first_valid_user_id_in_messages_is_used(
|
|
self, mock_client: AsyncMock, settings: PurviewSettings
|
|
) -> None:
|
|
"""Test that the first valid user_id found in messages is used for all."""
|
|
processor = ScopedContentProcessor(mock_client, settings)
|
|
|
|
messages = [
|
|
ChatMessage(role="user", text="First", author_name="Not a GUID"),
|
|
ChatMessage(
|
|
role="assistant",
|
|
text="Response",
|
|
additional_properties={"user_id": "66666666-6666-6666-6666-666666666666"},
|
|
),
|
|
ChatMessage(
|
|
role="user", text="Third", additional_properties={"user_id": "77777777-7777-7777-7777-777777777777"}
|
|
),
|
|
]
|
|
|
|
requests, user_id = await processor._map_messages(messages, Activity.UPLOAD_TEXT)
|
|
|
|
# First valid user_id (from second message) should be used
|
|
assert user_id == "66666666-6666-6666-6666-666666666666"
|
|
assert all(req.user_id == "66666666-6666-6666-6666-666666666666" for req in requests)
|
|
|
|
|
|
class TestScopedContentProcessorCaching:
|
|
"""Test caching functionality in ScopedContentProcessor."""
|
|
|
|
@pytest.fixture
|
|
def mock_client(self) -> AsyncMock:
|
|
"""Create a mock Purview client."""
|
|
client = AsyncMock()
|
|
client.get_user_info_from_token = AsyncMock(
|
|
return_value={
|
|
"tenant_id": "12345678-1234-1234-1234-123456789012",
|
|
"user_id": "12345678-1234-1234-1234-123456789012",
|
|
"client_id": "12345678-1234-1234-1234-123456789012",
|
|
}
|
|
)
|
|
client.get_protection_scopes = AsyncMock()
|
|
return client
|
|
|
|
@pytest.fixture
|
|
def settings(self) -> PurviewSettings:
|
|
"""Create test settings."""
|
|
location = PurviewAppLocation(location_type=PurviewLocationType.APPLICATION, location_value="app-id")
|
|
return PurviewSettings(
|
|
app_name="Test App",
|
|
tenant_id="12345678-1234-1234-1234-123456789012",
|
|
default_user_id="12345678-1234-1234-1234-123456789012",
|
|
purview_app_location=location,
|
|
)
|
|
|
|
async def test_protection_scopes_cached_on_first_call(
|
|
self, mock_client: AsyncMock, settings: PurviewSettings
|
|
) -> None:
|
|
"""Test that protection scopes response is cached after first call."""
|
|
from agent_framework_purview._cache import InMemoryCacheProvider
|
|
from agent_framework_purview._models import ProtectionScopesResponse
|
|
|
|
cache_provider = InMemoryCacheProvider()
|
|
processor = ScopedContentProcessor(mock_client, settings, cache_provider=cache_provider)
|
|
|
|
mock_client.get_protection_scopes.return_value = ProtectionScopesResponse(
|
|
scope_identifier="scope-123", scopes=[]
|
|
)
|
|
|
|
messages = [ChatMessage(role="user", text="Test")]
|
|
|
|
await processor.process_messages(messages, Activity.UPLOAD_TEXT, user_id="12345678-1234-1234-1234-123456789012")
|
|
|
|
mock_client.get_protection_scopes.assert_called_once()
|
|
|
|
await processor.process_messages(messages, Activity.UPLOAD_TEXT, user_id="12345678-1234-1234-1234-123456789012")
|
|
|
|
mock_client.get_protection_scopes.assert_called_once()
|
|
|
|
async def test_payment_required_exception_cached_at_tenant_level(
|
|
self, mock_client: AsyncMock, settings: PurviewSettings
|
|
) -> None:
|
|
"""Test that 402 payment required exceptions are cached at tenant level."""
|
|
from agent_framework_purview._cache import InMemoryCacheProvider
|
|
from agent_framework_purview._exceptions import PurviewPaymentRequiredError
|
|
|
|
cache_provider = InMemoryCacheProvider()
|
|
processor = ScopedContentProcessor(mock_client, settings, cache_provider=cache_provider)
|
|
|
|
mock_client.get_protection_scopes.side_effect = PurviewPaymentRequiredError("Payment required")
|
|
|
|
messages = [ChatMessage(role="user", text="Test")]
|
|
|
|
with pytest.raises(PurviewPaymentRequiredError):
|
|
await processor.process_messages(
|
|
messages, Activity.UPLOAD_TEXT, user_id="12345678-1234-1234-1234-123456789012"
|
|
)
|
|
|
|
mock_client.get_protection_scopes.assert_called_once()
|
|
|
|
with pytest.raises(PurviewPaymentRequiredError):
|
|
await processor.process_messages(
|
|
messages, Activity.UPLOAD_TEXT, user_id="12345678-1234-1234-1234-123456789012"
|
|
)
|
|
|
|
mock_client.get_protection_scopes.assert_called_once()
|
|
|
|
async def test_custom_cache_provider_used(self, mock_client: AsyncMock, settings: PurviewSettings) -> None:
|
|
"""Test that custom cache provider is used when provided."""
|
|
from agent_framework_purview._cache import InMemoryCacheProvider
|
|
|
|
custom_cache = InMemoryCacheProvider(default_ttl_seconds=60)
|
|
processor = ScopedContentProcessor(mock_client, settings, cache_provider=custom_cache)
|
|
|
|
assert processor._cache is custom_cache
|
|
assert processor._cache._default_ttl == 60
|