mirror of
https://github.com/microsoft/agent-framework.git
synced 2026-06-16 21:04:09 +08:00
Python: Add Purview Middleware (#1142)
* [Py Purview] Purview Python Initial Commit * [Py Purview] Purview Python Minor Fixes * [Py Purview] Purview Python Comment Fixesish * [Py Purview] Purview Python Agent Middleware Done * [Py Purview] Purview Python Agent Middleware Done * [Py Purview] Purview Python Lint Errors * [Py Purview] Purview Python Final Hopefully * [Py Purview] Purview Python Final Hopefully * [Py Purview] Purview Python Fix ReadMe * [Py Purview] Purview Python Fix MyPy * [Py Purview] Purview Python Minor Updates on comments * [Py Purview] Purview Python Fix Build Error --------- Co-authored-by: Dmytro Struk <13853051+dmytrostruk@users.noreply.github.com>
This commit is contained in:
Rishabh Chawla
committed by
GitHub
Unverified
GitHub
parent
76ae0a62ac
commit
59da578902
@@ -0,0 +1,68 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
"""Shared pytest fixtures for Purview tests."""
|
||||
|
||||
import pytest
|
||||
|
||||
from agent_framework_purview._models import (
|
||||
Activity,
|
||||
ActivityMetadata,
|
||||
ContentToProcess,
|
||||
DeviceMetadata,
|
||||
IntegratedAppMetadata,
|
||||
OperatingSystemSpecifications,
|
||||
PolicyLocation,
|
||||
ProcessContentRequest,
|
||||
ProcessConversationMetadata,
|
||||
ProtectedAppMetadata,
|
||||
PurviewTextContent,
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def content_to_process_factory():
|
||||
"""Factory fixture to create ContentToProcess objects with test data."""
|
||||
|
||||
def _create_content(text: str = "Test") -> ContentToProcess:
|
||||
text_content = PurviewTextContent(data=text)
|
||||
metadata = ProcessConversationMetadata(
|
||||
identifier="msg-1",
|
||||
content=text_content,
|
||||
name="Test",
|
||||
is_truncated=False,
|
||||
)
|
||||
activity_meta = ActivityMetadata(activity=Activity.UPLOAD_TEXT)
|
||||
device_meta = DeviceMetadata(
|
||||
operating_system_specifications=OperatingSystemSpecifications(
|
||||
operating_system_platform="Windows", operating_system_version="10"
|
||||
)
|
||||
)
|
||||
integrated_app = IntegratedAppMetadata(name="App", version="1.0")
|
||||
location = PolicyLocation(data_type="microsoft.graph.policyLocationApplication", value="app-id")
|
||||
protected_app = ProtectedAppMetadata(name="Protected", version="1.0", application_location=location)
|
||||
|
||||
return ContentToProcess(
|
||||
content_entries=[metadata],
|
||||
activity_metadata=activity_meta,
|
||||
device_metadata=device_meta,
|
||||
integrated_app_metadata=integrated_app,
|
||||
protected_app_metadata=protected_app,
|
||||
)
|
||||
|
||||
return _create_content
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def process_content_request_factory(content_to_process_factory):
|
||||
"""Factory fixture to create ProcessContentRequest objects with test data."""
|
||||
|
||||
def _create_request(
|
||||
text: str = "Test", user_id: str = "user-123", tenant_id: str = "tenant-456"
|
||||
) -> ProcessContentRequest:
|
||||
content = content_to_process_factory(text)
|
||||
return ProcessContentRequest(
|
||||
content_to_process=content,
|
||||
user_id=user_id,
|
||||
tenant_id=tenant_id,
|
||||
)
|
||||
|
||||
return _create_request
|
||||
@@ -0,0 +1,149 @@
|
||||
# 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, Role
|
||||
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(role=Role.USER, text="Hello")], chat_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(role=Role.ASSISTANT, text="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 == 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
|
||||
msg = chat_context.result[0] # type: ignore[index]
|
||||
assert msg.role in ("system", Role.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(role=Role.ASSISTANT, text="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", Role.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(role=Role.USER, text="Hello")],
|
||||
chat_options=chat_options,
|
||||
is_streaming=True,
|
||||
)
|
||||
with patch.object(middleware._processor, "process_messages", return_value=False) 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."""
|
||||
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:
|
||||
# Create a mock result with messages attribute
|
||||
result = MagicMock()
|
||||
result.messages = [ChatMessage(role=Role.ASSISTANT, text="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
|
||||
@@ -0,0 +1,238 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
"""Tests for Purview client."""
|
||||
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
from azure.core.credentials import AccessToken
|
||||
|
||||
from agent_framework_purview import PurviewSettings
|
||||
from agent_framework_purview._client import PurviewClient
|
||||
from agent_framework_purview._exceptions import (
|
||||
PurviewAuthenticationError,
|
||||
PurviewRateLimitError,
|
||||
PurviewRequestError,
|
||||
PurviewServiceError,
|
||||
)
|
||||
from agent_framework_purview._models import (
|
||||
PolicyLocation,
|
||||
ProcessContentRequest,
|
||||
ProtectionScopesRequest,
|
||||
)
|
||||
|
||||
|
||||
class TestPurviewClient:
|
||||
"""Test PurviewClient functionality."""
|
||||
|
||||
@pytest.fixture
|
||||
def mock_credential(self) -> MagicMock:
|
||||
"""Create a mock async credential."""
|
||||
from azure.core.credentials_async import AsyncTokenCredential
|
||||
|
||||
credential = MagicMock(spec=AsyncTokenCredential)
|
||||
mock_token = AccessToken("fake-token", 9999999999)
|
||||
|
||||
async def mock_get_token(*args, **kwargs):
|
||||
return mock_token
|
||||
|
||||
credential.get_token = mock_get_token
|
||||
return credential
|
||||
|
||||
@pytest.fixture
|
||||
def settings(self) -> PurviewSettings:
|
||||
"""Create test settings."""
|
||||
return PurviewSettings(app_name="Test App", tenant_id="test-tenant", default_user_id="test-user")
|
||||
|
||||
@pytest.fixture
|
||||
async def client(self, mock_credential: MagicMock, settings: PurviewSettings) -> PurviewClient:
|
||||
"""Create a PurviewClient with mock credential."""
|
||||
client = PurviewClient(mock_credential, settings, timeout=10.0)
|
||||
yield client
|
||||
await client.close()
|
||||
|
||||
async def test_client_initialization(self, mock_credential: MagicMock, settings: PurviewSettings) -> None:
|
||||
"""Test PurviewClient initialization."""
|
||||
client = PurviewClient(mock_credential, settings)
|
||||
|
||||
assert client._credential == mock_credential
|
||||
assert client._settings == settings
|
||||
assert client._graph_uri == "https://graph.microsoft.com/v1.0"
|
||||
assert client._timeout == 10.0
|
||||
|
||||
await client.close()
|
||||
|
||||
async def test_get_token_async_credential(self, client: PurviewClient, mock_credential: MagicMock) -> None:
|
||||
"""Test _get_token with async credential."""
|
||||
token = await client._get_token(tenant_id="test-tenant")
|
||||
|
||||
assert token == "fake-token"
|
||||
|
||||
async def test_get_token_sync_credential(self, settings: PurviewSettings) -> None:
|
||||
"""Test _get_token with sync credential."""
|
||||
sync_credential = MagicMock()
|
||||
sync_credential.get_token = MagicMock(return_value=AccessToken("sync-token", 9999999999))
|
||||
|
||||
client = PurviewClient(sync_credential, settings)
|
||||
|
||||
with patch("asyncio.get_running_loop") as mock_loop:
|
||||
mock_executor = AsyncMock()
|
||||
mock_executor.return_value = AccessToken("sync-token", 9999999999)
|
||||
mock_loop.return_value.run_in_executor = mock_executor
|
||||
|
||||
token = await client._get_token(tenant_id="test-tenant")
|
||||
|
||||
assert token == "sync-token"
|
||||
|
||||
await client.close()
|
||||
|
||||
async def test_get_user_info_from_token(self, client: PurviewClient) -> None:
|
||||
"""Test get_user_info_from_token extracts user info."""
|
||||
import base64
|
||||
import json
|
||||
|
||||
payload = {"tid": "test-tenant", "oid": "test-user", "idtyp": "user"}
|
||||
payload_str = json.dumps(payload)
|
||||
payload_bytes = payload_str.encode("utf-8")
|
||||
payload_b64 = base64.urlsafe_b64encode(payload_bytes).decode("utf-8").rstrip("=")
|
||||
fake_token = f"header.{payload_b64}.signature"
|
||||
|
||||
with patch.object(client, "_get_token", return_value=fake_token):
|
||||
user_info = await client.get_user_info_from_token(tenant_id="test-tenant")
|
||||
|
||||
assert user_info["tenant_id"] == "test-tenant"
|
||||
assert user_info["user_id"] == "test-user"
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"status_code,exception_type",
|
||||
[
|
||||
(401, PurviewAuthenticationError),
|
||||
(403, PurviewAuthenticationError),
|
||||
(429, PurviewRateLimitError),
|
||||
(400, PurviewRequestError),
|
||||
(404, PurviewRequestError),
|
||||
(500, PurviewServiceError),
|
||||
(502, PurviewServiceError),
|
||||
],
|
||||
)
|
||||
async def test_post_error_handling(
|
||||
self, client: PurviewClient, content_to_process_factory, status_code: int, exception_type: type[Exception]
|
||||
) -> None:
|
||||
"""Test _post method handles different HTTP errors correctly."""
|
||||
from agent_framework_purview._models import ProcessContentResponse
|
||||
|
||||
content = content_to_process_factory()
|
||||
request = ProcessContentRequest(
|
||||
content_to_process=content,
|
||||
user_id="user-123",
|
||||
tenant_id="tenant-456",
|
||||
)
|
||||
|
||||
mock_response = MagicMock(spec=httpx.Response)
|
||||
mock_response.status_code = status_code
|
||||
mock_response.text = "Error message"
|
||||
mock_response.raise_for_status.side_effect = httpx.HTTPStatusError(
|
||||
"Error", request=MagicMock(), response=mock_response
|
||||
)
|
||||
|
||||
with patch.object(client._client, "post", return_value=mock_response), pytest.raises(exception_type):
|
||||
await client._post(
|
||||
"https://graph.microsoft.com/v1.0/test",
|
||||
request,
|
||||
ProcessContentResponse,
|
||||
"fake-token",
|
||||
)
|
||||
|
||||
async def test_process_content_success(
|
||||
self, client: PurviewClient, content_to_process_factory, mock_credential: MagicMock
|
||||
) -> None:
|
||||
"""Test process_content method success path."""
|
||||
content = content_to_process_factory("Test message")
|
||||
request = ProcessContentRequest(
|
||||
content_to_process=content,
|
||||
user_id="user-123",
|
||||
tenant_id="tenant-456",
|
||||
)
|
||||
|
||||
mock_response = MagicMock(spec=httpx.Response)
|
||||
mock_response.status_code = 200
|
||||
mock_response.json.return_value = {"id": "response-123", "protectionScopeState": "notModified"}
|
||||
|
||||
with patch.object(client._client, "post", return_value=mock_response):
|
||||
response = await client.process_content(request)
|
||||
|
||||
assert response.id == "response-123"
|
||||
assert response.protection_scope_state == "notModified"
|
||||
|
||||
async def test_get_protection_scopes_success(self, client: PurviewClient) -> None:
|
||||
"""Test get_protection_scopes method success path."""
|
||||
location = PolicyLocation(**{"@odata.type": "microsoft.graph.policyLocationApplication", "value": "app-id"})
|
||||
request = ProtectionScopesRequest(
|
||||
user_id="user-123", tenant_id="tenant-456", locations=[location], correlation_id="corr-789"
|
||||
)
|
||||
|
||||
mock_response = MagicMock(spec=httpx.Response)
|
||||
mock_response.status_code = 200
|
||||
mock_response.json.return_value = {"scopeIdentifier": "scope-123", "value": []}
|
||||
|
||||
with patch.object(client._client, "post", return_value=mock_response):
|
||||
response = await client.get_protection_scopes(request)
|
||||
|
||||
assert response.scope_identifier == "scope-123"
|
||||
assert response.scopes == []
|
||||
|
||||
async def test_client_close(self, mock_credential: AsyncMock, settings: PurviewSettings) -> None:
|
||||
"""Test client properly closes HTTP client."""
|
||||
client = PurviewClient(mock_credential, settings)
|
||||
|
||||
with patch.object(client._client, "aclose", new_callable=AsyncMock) as mock_close:
|
||||
await client.close()
|
||||
mock_close.assert_called_once()
|
||||
|
||||
async def test_invalid_jwt_token_format(self, client: PurviewClient) -> None:
|
||||
"""Test that invalid JWT token format raises ValueError."""
|
||||
with pytest.raises(ValueError, match="Invalid JWT token format"):
|
||||
client._extract_token_info("invalid-token-without-dots")
|
||||
|
||||
async def test_rate_limit_error(self, client: PurviewClient) -> None:
|
||||
"""Test that 429 status code raises PurviewRateLimitError."""
|
||||
request = ProcessContentRequest(
|
||||
user_id="test-user",
|
||||
tenant_id="test-tenant",
|
||||
content_to_process=[],
|
||||
correlation_id="test-correlation-id",
|
||||
)
|
||||
|
||||
with (
|
||||
patch.object(client, "_get_token", return_value="fake-token"),
|
||||
patch.object(
|
||||
client._client,
|
||||
"post",
|
||||
return_value=httpx.Response(429, text="Rate limited", request=httpx.Request("POST", "http://test")),
|
||||
),
|
||||
pytest.raises(PurviewRateLimitError, match="Rate limited"),
|
||||
):
|
||||
await client.process_content(request)
|
||||
|
||||
async def test_generic_request_error(self, client: PurviewClient) -> None:
|
||||
"""Test that non-200/201/202 status codes raise PurviewRequestError."""
|
||||
request = ProcessContentRequest(
|
||||
user_id="test-user",
|
||||
tenant_id="test-tenant",
|
||||
content_to_process=[],
|
||||
correlation_id="test-correlation-id",
|
||||
)
|
||||
|
||||
with (
|
||||
patch.object(client, "_get_token", return_value="fake-token"),
|
||||
patch.object(
|
||||
client._client,
|
||||
"post",
|
||||
return_value=httpx.Response(
|
||||
500, text="Internal server error", request=httpx.Request("POST", "http://test")
|
||||
),
|
||||
),
|
||||
pytest.raises(PurviewRequestError, match="Purview request failed"),
|
||||
):
|
||||
await client.process_content(request)
|
||||
@@ -0,0 +1,38 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
"""Tests for Purview exceptions."""
|
||||
|
||||
from agent_framework_purview import (
|
||||
PurviewAuthenticationError,
|
||||
PurviewRateLimitError,
|
||||
PurviewRequestError,
|
||||
PurviewServiceError,
|
||||
)
|
||||
|
||||
|
||||
class TestPurviewExceptions:
|
||||
"""Test custom Purview exception classes."""
|
||||
|
||||
def test_purview_service_error(self) -> None:
|
||||
"""Test PurviewServiceError base exception."""
|
||||
error = PurviewServiceError("Service error occurred")
|
||||
assert str(error) == "Service error occurred"
|
||||
assert isinstance(error, Exception)
|
||||
|
||||
def test_purview_authentication_error(self) -> None:
|
||||
"""Test PurviewAuthenticationError exception."""
|
||||
error = PurviewAuthenticationError("Authentication failed")
|
||||
assert str(error) == "Authentication failed"
|
||||
assert isinstance(error, PurviewServiceError)
|
||||
|
||||
def test_purview_rate_limit_error(self) -> None:
|
||||
"""Test PurviewRateLimitError exception."""
|
||||
error = PurviewRateLimitError("Rate limit exceeded")
|
||||
assert str(error) == "Rate limit exceeded"
|
||||
assert isinstance(error, PurviewServiceError)
|
||||
|
||||
def test_purview_request_error(self) -> None:
|
||||
"""Test PurviewRequestError exception."""
|
||||
error = PurviewRequestError("Request failed")
|
||||
assert str(error) == "Request failed"
|
||||
assert isinstance(error, PurviewServiceError)
|
||||
@@ -0,0 +1,201 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
"""Tests for Purview middleware."""
|
||||
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import pytest
|
||||
from agent_framework import AgentRunContext, AgentRunResponse, ChatMessage, Role
|
||||
from azure.core.credentials import AccessToken
|
||||
|
||||
from agent_framework_purview import PurviewPolicyMiddleware, PurviewSettings
|
||||
|
||||
|
||||
class TestPurviewPolicyMiddleware:
|
||||
"""Test PurviewPolicyMiddleware functionality."""
|
||||
|
||||
@pytest.fixture
|
||||
def mock_credential(self) -> AsyncMock:
|
||||
"""Create a mock async credential."""
|
||||
credential = AsyncMock()
|
||||
credential.get_token = AsyncMock(return_value=AccessToken("fake-token", 9999999999))
|
||||
return credential
|
||||
|
||||
@pytest.fixture
|
||||
def settings(self) -> PurviewSettings:
|
||||
"""Create test settings."""
|
||||
return PurviewSettings(app_name="Test App", tenant_id="test-tenant")
|
||||
|
||||
@pytest.fixture
|
||||
def middleware(self, mock_credential: AsyncMock, settings: PurviewSettings) -> PurviewPolicyMiddleware:
|
||||
"""Create PurviewPolicyMiddleware instance."""
|
||||
return PurviewPolicyMiddleware(mock_credential, settings)
|
||||
|
||||
@pytest.fixture
|
||||
def mock_agent(self) -> MagicMock:
|
||||
"""Create a mock agent."""
|
||||
agent = MagicMock()
|
||||
agent.name = "test-agent"
|
||||
return agent
|
||||
|
||||
def test_middleware_initialization(self, mock_credential: AsyncMock, settings: PurviewSettings) -> None:
|
||||
"""Test PurviewPolicyMiddleware initialization."""
|
||||
middleware = PurviewPolicyMiddleware(mock_credential, settings)
|
||||
|
||||
assert middleware._client is not None
|
||||
assert middleware._processor is not None
|
||||
|
||||
async def test_middleware_allows_clean_prompt(
|
||||
self, middleware: PurviewPolicyMiddleware, mock_agent: MagicMock
|
||||
) -> None:
|
||||
"""Test middleware allows prompt that passes policy check."""
|
||||
context = AgentRunContext(agent=mock_agent, messages=[ChatMessage(role=Role.USER, text="Hello, how are you?")])
|
||||
|
||||
with patch.object(middleware._processor, "process_messages", return_value=(False, "user-123")):
|
||||
next_called = False
|
||||
|
||||
async def mock_next(ctx: AgentRunContext) -> None:
|
||||
nonlocal next_called
|
||||
next_called = True
|
||||
ctx.result = AgentRunResponse(messages=[ChatMessage(role=Role.ASSISTANT, text="I'm good, thanks!")])
|
||||
|
||||
await middleware.process(context, mock_next)
|
||||
|
||||
assert next_called
|
||||
assert context.result is not None
|
||||
assert not context.terminate
|
||||
|
||||
async def test_middleware_blocks_prompt_on_policy_violation(
|
||||
self, middleware: PurviewPolicyMiddleware, mock_agent: MagicMock
|
||||
) -> None:
|
||||
"""Test middleware blocks prompt that violates policy."""
|
||||
context = AgentRunContext(
|
||||
agent=mock_agent, messages=[ChatMessage(role=Role.USER, text="Sensitive information")]
|
||||
)
|
||||
|
||||
with patch.object(middleware._processor, "process_messages", return_value=(True, "user-123")):
|
||||
next_called = False
|
||||
|
||||
async def mock_next(ctx: AgentRunContext) -> None:
|
||||
nonlocal next_called
|
||||
next_called = True
|
||||
|
||||
await middleware.process(context, mock_next)
|
||||
|
||||
assert not next_called
|
||||
assert context.result is not None
|
||||
assert context.terminate
|
||||
assert len(context.result.messages) == 1
|
||||
assert context.result.messages[0].role == Role.SYSTEM
|
||||
assert "blocked by policy" in context.result.messages[0].text.lower()
|
||||
|
||||
async def test_middleware_checks_response(self, middleware: PurviewPolicyMiddleware, mock_agent: MagicMock) -> None:
|
||||
"""Test middleware checks agent response for policy violations."""
|
||||
context = AgentRunContext(agent=mock_agent, messages=[ChatMessage(role=Role.USER, text="Hello")])
|
||||
|
||||
call_count = 0
|
||||
|
||||
async def mock_process_messages(messages, activity, user_id=None):
|
||||
nonlocal call_count
|
||||
call_count += 1
|
||||
should_block = call_count != 1
|
||||
return (should_block, "user-123")
|
||||
|
||||
with patch.object(middleware._processor, "process_messages", side_effect=mock_process_messages):
|
||||
|
||||
async def mock_next(ctx: AgentRunContext) -> None:
|
||||
ctx.result = AgentRunResponse(
|
||||
messages=[ChatMessage(role=Role.ASSISTANT, text="Here's some sensitive information")]
|
||||
)
|
||||
|
||||
await middleware.process(context, mock_next)
|
||||
|
||||
assert call_count == 2
|
||||
assert context.result is not None
|
||||
assert len(context.result.messages) == 1
|
||||
assert context.result.messages[0].role == Role.SYSTEM
|
||||
assert "blocked by policy" in context.result.messages[0].text.lower()
|
||||
|
||||
async def test_middleware_handles_result_without_messages(
|
||||
self, middleware: PurviewPolicyMiddleware, mock_agent: MagicMock
|
||||
) -> None:
|
||||
"""Test middleware handles result that doesn't have messages attribute."""
|
||||
context = AgentRunContext(agent=mock_agent, messages=[ChatMessage(role=Role.USER, text="Hello")])
|
||||
|
||||
with patch.object(middleware._processor, "process_messages", return_value=(False, "user-123")):
|
||||
|
||||
async def mock_next(ctx: AgentRunContext) -> None:
|
||||
ctx.result = "Some non-standard result"
|
||||
|
||||
await middleware.process(context, mock_next)
|
||||
|
||||
assert context.result == "Some non-standard result"
|
||||
|
||||
async def test_middleware_processor_receives_correct_activity(
|
||||
self, middleware: PurviewPolicyMiddleware, mock_agent: MagicMock
|
||||
) -> None:
|
||||
"""Test middleware passes correct activity type to processor."""
|
||||
from agent_framework_purview._models import Activity
|
||||
|
||||
context = AgentRunContext(agent=mock_agent, messages=[ChatMessage(role=Role.USER, text="Test")])
|
||||
|
||||
with patch.object(middleware._processor, "process_messages", return_value=(False, "user-123")) as mock_process:
|
||||
|
||||
async def mock_next(ctx: AgentRunContext) -> None:
|
||||
ctx.result = AgentRunResponse(messages=[ChatMessage(role=Role.ASSISTANT, text="Response")])
|
||||
|
||||
await middleware.process(context, mock_next)
|
||||
|
||||
assert mock_process.call_count == 2
|
||||
for call in mock_process.call_args_list:
|
||||
assert call[0][1] == Activity.UPLOAD_TEXT
|
||||
|
||||
async def test_middleware_handles_pre_check_exception(
|
||||
self, middleware: PurviewPolicyMiddleware, mock_agent: MagicMock
|
||||
) -> None:
|
||||
"""Test that exceptions in pre-check are logged but don't stop processing."""
|
||||
context = AgentRunContext(agent=mock_agent, messages=[ChatMessage(role=Role.USER, text="Test")])
|
||||
|
||||
with patch.object(
|
||||
middleware._processor, "process_messages", side_effect=Exception("Pre-check error")
|
||||
) as mock_process:
|
||||
|
||||
async def mock_next(ctx: AgentRunContext) -> None:
|
||||
ctx.result = AgentRunResponse(messages=[ChatMessage(role=Role.ASSISTANT, text="Response")])
|
||||
|
||||
await middleware.process(context, mock_next)
|
||||
|
||||
# Should have been called twice (pre-check raises, then post-check also raises)
|
||||
assert mock_process.call_count == 2
|
||||
# Context should not be terminated
|
||||
assert not context.terminate
|
||||
# Result should be set by mock_next
|
||||
assert context.result is not None
|
||||
|
||||
async def test_middleware_handles_post_check_exception(
|
||||
self, middleware: PurviewPolicyMiddleware, mock_agent: MagicMock
|
||||
) -> None:
|
||||
"""Test that exceptions in post-check are logged but don't affect result."""
|
||||
context = AgentRunContext(agent=mock_agent, messages=[ChatMessage(role=Role.USER, text="Test")])
|
||||
|
||||
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: AgentRunContext) -> None:
|
||||
ctx.result = AgentRunResponse(messages=[ChatMessage(role=Role.ASSISTANT, text="Response")])
|
||||
|
||||
await middleware.process(context, mock_next)
|
||||
|
||||
# Should have been called twice (pre and post)
|
||||
assert call_count == 2
|
||||
# Result should still be set
|
||||
assert context.result is not None
|
||||
assert hasattr(context.result, "messages")
|
||||
@@ -0,0 +1,246 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
"""Tests for Purview models and serialization."""
|
||||
|
||||
from agent_framework_purview._models import (
|
||||
Activity,
|
||||
ActivityMetadata,
|
||||
ContentToProcess,
|
||||
DeviceMetadata,
|
||||
IntegratedAppMetadata,
|
||||
OperatingSystemSpecifications,
|
||||
PolicyLocation,
|
||||
ProcessContentRequest,
|
||||
ProcessContentResponse,
|
||||
ProcessConversationMetadata,
|
||||
ProtectedAppMetadata,
|
||||
ProtectionScopeActivities,
|
||||
ProtectionScopesRequest,
|
||||
ProtectionScopesResponse,
|
||||
PurviewTextContent,
|
||||
deserialize_flag,
|
||||
serialize_flag,
|
||||
)
|
||||
|
||||
|
||||
class TestFlagOperations:
|
||||
"""Test flag serialization and deserialization operations."""
|
||||
|
||||
def test_protection_scope_activities_flag_combination(self) -> None:
|
||||
"""Test combining flags."""
|
||||
combined = ProtectionScopeActivities.UPLOAD_TEXT | ProtectionScopeActivities.UPLOAD_FILE
|
||||
assert combined.value == 3
|
||||
assert ProtectionScopeActivities.UPLOAD_TEXT in combined
|
||||
assert ProtectionScopeActivities.UPLOAD_FILE in combined
|
||||
|
||||
def test_deserialize_flag_with_string(self) -> None:
|
||||
"""Test deserializing flag from comma-separated string."""
|
||||
mapping = {
|
||||
"uploadText": ProtectionScopeActivities.UPLOAD_TEXT,
|
||||
"uploadFile": ProtectionScopeActivities.UPLOAD_FILE,
|
||||
}
|
||||
|
||||
result = deserialize_flag("uploadText,uploadFile", mapping, ProtectionScopeActivities)
|
||||
assert result is not None
|
||||
assert ProtectionScopeActivities.UPLOAD_TEXT in result
|
||||
assert ProtectionScopeActivities.UPLOAD_FILE in result
|
||||
|
||||
def test_deserialize_flag_with_none(self) -> None:
|
||||
"""Test deserializing None returns None."""
|
||||
mapping = {"uploadText": ProtectionScopeActivities.UPLOAD_TEXT}
|
||||
result = deserialize_flag(None, mapping, ProtectionScopeActivities)
|
||||
assert result is None
|
||||
|
||||
def test_serialize_flag_with_none(self) -> None:
|
||||
"""Test serializing None returns None."""
|
||||
result = serialize_flag(None, [])
|
||||
assert result is None
|
||||
|
||||
def test_serialize_flag_with_values(self) -> None:
|
||||
"""Test serializing flag with values."""
|
||||
flag = ProtectionScopeActivities.UPLOAD_TEXT | ProtectionScopeActivities.UPLOAD_FILE
|
||||
ordered = [
|
||||
("uploadText", ProtectionScopeActivities.UPLOAD_TEXT),
|
||||
("uploadFile", ProtectionScopeActivities.UPLOAD_FILE),
|
||||
]
|
||||
result = serialize_flag(flag, ordered)
|
||||
assert result == "uploadText,uploadFile"
|
||||
|
||||
|
||||
class TestComplexModels:
|
||||
"""Test complex models with nested structures."""
|
||||
|
||||
def test_content_to_process_with_nested_structures(self) -> None:
|
||||
"""Test ContentToProcess with all nested structures."""
|
||||
text_content = PurviewTextContent(data="Test")
|
||||
metadata = ProcessConversationMetadata(
|
||||
identifier="msg-1",
|
||||
content=text_content,
|
||||
name="Test",
|
||||
is_truncated=False,
|
||||
)
|
||||
|
||||
activity_meta = ActivityMetadata(activity=Activity.UPLOAD_TEXT)
|
||||
device_meta = DeviceMetadata(
|
||||
operating_system_specifications=OperatingSystemSpecifications(
|
||||
operating_system_platform="Windows", operating_system_version="10"
|
||||
)
|
||||
)
|
||||
integrated_app = IntegratedAppMetadata(name="App", version="1.0")
|
||||
location = PolicyLocation(data_type="microsoft.graph.policyLocationApplication", value="app-id")
|
||||
protected_app = ProtectedAppMetadata(name="Protected", version="1.0", application_location=location)
|
||||
|
||||
content = ContentToProcess(
|
||||
content_entries=[metadata],
|
||||
activity_metadata=activity_meta,
|
||||
device_metadata=device_meta,
|
||||
integrated_app_metadata=integrated_app,
|
||||
protected_app_metadata=protected_app,
|
||||
)
|
||||
|
||||
assert len(content.content_entries) == 1
|
||||
assert content.activity_metadata.activity == Activity.UPLOAD_TEXT
|
||||
assert content.device_metadata.operating_system_specifications.operating_system_platform == "Windows"
|
||||
assert content.integrated_app_metadata.name == "App"
|
||||
assert content.protected_app_metadata.name == "Protected"
|
||||
|
||||
|
||||
class TestRequestResponseSerialization:
|
||||
"""Test request/response serialization with aliases."""
|
||||
|
||||
def test_protection_scopes_request_serialization(self) -> None:
|
||||
"""Test ProtectionScopesRequest serializes activities correctly."""
|
||||
location = PolicyLocation(data_type="microsoft.graph.policyLocationApplication", value="app-id")
|
||||
|
||||
request = ProtectionScopesRequest(
|
||||
user_id="user-123",
|
||||
tenant_id="tenant-456",
|
||||
activities=ProtectionScopeActivities.UPLOAD_TEXT | ProtectionScopeActivities.UPLOAD_FILE,
|
||||
locations=[location],
|
||||
)
|
||||
|
||||
dumped = request.model_dump(by_alias=True, exclude_none=True, mode="json")
|
||||
|
||||
assert "activities" in dumped
|
||||
assert isinstance(dumped["activities"], str)
|
||||
assert "uploadText" in dumped["activities"]
|
||||
|
||||
|
||||
class TestModelDeserialization:
|
||||
"""Test model deserialization from API responses."""
|
||||
|
||||
def test_protection_scopes_response_deserialization(self) -> None:
|
||||
"""Test ProtectionScopesResponse deserializes 'value' to 'scopes'."""
|
||||
api_data = {
|
||||
"scopeIdentifier": "scope-123",
|
||||
"value": [
|
||||
{
|
||||
"activities": "uploadText,downloadText",
|
||||
"locations": [{"@odata.type": "location.type", "value": "/path"}],
|
||||
"policyActions": [{"action": "warn", "restrictionAction": "blockAccess"}],
|
||||
"executionMode": "evaluateInline",
|
||||
}
|
||||
],
|
||||
}
|
||||
|
||||
response = ProtectionScopesResponse.model_validate(api_data)
|
||||
|
||||
assert response.scope_identifier == "scope-123"
|
||||
assert response.scopes is not None
|
||||
assert len(response.scopes) == 1
|
||||
assert response.scopes[0].execution_mode == "evaluateInline"
|
||||
|
||||
def test_process_content_response_deserialization(self) -> None:
|
||||
"""Test ProcessContentResponse deserializes aliased fields correctly."""
|
||||
api_data = {
|
||||
"id": "response-123",
|
||||
"protectionScopeState": "blocked",
|
||||
"policyActions": [{"action": "block", "restrictionAction": "blockAccess"}],
|
||||
}
|
||||
|
||||
response = ProcessContentResponse.model_validate(api_data)
|
||||
|
||||
assert response.id == "response-123"
|
||||
assert response.protection_scope_state == "blocked"
|
||||
assert len(response.policy_actions) == 1
|
||||
|
||||
def test_content_serialization_uses_aliases(self) -> None:
|
||||
"""Test ContentToProcess serializes with camelCase aliases."""
|
||||
text_content = PurviewTextContent(data="Test")
|
||||
metadata = ProcessConversationMetadata(
|
||||
identifier="msg-1",
|
||||
content=text_content,
|
||||
name="Test",
|
||||
is_truncated=False,
|
||||
)
|
||||
|
||||
activity_meta = ActivityMetadata(activity=Activity.UPLOAD_TEXT)
|
||||
device_meta = DeviceMetadata(
|
||||
operating_system_specifications=OperatingSystemSpecifications(
|
||||
operating_system_platform="Windows", operating_system_version="10"
|
||||
)
|
||||
)
|
||||
integrated_app = IntegratedAppMetadata(name="App", version="1.0")
|
||||
location = PolicyLocation(data_type="microsoft.graph.policyLocationApplication", value="app-id")
|
||||
protected_app = ProtectedAppMetadata(name="Protected", version="1.0", application_location=location)
|
||||
|
||||
content = ContentToProcess(
|
||||
content_entries=[metadata],
|
||||
activity_metadata=activity_meta,
|
||||
device_metadata=device_meta,
|
||||
integrated_app_metadata=integrated_app,
|
||||
protected_app_metadata=protected_app,
|
||||
)
|
||||
|
||||
dumped = content.model_dump(by_alias=True, exclude_none=True, mode="json")
|
||||
|
||||
assert "contentEntries" in dumped
|
||||
assert "activityMetadata" in dumped
|
||||
assert "deviceMetadata" in dumped
|
||||
assert "integratedAppMetadata" in dumped
|
||||
assert "protectedAppMetadata" in dumped
|
||||
|
||||
def test_process_content_request_excludes_private_fields(self) -> None:
|
||||
"""Test ProcessContentRequest excludes private fields when serializing."""
|
||||
text_content = PurviewTextContent(data="Test")
|
||||
metadata = ProcessConversationMetadata(
|
||||
identifier="msg-1",
|
||||
content=text_content,
|
||||
name="Test",
|
||||
is_truncated=False,
|
||||
)
|
||||
|
||||
activity_meta = ActivityMetadata(activity=Activity.UPLOAD_TEXT)
|
||||
device_meta = DeviceMetadata(
|
||||
operating_system_specifications=OperatingSystemSpecifications(
|
||||
operating_system_platform="Windows", operating_system_version="10"
|
||||
)
|
||||
)
|
||||
integrated_app = IntegratedAppMetadata(name="App", version="1.0")
|
||||
location = PolicyLocation(data_type="microsoft.graph.policyLocationApplication", value="app-id")
|
||||
protected_app = ProtectedAppMetadata(name="Protected", version="1.0", application_location=location)
|
||||
|
||||
content = ContentToProcess(
|
||||
content_entries=[metadata],
|
||||
activity_metadata=activity_meta,
|
||||
device_metadata=device_meta,
|
||||
integrated_app_metadata=integrated_app,
|
||||
protected_app_metadata=protected_app,
|
||||
)
|
||||
|
||||
request = ProcessContentRequest(
|
||||
content_to_process=content,
|
||||
user_id="user-123",
|
||||
tenant_id="tenant-456",
|
||||
correlation_id="corr-789",
|
||||
)
|
||||
|
||||
dumped = request.model_dump(by_alias=True, exclude_none=True, mode="json")
|
||||
|
||||
# Check that excluded fields are not present
|
||||
assert "user_id" not in dumped
|
||||
assert "tenant_id" not in dumped
|
||||
assert "correlation_id" not in dumped
|
||||
|
||||
# Check that content is present
|
||||
assert "contentToProcess" in dumped
|
||||
@@ -0,0 +1,369 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
"""Tests for Purview processor."""
|
||||
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import pytest
|
||||
from agent_framework import ChatMessage, Role
|
||||
|
||||
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=Role.USER, text="Hello"),
|
||||
ChatMessage(role=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=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=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=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=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 = 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 = 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 and process_content."""
|
||||
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()
|
||||
mock_client.process_content.assert_not_called()
|
||||
mock_client.send_content_activities.assert_called_once()
|
||||
assert response.id is None
|
||||
|
||||
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=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=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=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 content activities are sent when scopes don't apply."""
|
||||
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
|
||||
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()
|
||||
mock_client.send_content_activities.assert_called_once()
|
||||
# When content activities succeed, response has no errors (processing_errors can be None or empty)
|
||||
assert response.processing_errors is None or response.processing_errors == []
|
||||
|
||||
async def test_process_content_handles_activities_error(
|
||||
self, mock_client: AsyncMock, process_content_request_factory
|
||||
) -> None:
|
||||
"""Test error handling when content activities fail."""
|
||||
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
|
||||
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)
|
||||
|
||||
assert len(response.processing_errors) == 1
|
||||
assert response.processing_errors[0].message == "Test error message"
|
||||
@@ -0,0 +1,85 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
"""Tests for Purview settings."""
|
||||
|
||||
import pytest
|
||||
|
||||
from agent_framework_purview import PurviewAppLocation, PurviewLocationType, PurviewSettings
|
||||
|
||||
|
||||
class TestPurviewSettings:
|
||||
"""Test PurviewSettings configuration."""
|
||||
|
||||
def test_settings_defaults(self) -> None:
|
||||
"""Test PurviewSettings with default values."""
|
||||
settings = PurviewSettings(app_name="Test App")
|
||||
|
||||
assert settings.app_name == "Test App"
|
||||
assert settings.graph_base_uri == "https://graph.microsoft.com/v1.0/"
|
||||
assert settings.tenant_id is None
|
||||
assert settings.purview_app_location is None
|
||||
assert settings.process_inline is False
|
||||
|
||||
def test_settings_with_custom_values(self) -> None:
|
||||
"""Test PurviewSettings with custom values."""
|
||||
app_location = PurviewAppLocation(location_type=PurviewLocationType.APPLICATION, location_value="app-123")
|
||||
|
||||
settings = PurviewSettings(
|
||||
app_name="Test App",
|
||||
graph_base_uri="https://graph.microsoft-ppe.com",
|
||||
tenant_id="test-tenant-id",
|
||||
process_inline=True,
|
||||
purview_app_location=app_location,
|
||||
)
|
||||
|
||||
assert settings.graph_base_uri == "https://graph.microsoft-ppe.com"
|
||||
assert settings.tenant_id == "test-tenant-id"
|
||||
assert settings.process_inline is True
|
||||
assert settings.purview_app_location.location_value == "app-123"
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"graph_uri,expected_scope",
|
||||
[
|
||||
("https://graph.microsoft.com/v1.0/", "https://graph.microsoft.com/.default"),
|
||||
("https://graph.microsoft-ppe.com/v1.0/", "https://graph.microsoft-ppe.com/.default"),
|
||||
],
|
||||
)
|
||||
def test_get_scopes(self, graph_uri: str, expected_scope: str) -> None:
|
||||
"""Test get_scopes returns correct scope for different URIs."""
|
||||
settings = PurviewSettings(app_name="Test App", graph_base_uri=graph_uri)
|
||||
scopes = settings.get_scopes()
|
||||
|
||||
assert len(scopes) == 1
|
||||
assert expected_scope in scopes
|
||||
|
||||
|
||||
class TestPurviewAppLocation:
|
||||
"""Test PurviewAppLocation configuration."""
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"location_type,location_value,expected_odata_type",
|
||||
[
|
||||
(PurviewLocationType.APPLICATION, "app-123", "microsoft.graph.policyLocationApplication"),
|
||||
(PurviewLocationType.URI, "https://example.com", "microsoft.graph.policyLocationUrl"),
|
||||
(PurviewLocationType.DOMAIN, "example.com", "microsoft.graph.policyLocationDomain"),
|
||||
],
|
||||
)
|
||||
def test_get_policy_location(
|
||||
self, location_type: PurviewLocationType, location_value: str, expected_odata_type: str
|
||||
) -> None:
|
||||
"""Test get_policy_location returns correct structure for all location types."""
|
||||
location = PurviewAppLocation(location_type=location_type, location_value=location_value)
|
||||
policy_location = location.get_policy_location()
|
||||
|
||||
assert policy_location["@odata.type"] == expected_odata_type
|
||||
assert policy_location["value"] == location_value
|
||||
|
||||
|
||||
class TestPurviewLocationType:
|
||||
"""Test PurviewLocationType enum."""
|
||||
|
||||
def test_location_type_values(self) -> None:
|
||||
"""Test PurviewLocationType enum has expected values."""
|
||||
assert PurviewLocationType.APPLICATION == "application"
|
||||
assert PurviewLocationType.URI == "uri"
|
||||
assert PurviewLocationType.DOMAIN == "domain"
|
||||
Reference in New Issue
Block a user