Files
wehub-resource-sync db620d33df
CodeQL / Analyze (csharp) (push) Waiting to run
CodeQL / Analyze (python) (push) Waiting to run
dotnet-build-and-test / dotnet-test-functions (push) Has been cancelled
dotnet-build-and-test / paths-filter (push) Has been cancelled
dotnet-build-and-test / dotnet-build (Debug, windows-latest, net9.0) (push) Has been cancelled
dotnet-build-and-test / dotnet-build (Release, ubuntu-latest, net10.0) (push) Has been cancelled
dotnet-build-and-test / dotnet-build (Release, ubuntu-latest, net8.0) (push) Has been cancelled
dotnet-build-and-test / dotnet-build (Release, windows-latest, net472) (push) Has been cancelled
dotnet-build-and-test / dotnet-test (Release, integration, true, ubuntu-latest, net10.0) (push) Has been cancelled
dotnet-build-and-test / dotnet-test (Release, integration, true, windows-latest, net472) (push) Has been cancelled
dotnet-build-and-test / dotnet-foundry-hosted-it (push) Has been cancelled
dotnet-build-and-test / dotnet-build-and-test-check (push) Has been cancelled
dotnet-build-and-test / Integration Test Report (push) Has been cancelled
chore: import upstream snapshot with attribution
2026-07-13 13:39:25 +08:00

1827 lines
72 KiB
Python

# Copyright (c) Microsoft. All rights reserved.
from __future__ import annotations
import asyncio
import base64
import json
import re
from typing import Any, cast
from unittest.mock import AsyncMock, MagicMock
from agent_framework import Content, Message, SessionContext
from agent_framework._sessions import AgentSession
from azure.ai.contentunderstanding.models import AnalysisResult
from agent_framework_azure_contentunderstanding import (
ContentUnderstandingContextProvider,
DocumentStatus,
)
from agent_framework_azure_contentunderstanding._detection import SUPPORTED_MEDIA_TYPES, derive_doc_key
# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------
_SAMPLE_PDF_BYTES = b"%PDF-1.4 fake content for testing"
def _make_mock_poller(result: AnalysisResult) -> AsyncMock:
"""Create a mock poller that returns the given result immediately."""
poller = AsyncMock()
poller.result = AsyncMock(return_value=result)
return poller
def _make_slow_poller(result: AnalysisResult, delay: float = 10.0) -> MagicMock:
"""Create a mock poller that simulates a timeout then eventually returns."""
poller = MagicMock()
async def slow_result() -> AnalysisResult:
await asyncio.sleep(delay)
return result
poller.result = slow_result
poller.continuation_token = MagicMock(return_value="mock_slow_continuation_token")
return poller
def _make_failing_poller(error: Exception) -> AsyncMock:
"""Create a mock poller that raises an exception."""
poller = AsyncMock()
poller.result = AsyncMock(side_effect=error)
return poller
def _make_data_uri(data: bytes, media_type: str) -> str:
return f"data:{media_type};base64,{base64.b64encode(data).decode('ascii')}"
def _make_content_from_data(data: bytes, media_type: str, filename: str | None = None) -> Content:
props = {"filename": filename} if filename else None
return Content.from_data(data, media_type, additional_properties=props)
def _make_context(messages: list[Message]) -> SessionContext:
return SessionContext(input_messages=messages)
def _make_provider(
mock_client: AsyncMock | None = None,
**kwargs: Any,
) -> ContentUnderstandingContextProvider:
provider = ContentUnderstandingContextProvider(
endpoint="https://test.cognitiveservices.azure.com/",
credential=AsyncMock(),
**kwargs,
)
if mock_client:
provider._client = mock_client # type: ignore[assignment]
return provider
def _make_mock_agent() -> MagicMock:
return MagicMock()
# ===========================================================================
# Test Classes
# ===========================================================================
class TestInit:
def test_default_values(self) -> None:
provider = ContentUnderstandingContextProvider(
endpoint="https://test.cognitiveservices.azure.com/",
credential=AsyncMock(),
)
assert provider.analyzer_id is None
assert provider.max_wait == 5.0
assert provider.output_sections == ["markdown", "fields"]
assert provider.source_id == "azure_contentunderstanding"
def test_custom_values(self) -> None:
provider = ContentUnderstandingContextProvider(
endpoint="https://custom.cognitiveservices.azure.com/",
credential=AsyncMock(),
analyzer_id="prebuilt-invoice",
max_wait=10.0,
output_sections=["markdown"],
source_id="custom_cu",
)
assert provider.analyzer_id == "prebuilt-invoice"
assert provider.max_wait == 10.0
assert provider.output_sections == ["markdown"]
assert provider.source_id == "custom_cu"
def test_max_wait_none(self) -> None:
provider = ContentUnderstandingContextProvider(
endpoint="https://test.cognitiveservices.azure.com/",
credential=AsyncMock(),
max_wait=None,
)
assert provider.max_wait is None
def test_endpoint_from_env_var(self, monkeypatch: Any) -> None:
"""Endpoint can be loaded from AZURE_CONTENTUNDERSTANDING_ENDPOINT env var."""
monkeypatch.setenv(
"AZURE_CONTENTUNDERSTANDING_ENDPOINT",
"https://env-test.cognitiveservices.azure.com/",
)
provider = ContentUnderstandingContextProvider(credential=AsyncMock())
assert provider._endpoint == "https://env-test.cognitiveservices.azure.com/"
def test_explicit_endpoint_overrides_env_var(self, monkeypatch: Any) -> None:
"""Explicit endpoint kwarg takes priority over env var."""
monkeypatch.setenv(
"AZURE_CONTENTUNDERSTANDING_ENDPOINT",
"https://env-test.cognitiveservices.azure.com/",
)
provider = ContentUnderstandingContextProvider(
endpoint="https://explicit.cognitiveservices.azure.com/",
credential=AsyncMock(),
)
assert provider._endpoint == "https://explicit.cognitiveservices.azure.com/"
def test_missing_endpoint_raises(self) -> None:
"""Missing endpoint (no kwarg, no env var) raises an error."""
# Clear env var to ensure load_settings raises
import os
import pytest as _pytest
from agent_framework.exceptions import SettingNotFoundError
env_key = "AZURE_CONTENTUNDERSTANDING_ENDPOINT"
old_val = os.environ.pop(env_key, None)
try:
with _pytest.raises(SettingNotFoundError, match="endpoint"):
ContentUnderstandingContextProvider(credential=AsyncMock())
finally:
if old_val is not None:
os.environ[env_key] = old_val
def test_missing_credential_raises(self) -> None:
"""Missing credential raises ValueError."""
import pytest as _pytest
with _pytest.raises(ValueError, match="credential is required"):
ContentUnderstandingContextProvider(
endpoint="https://test.cognitiveservices.azure.com/",
)
class TestAsyncContextManager:
async def test_aenter_returns_self(self) -> None:
provider = ContentUnderstandingContextProvider(
endpoint="https://test.cognitiveservices.azure.com/",
credential=AsyncMock(),
)
result = await provider.__aenter__()
assert result is provider
await provider.__aexit__(None, None, None)
async def test_aexit_closes_client(self) -> None:
provider = ContentUnderstandingContextProvider(
endpoint="https://test.cognitiveservices.azure.com/",
credential=AsyncMock(),
)
mock_client = AsyncMock()
provider._client = mock_client # type: ignore[assignment]
await provider.__aexit__(None, None, None)
mock_client.close.assert_called_once()
class TestBeforeRunNewFile:
async def test_single_pdf_analyzed(
self,
mock_cu_client: AsyncMock,
pdf_analysis_result: AnalysisResult,
) -> None:
mock_cu_client.begin_analyze_binary = AsyncMock(return_value=_make_mock_poller(pdf_analysis_result))
provider = _make_provider(mock_client=mock_cu_client)
msg = Message(
role="user",
contents=[
Content.from_text("What's on this invoice?"),
_make_content_from_data(_SAMPLE_PDF_BYTES, "application/pdf", "invoice.pdf"),
],
)
context = _make_context([msg])
state: dict[str, Any] = {}
session = AgentSession()
await provider.before_run(agent=_make_mock_agent(), session=session, context=context, state=state)
# Document should be in state
assert "documents" in state
assert "invoice.pdf" in state["documents"]
assert state["documents"]["invoice.pdf"]["status"] == DocumentStatus.READY
# Binary should be stripped from input
for m in context.input_messages:
for c in m.contents:
assert c.media_type != "application/pdf"
# Context should have messages injected
assert len(context.context_messages) > 0
async def test_url_input_analyzed(
self,
mock_cu_client: AsyncMock,
pdf_analysis_result: AnalysisResult,
) -> None:
mock_cu_client.begin_analyze = AsyncMock(return_value=_make_mock_poller(pdf_analysis_result))
provider = _make_provider(mock_client=mock_cu_client)
msg = Message(
role="user",
contents=[
Content.from_text("Analyze this document"),
Content.from_uri("https://example.com/report.pdf", media_type="application/pdf"),
],
)
context = _make_context([msg])
state: dict[str, Any] = {}
session = AgentSession()
await provider.before_run(agent=_make_mock_agent(), session=session, context=context, state=state)
# URL input should use begin_analyze
mock_cu_client.begin_analyze.assert_called_once()
assert "report.pdf" in state["documents"]
assert state["documents"]["report.pdf"]["status"] == DocumentStatus.READY
async def test_text_only_skipped(self, mock_cu_client: AsyncMock) -> None:
provider = _make_provider(mock_client=mock_cu_client)
msg = Message(role="user", contents=[Content.from_text("What's the weather?")])
context = _make_context([msg])
state: dict[str, Any] = {}
session = AgentSession()
await provider.before_run(agent=_make_mock_agent(), session=session, context=context, state=state)
# No CU calls
mock_cu_client.begin_analyze.assert_not_called()
mock_cu_client.begin_analyze_binary.assert_not_called()
# No documents
assert state.get("documents", {}) == {}
class TestBeforeRunMultiFile:
async def test_two_files_both_analyzed(
self,
mock_cu_client: AsyncMock,
pdf_analysis_result: AnalysisResult,
image_analysis_result: AnalysisResult,
) -> None:
mock_cu_client.begin_analyze_binary = AsyncMock(
side_effect=[
_make_mock_poller(pdf_analysis_result),
_make_mock_poller(image_analysis_result),
]
)
provider = _make_provider(mock_client=mock_cu_client)
msg = Message(
role="user",
contents=[
Content.from_text("Compare these documents"),
_make_content_from_data(_SAMPLE_PDF_BYTES, "application/pdf", "doc1.pdf"),
_make_content_from_data(b"\x89PNG fake", "image/png", "chart.png"),
],
)
context = _make_context([msg])
state: dict[str, Any] = {}
session = AgentSession()
await provider.before_run(agent=_make_mock_agent(), session=session, context=context, state=state)
assert len(state["documents"]) == 2
assert state["documents"]["doc1.pdf"]["status"] == DocumentStatus.READY
assert state["documents"]["chart.png"]["status"] == DocumentStatus.READY
class TestBeforeRunTimeout:
async def test_exceeds_max_wait_defers_to_background(
self,
mock_cu_client: AsyncMock,
pdf_analysis_result: AnalysisResult,
) -> None:
mock_cu_client.begin_analyze_binary = AsyncMock(return_value=_make_slow_poller(pdf_analysis_result, delay=10.0))
provider = _make_provider(mock_client=mock_cu_client, max_wait=0.1)
msg = Message(
role="user",
contents=[
Content.from_text("Analyze this"),
_make_content_from_data(_SAMPLE_PDF_BYTES, "application/pdf", "big_doc.pdf"),
],
)
context = _make_context([msg])
state: dict[str, Any] = {}
session = AgentSession()
await provider.before_run(agent=_make_mock_agent(), session=session, context=context, state=state)
assert state["documents"]["big_doc.pdf"]["status"] == DocumentStatus.ANALYZING
assert "big_doc.pdf" in state.get("_pending_tokens", {})
token_info = state["_pending_tokens"]["big_doc.pdf"]
assert "continuation_token" in token_info
assert "analyzer_id" in token_info
# Context messages should mention analyzing
assert any("being analyzed" in m.text for msgs in context.context_messages.values() for m in msgs)
class TestBeforeRunPendingResolution:
async def test_pending_completes_on_next_turn(
self,
mock_cu_client: AsyncMock,
pdf_analysis_result: AnalysisResult,
) -> None:
# Mock begin_analyze to return a completed poller when called with continuation_token
mock_poller = _make_mock_poller(pdf_analysis_result)
mock_poller.done = MagicMock(return_value=True)
mock_cu_client.begin_analyze = AsyncMock(return_value=mock_poller)
provider = _make_provider(mock_client=mock_cu_client)
state: dict[str, Any] = {
"_pending_tokens": {
"report.pdf": {"continuation_token": "tok_123", "analyzer_id": "prebuilt-documentSearch"}
},
"documents": {
"report.pdf": {
"status": DocumentStatus.ANALYZING,
"filename": "report.pdf",
"media_type": "application/pdf",
"analyzer_id": "prebuilt-documentSearch",
"analyzed_at": None,
"analysis_duration_s": None,
"upload_duration_s": None,
"result": None,
"error": None,
},
},
}
msg = Message(role="user", contents=[Content.from_text("Is the report ready?")])
context = _make_context([msg])
session = AgentSession()
await provider.before_run(agent=_make_mock_agent(), session=session, context=context, state=state)
assert state["documents"]["report.pdf"]["status"] == DocumentStatus.READY
assert state["documents"]["report.pdf"]["result"] is not None
assert "report.pdf" not in state.get("_pending_tokens", {})
class TestBeforeRunPendingFailure:
async def test_pending_task_failure_updates_state(
self,
mock_cu_client: AsyncMock,
) -> None:
# Mock begin_analyze to raise when resuming from continuation token
mock_cu_client.begin_analyze = AsyncMock(side_effect=RuntimeError("CU service unavailable"))
provider = _make_provider(mock_client=mock_cu_client)
state: dict[str, Any] = {
"_pending_tokens": {
"bad_doc.pdf": {"continuation_token": "tok_fail", "analyzer_id": "prebuilt-documentSearch"}
},
"documents": {
"bad_doc.pdf": {
"status": DocumentStatus.ANALYZING,
"filename": "bad_doc.pdf",
"media_type": "application/pdf",
"analyzer_id": "prebuilt-documentSearch",
"analyzed_at": None,
"analysis_duration_s": None,
"upload_duration_s": None,
"result": None,
"error": None,
},
},
}
msg = Message(role="user", contents=[Content.from_text("Check status")])
context = _make_context([msg])
session = AgentSession()
await provider.before_run(agent=_make_mock_agent(), session=session, context=context, state=state)
assert state["documents"]["bad_doc.pdf"]["status"] == DocumentStatus.FAILED
assert "CU service unavailable" in (state["documents"]["bad_doc.pdf"]["error"] or "")
class TestDocumentKeyDerivation:
def test_filename_from_additional_properties(self) -> None:
content = _make_content_from_data(_SAMPLE_PDF_BYTES, "application/pdf", "my_report.pdf")
key = derive_doc_key(content)
assert key == "my_report.pdf"
def test_url_basename(self) -> None:
content = Content.from_uri("https://example.com/docs/annual_report.pdf", media_type="application/pdf")
key = derive_doc_key(content)
assert key == "annual_report.pdf"
def test_content_hash_fallback(self) -> None:
content = Content.from_data(_SAMPLE_PDF_BYTES, "application/pdf")
key = derive_doc_key(content)
assert key.startswith("doc_")
assert len(key) == 12 # "doc_" + 8 hex chars
class TestSessionState:
async def test_documents_persist_across_turns(
self,
mock_cu_client: AsyncMock,
pdf_analysis_result: AnalysisResult,
) -> None:
mock_cu_client.begin_analyze_binary = AsyncMock(return_value=_make_mock_poller(pdf_analysis_result))
provider = _make_provider(mock_client=mock_cu_client)
state: dict[str, Any] = {}
session = AgentSession()
# Turn 1: upload
msg1 = Message(
role="user",
contents=[
Content.from_text("Analyze this"),
_make_content_from_data(_SAMPLE_PDF_BYTES, "application/pdf", "doc.pdf"),
],
)
ctx1 = _make_context([msg1])
await provider.before_run(agent=_make_mock_agent(), session=session, context=ctx1, state=state)
assert "doc.pdf" in state["documents"]
# Turn 2: follow-up (no file)
msg2 = Message(role="user", contents=[Content.from_text("What's the total?")])
ctx2 = _make_context([msg2])
await provider.before_run(agent=_make_mock_agent(), session=session, context=ctx2, state=state)
# Document should still be there
assert "doc.pdf" in state["documents"]
assert state["documents"]["doc.pdf"]["status"] == DocumentStatus.READY
class TestListDocumentsTool:
async def test_returns_all_docs_with_status(
self,
mock_cu_client: AsyncMock,
pdf_analysis_result: AnalysisResult,
) -> None:
mock_cu_client.begin_analyze_binary = AsyncMock(return_value=_make_mock_poller(pdf_analysis_result))
provider = _make_provider(mock_client=mock_cu_client)
state: dict[str, Any] = {}
session = AgentSession()
msg = Message(
role="user",
contents=[
Content.from_text("Analyze this"),
_make_content_from_data(_SAMPLE_PDF_BYTES, "application/pdf", "test.pdf"),
],
)
context = _make_context([msg])
await provider.before_run(agent=_make_mock_agent(), session=session, context=context, state=state)
# Find the list_documents tool
list_tool = None
for tool in context.tools:
if getattr(tool, "name", None) == "list_documents":
list_tool = tool
break
assert list_tool is not None
result = list_tool.func() # type: ignore[union-attr]
parsed = json.loads(result)
assert len(parsed) == 1
assert parsed[0]["name"] == "test.pdf"
assert parsed[0]["status"] == DocumentStatus.READY
class TestOutputFiltering:
"""Validate that output_sections controls what `_render_for_llm` emits.
Rendering is delegated to ``azure.ai.contentunderstanding.to_llm_input``:
- ``"markdown" in output_sections`` -> ``include_markdown=True``.
- ``"fields" in output_sections`` -> ``include_fields=True``.
- ``metadata={"source": <filename>}`` is always supplied.
Note: detailed field/JSON shape is owned by the SDK and exercised in the
SDK's own ``to_llm_input`` tests. We only assert MAF-level wiring here.
"""
def test_default_markdown_and_fields(self, pdf_analysis_result: AnalysisResult) -> None:
provider = _make_provider()
rendered = provider._render_for_llm(pdf_analysis_result, "report.pdf")
# YAML front matter with source key.
assert "source: report.pdf" in rendered
# PDF fixture contains "Contoso" in its markdown body.
assert "Contoso" in rendered
def test_markdown_only(self, pdf_analysis_result: AnalysisResult) -> None:
provider = _make_provider(output_sections=["markdown"])
rendered = provider._render_for_llm(pdf_analysis_result, "report.pdf")
# Markdown body still present; no ``fields:`` front-matter section.
assert "Contoso" in rendered
assert "\nfields:" not in rendered
assert not rendered.startswith("fields:")
def test_fields_only(self, invoice_analysis_result: AnalysisResult) -> None:
provider = _make_provider(output_sections=["fields"])
rendered = provider._render_for_llm(invoice_analysis_result, "invoice.pdf")
# ``fields:`` YAML key is emitted; vendor name appears under it.
assert "fields:" in rendered
assert "VendorName" in rendered
assert "TechServe Global Partners" in rendered
def test_field_values_extracted(self, invoice_analysis_result: AnalysisResult) -> None:
provider = _make_provider()
rendered = provider._render_for_llm(invoice_analysis_result, "invoice.pdf")
# Both sections present.
assert "fields:" in rendered
# Field values visible to the LLM (vendor + a known line-item description).
assert "TechServe Global Partners" in rendered
assert "Consulting Services" in rendered
def test_source_metadata_uses_filename(self, pdf_analysis_result: AnalysisResult) -> None:
"""Per-document ``source`` key carries the original filename."""
provider = _make_provider()
rendered = provider._render_for_llm(pdf_analysis_result, "custom_name.pdf")
assert "source: custom_name.pdf" in rendered
def test_page_markers_passed_through_to_llm_input(self, pdf_analysis_result: AnalysisResult) -> None:
"""Decision H: MAF must not strip page markers emitted by the SDK helper.
Today the SDK helper (``azure.ai.contentunderstanding.to_llm_input``)
injects ``<!-- page N -->`` markers per page. Per
``cognitive-services/ContentUnderstanding-Docs#249`` (Decision 4) it
will switch to ``<!-- InputPageNumber: N -->`` once the service ships
the marker natively. Either format must reach the LLM unchanged --
this test guards against MAF accidentally regex-stripping them.
"""
provider = _make_provider()
rendered = provider._render_for_llm(pdf_analysis_result, "report.pdf")
legacy = re.findall(r"<!--\s*page\s+\d+\s*-->", rendered)
future = re.findall(r"<!--\s*InputPageNumber:\s*\d+\s*-->", rendered)
# PDF fixture has 5 pages; expect 5 markers in whichever format is in use.
assert len(legacy) == 5 or len(future) == 5, (
"Expected SDK-injected page markers to be passed through to LLM input. "
f"Found legacy={len(legacy)}, future={len(future)}."
)
class TestDuplicateDocumentKey:
async def test_duplicate_filename_rejected(
self,
mock_cu_client: AsyncMock,
pdf_analysis_result: AnalysisResult,
) -> None:
"""Uploading the same filename twice in the same session should reject the second."""
mock_cu_client.begin_analyze_binary = AsyncMock(return_value=_make_mock_poller(pdf_analysis_result))
provider = _make_provider(mock_client=mock_cu_client)
# Turn 1: upload invoice.pdf
msg1 = Message(
role="user",
contents=[
Content.from_text("Analyze this"),
_make_content_from_data(_SAMPLE_PDF_BYTES, "application/pdf", "invoice.pdf"),
],
)
context1 = _make_context([msg1])
state: dict[str, Any] = {}
session = AgentSession()
await provider.before_run(agent=_make_mock_agent(), session=session, context=context1, state=state)
assert "invoice.pdf" in state["documents"]
assert state["documents"]["invoice.pdf"]["status"] == DocumentStatus.READY
# Turn 2: upload invoice.pdf again (different content but same filename)
msg2 = Message(
role="user",
contents=[
Content.from_text("Analyze this too"),
_make_content_from_data(b"different-content", "application/pdf", "invoice.pdf"),
],
)
context2 = _make_context([msg2])
await provider.before_run(agent=_make_mock_agent(), session=session, context=context2, state=state)
# Should still have only one document, not re-analyzed
assert mock_cu_client.begin_analyze_binary.call_count == 1
# Context messages should mention duplicate
assert any("already uploaded" in m.text for msgs in context2.context_messages.values() for m in msgs)
async def test_duplicate_in_same_turn_rejected(
self,
mock_cu_client: AsyncMock,
pdf_analysis_result: AnalysisResult,
) -> None:
"""Two files with the same filename in the same turn: first wins, second rejected."""
mock_cu_client.begin_analyze_binary = AsyncMock(return_value=_make_mock_poller(pdf_analysis_result))
provider = _make_provider(mock_client=mock_cu_client)
msg = Message(
role="user",
contents=[
Content.from_text("Analyze both"),
_make_content_from_data(_SAMPLE_PDF_BYTES, "application/pdf", "report.pdf"),
_make_content_from_data(b"other-content", "application/pdf", "report.pdf"),
],
)
context = _make_context([msg])
state: dict[str, Any] = {}
session = AgentSession()
await provider.before_run(agent=_make_mock_agent(), session=session, context=context, state=state)
# Only analyzed once (first one wins)
assert mock_cu_client.begin_analyze_binary.call_count == 1
assert "report.pdf" in state["documents"]
assert any("already uploaded" in m.text for msgs in context.context_messages.values() for m in msgs)
class TestBinaryStripping:
async def test_supported_files_stripped(
self,
mock_cu_client: AsyncMock,
pdf_analysis_result: AnalysisResult,
) -> None:
mock_cu_client.begin_analyze_binary = AsyncMock(return_value=_make_mock_poller(pdf_analysis_result))
provider = _make_provider(mock_client=mock_cu_client)
msg = Message(
role="user",
contents=[
Content.from_text("What's in here?"),
_make_content_from_data(_SAMPLE_PDF_BYTES, "application/pdf", "doc.pdf"),
],
)
context = _make_context([msg])
state: dict[str, Any] = {}
session = AgentSession()
await provider.before_run(agent=_make_mock_agent(), session=session, context=context, state=state)
# PDF should be stripped; text should remain
for m in context.input_messages:
for c in m.contents:
assert c.media_type != "application/pdf"
assert any(c.text and "What's in here?" in c.text for c in m.contents)
async def test_unsupported_files_left_in_place(self, mock_cu_client: AsyncMock) -> None:
provider = _make_provider(mock_client=mock_cu_client)
msg = Message(
role="user",
contents=[
Content.from_text("What's in this zip?"),
Content.from_data(
b"PK\x03\x04fake",
"application/zip",
additional_properties={"filename": "archive.zip"},
),
],
)
context = _make_context([msg])
state: dict[str, Any] = {}
session = AgentSession()
await provider.before_run(agent=_make_mock_agent(), session=session, context=context, state=state)
# Zip should NOT be stripped (unsupported)
found_zip = False
for m in context.input_messages:
for c in m.contents:
if c.media_type == "application/zip":
found_zip = True
assert found_zip
# Real magic-byte headers for binary sniffing tests
_MP4_MAGIC = b"\x00\x00\x00\x1cftypisom" + b"\x00" * 250
_WAV_MAGIC = b"RIFF\x00\x00\x00\x00WAVE" + b"\x00" * 250
_MP3_MAGIC = b"ID3\x04\x00\x00" + b"\x00" * 250
_FLAC_MAGIC = b"fLaC\x00\x00\x00\x00" + b"\x00" * 250
_OGG_MAGIC = b"OggS\x00\x02" + b"\x00" * 250
_AVI_MAGIC = b"RIFF\x00\x00\x00\x00AVI " + b"\x00" * 250
_MOV_MAGIC = b"\x00\x00\x00\x14ftypqt " + b"\x00" * 250
class TestMimeSniffing:
"""Tests for binary MIME sniffing via filetype when upstream MIME is unreliable."""
async def test_octet_stream_mp4_detected_and_stripped(
self,
mock_cu_client: AsyncMock,
pdf_analysis_result: AnalysisResult,
) -> None:
"""MP4 uploaded as application/octet-stream should be sniffed, corrected, and stripped."""
mock_cu_client.begin_analyze_binary = AsyncMock(return_value=_make_mock_poller(pdf_analysis_result))
provider = _make_provider(mock_client=mock_cu_client)
msg = Message(
role="user",
contents=[
Content.from_text("What's in this file?"),
_make_content_from_data(_MP4_MAGIC, "application/octet-stream", "video.mp4"),
],
)
context = _make_context([msg])
state: dict[str, Any] = {}
session = AgentSession()
await provider.before_run(agent=_make_mock_agent(), session=session, context=context, state=state)
# MP4 should be stripped from input
for m in context.input_messages:
for c in m.contents:
assert c.media_type != "application/octet-stream", "octet-stream content should be stripped"
# CU should have been called
assert mock_cu_client.begin_analyze_binary.called
async def test_octet_stream_wav_detected_via_sniff(
self,
mock_cu_client: AsyncMock,
pdf_analysis_result: AnalysisResult,
) -> None:
"""WAV uploaded as application/octet-stream should be detected via filetype sniffing."""
mock_cu_client.begin_analyze_binary = AsyncMock(return_value=_make_mock_poller(pdf_analysis_result))
provider = _make_provider(mock_client=mock_cu_client)
msg = Message(
role="user",
contents=[
Content.from_text("Transcribe"),
_make_content_from_data(_WAV_MAGIC, "application/octet-stream", "audio.wav"),
],
)
context = _make_context([msg])
state: dict[str, Any] = {}
session = AgentSession()
await provider.before_run(agent=_make_mock_agent(), session=session, context=context, state=state)
# Should be detected and analyzed
assert "audio.wav" in state["documents"]
# The media_type should be corrected to audio/wav (via _MIME_ALIASES)
assert state["documents"]["audio.wav"]["media_type"] == "audio/wav"
async def test_octet_stream_mp3_detected_via_sniff(
self,
mock_cu_client: AsyncMock,
pdf_analysis_result: AnalysisResult,
) -> None:
"""MP3 uploaded as application/octet-stream should be detected as audio/mpeg."""
mock_cu_client.begin_analyze_binary = AsyncMock(return_value=_make_mock_poller(pdf_analysis_result))
provider = _make_provider(mock_client=mock_cu_client)
msg = Message(
role="user",
contents=[
Content.from_text("Transcribe"),
_make_content_from_data(_MP3_MAGIC, "application/octet-stream", "song.mp3"),
],
)
context = _make_context([msg])
state: dict[str, Any] = {}
session = AgentSession()
await provider.before_run(agent=_make_mock_agent(), session=session, context=context, state=state)
assert "song.mp3" in state["documents"]
assert state["documents"]["song.mp3"]["media_type"] == "audio/mpeg"
async def test_octet_stream_flac_alias_normalized(
self,
mock_cu_client: AsyncMock,
pdf_analysis_result: AnalysisResult,
) -> None:
"""FLAC sniffed as audio/x-flac should be normalized to audio/flac."""
mock_cu_client.begin_analyze_binary = AsyncMock(return_value=_make_mock_poller(pdf_analysis_result))
provider = _make_provider(mock_client=mock_cu_client)
msg = Message(
role="user",
contents=[
Content.from_text("Transcribe"),
_make_content_from_data(_FLAC_MAGIC, "application/octet-stream", "music.flac"),
],
)
context = _make_context([msg])
state: dict[str, Any] = {}
session = AgentSession()
await provider.before_run(agent=_make_mock_agent(), session=session, context=context, state=state)
assert "music.flac" in state["documents"]
assert state["documents"]["music.flac"]["media_type"] == "audio/flac"
async def test_octet_stream_unknown_binary_not_stripped(
self,
mock_cu_client: AsyncMock,
) -> None:
"""Unknown binary with application/octet-stream should NOT be stripped."""
provider = _make_provider(mock_client=mock_cu_client)
unknown_bytes = b"\x00\x01\x02\x03random garbage" + b"\x00" * 250
msg = Message(
role="user",
contents=[
Content.from_text("What is this?"),
_make_content_from_data(unknown_bytes, "application/octet-stream", "mystery.bin"),
],
)
context = _make_context([msg])
state: dict[str, Any] = {}
session = AgentSession()
await provider.before_run(agent=_make_mock_agent(), session=session, context=context, state=state)
# Unknown file should NOT be stripped
found_octet = False
for m in context.input_messages:
for c in m.contents:
if c.media_type == "application/octet-stream":
found_octet = True
assert found_octet
async def test_missing_mime_falls_back_to_filename(
self,
mock_cu_client: AsyncMock,
pdf_analysis_result: AnalysisResult,
) -> None:
"""Content with empty MIME but a .mp4 filename should be detected via mimetypes fallback."""
mock_cu_client.begin_analyze_binary = AsyncMock(return_value=_make_mock_poller(pdf_analysis_result))
provider = _make_provider(mock_client=mock_cu_client)
# Use garbage binary (filetype won't detect) but filename has .mp4
garbage = b"\x00" * 300
content = Content.from_data(garbage, "", additional_properties={"filename": "recording.mp4"})
msg = Message(
role="user",
contents=[Content.from_text("Analyze"), content],
)
context = _make_context([msg])
state: dict[str, Any] = {}
session = AgentSession()
await provider.before_run(agent=_make_mock_agent(), session=session, context=context, state=state)
# Should be detected via filename and analyzed
assert "recording.mp4" in state["documents"]
async def test_correct_mime_not_sniffed(
self,
mock_cu_client: AsyncMock,
pdf_analysis_result: AnalysisResult,
) -> None:
"""Files with correct MIME type should go through fast path without sniffing."""
mock_cu_client.begin_analyze_binary = AsyncMock(return_value=_make_mock_poller(pdf_analysis_result))
provider = _make_provider(mock_client=mock_cu_client)
msg = Message(
role="user",
contents=[
Content.from_text("Analyze"),
_make_content_from_data(_SAMPLE_PDF_BYTES, "application/pdf", "doc.pdf"),
],
)
context = _make_context([msg])
state: dict[str, Any] = {}
session = AgentSession()
await provider.before_run(agent=_make_mock_agent(), session=session, context=context, state=state)
assert "doc.pdf" in state["documents"]
assert state["documents"]["doc.pdf"]["media_type"] == "application/pdf"
async def test_sniffed_video_uses_correct_analyzer(
self,
mock_cu_client: AsyncMock,
pdf_analysis_result: AnalysisResult,
) -> None:
"""MP4 sniffed from octet-stream should use prebuilt-videoSearch analyzer."""
mock_cu_client.begin_analyze_binary = AsyncMock(return_value=_make_mock_poller(pdf_analysis_result))
provider = _make_provider(mock_client=mock_cu_client) # analyzer_id=None → auto-detect
msg = Message(
role="user",
contents=[
Content.from_text("What's in this video?"),
_make_content_from_data(_MP4_MAGIC, "application/octet-stream", "demo.mp4"),
],
)
context = _make_context([msg])
state: dict[str, Any] = {}
session = AgentSession()
await provider.before_run(agent=_make_mock_agent(), session=session, context=context, state=state)
assert state["documents"]["demo.mp4"]["analyzer_id"] == "prebuilt-videoSearch"
class TestErrorHandling:
async def test_cu_service_error(self, mock_cu_client: AsyncMock) -> None:
mock_cu_client.begin_analyze_binary = AsyncMock(
return_value=_make_failing_poller(RuntimeError("Service unavailable"))
)
provider = _make_provider(mock_client=mock_cu_client)
msg = Message(
role="user",
contents=[
Content.from_text("Analyze this"),
_make_content_from_data(_SAMPLE_PDF_BYTES, "application/pdf", "error.pdf"),
],
)
context = _make_context([msg])
state: dict[str, Any] = {}
session = AgentSession()
await provider.before_run(agent=_make_mock_agent(), session=session, context=context, state=state)
assert state["documents"]["error.pdf"]["status"] == DocumentStatus.FAILED
assert "Service unavailable" in (state["documents"]["error.pdf"]["error"] or "")
async def test_lazy_initialization_on_before_run(self) -> None:
"""before_run works with eagerly-initialized client."""
provider = ContentUnderstandingContextProvider(
endpoint="https://test.cognitiveservices.azure.com/",
credential=AsyncMock(),
)
assert provider._client is not None
mock_client = AsyncMock()
mock_client.begin_analyze_binary = AsyncMock(
side_effect=Exception("mock error"),
)
provider._client = mock_client # type: ignore[assignment]
msg = Message(
role="user",
contents=[
Content.from_text("Analyze this"),
_make_content_from_data(_SAMPLE_PDF_BYTES, "application/pdf", "doc.pdf"),
],
)
context = _make_context([msg])
state: dict[str, Any] = {}
session = AgentSession()
await provider.before_run(agent=_make_mock_agent(), session=session, context=context, state=state)
# Client should still be set
assert provider._client is not None
class TestMultiModalFixtures:
"""Verify ``_render_for_llm`` produces sensible output for each modality.
Detailed shape of the YAML/Markdown payload is the SDK's responsibility and
is exercised by ``azure-ai-contentunderstanding`` tests. Here we only check
that the MAF wiring (filename surfaced as ``source``, key content visible)
works for each fixture kind.
"""
def test_pdf_fixture_loads(self, pdf_analysis_result: AnalysisResult) -> None:
provider = _make_provider()
rendered = provider._render_for_llm(pdf_analysis_result, "report.pdf")
assert "source: report.pdf" in rendered
assert "Contoso" in rendered
def test_audio_fixture_loads(self, audio_analysis_result: AnalysisResult) -> None:
provider = _make_provider()
rendered = provider._render_for_llm(audio_analysis_result, "call.mp3")
assert "source: call.mp3" in rendered
assert "Call Center" in rendered
def test_video_fixture_loads(self, video_analysis_result: AnalysisResult) -> None:
provider = _make_provider()
rendered = provider._render_for_llm(video_analysis_result, "demo.mp4")
assert "source: demo.mp4" in rendered
# All 3 segments should be visible in the rendered text.
assert "Contoso Product Demo" in rendered
assert "real-time monitoring" in rendered
assert "contoso.com/cloud-manager" in rendered
# Each segment must render its own YAML front matter with a timeRange entry.
# This guards against multi-segment results being collapsed into one block.
assert rendered.count("timeRange:") == 3
# Segments must be rendered in chronological order (1s, 15s, 36s starts).
assert (
rendered.index("Contoso Product Demo")
< rendered.index("real-time monitoring")
< rendered.index("contoso.com/cloud-manager")
)
def test_image_fixture_loads(self, image_analysis_result: AnalysisResult) -> None:
provider = _make_provider()
rendered = provider._render_for_llm(image_analysis_result, "image.png")
assert "source: image.png" in rendered
# Non-empty body (image markdown caption from CU).
assert len(rendered) > len("source: image.png")
def test_invoice_fixture_loads(self, invoice_analysis_result: AnalysisResult) -> None:
provider = _make_provider()
rendered = provider._render_for_llm(invoice_analysis_result, "invoice.pdf")
assert "source: invoice.pdf" in rendered
assert "fields:" in rendered
assert "VendorName" in rendered
# NOTE: ``TestFormatResult`` (4 tests) was deleted as part of the migration to
# ``azure.ai.contentunderstanding.to_llm_input``. The legacy ``format_result``
# helper no longer exists; rendering shape (YAML front matter + Markdown body,
# segment serialization, reserved-key handling) is owned and tested by the SDK.
class TestSupportedMediaTypes:
def test_pdf_supported(self) -> None:
assert "application/pdf" in SUPPORTED_MEDIA_TYPES
def test_audio_supported(self) -> None:
assert "audio/mp3" in SUPPORTED_MEDIA_TYPES
assert "audio/wav" in SUPPORTED_MEDIA_TYPES
def test_video_supported(self) -> None:
assert "video/mp4" in SUPPORTED_MEDIA_TYPES
def test_zip_not_supported(self) -> None:
assert "application/zip" not in SUPPORTED_MEDIA_TYPES
class TestAnalyzerAutoDetection:
"""Verify _resolve_analyzer_id auto-selects the right analyzer by media type."""
def test_explicit_analyzer_always_wins(self) -> None:
provider = _make_provider(analyzer_id="prebuilt-invoice")
assert provider._resolve_analyzer_id("audio/mp3") == "prebuilt-invoice"
assert provider._resolve_analyzer_id("video/mp4") == "prebuilt-invoice"
assert provider._resolve_analyzer_id("application/pdf") == "prebuilt-invoice"
def test_auto_detect_pdf(self) -> None:
provider = _make_provider() # analyzer_id=None
assert provider._resolve_analyzer_id("application/pdf") == "prebuilt-documentSearch"
def test_auto_detect_image(self) -> None:
provider = _make_provider()
assert provider._resolve_analyzer_id("image/jpeg") == "prebuilt-documentSearch"
assert provider._resolve_analyzer_id("image/png") == "prebuilt-documentSearch"
def test_auto_detect_audio(self) -> None:
provider = _make_provider()
assert provider._resolve_analyzer_id("audio/mp3") == "prebuilt-audioSearch"
assert provider._resolve_analyzer_id("audio/wav") == "prebuilt-audioSearch"
assert provider._resolve_analyzer_id("audio/mpeg") == "prebuilt-audioSearch"
def test_auto_detect_video(self) -> None:
provider = _make_provider()
assert provider._resolve_analyzer_id("video/mp4") == "prebuilt-videoSearch"
assert provider._resolve_analyzer_id("video/webm") == "prebuilt-videoSearch"
def test_auto_detect_unknown_falls_back_to_document(self) -> None:
provider = _make_provider()
assert provider._resolve_analyzer_id("application/octet-stream") == "prebuilt-documentSearch"
class TestFileSearchIntegration:
_FILE_SEARCH_TOOL = {"type": "file_search", "vector_store_ids": ["vs_test123"]}
def _make_mock_backend(self) -> AsyncMock:
"""Create a mock FileSearchBackend."""
backend = AsyncMock()
backend.upload_file = AsyncMock(return_value="file_test456")
backend.delete_file = AsyncMock()
return backend
def _make_file_search_config(self, backend: AsyncMock | None = None) -> Any:
from agent_framework_azure_contentunderstanding import FileSearchConfig
return FileSearchConfig(
backend=backend or self._make_mock_backend(),
vector_store_id="vs_test123",
file_search_tool=self._FILE_SEARCH_TOOL,
)
async def test_file_search_uploads_to_vector_store(
self,
mock_cu_client: AsyncMock,
pdf_analysis_result: AnalysisResult,
) -> None:
mock_backend = self._make_mock_backend()
config = self._make_file_search_config(mock_backend)
mock_cu_client.begin_analyze_binary = AsyncMock(
return_value=_make_mock_poller(pdf_analysis_result),
)
provider = _make_provider(
mock_client=mock_cu_client,
file_search=config,
)
msg = Message(
role="user",
contents=[
Content.from_text("Analyze this"),
_make_content_from_data(_SAMPLE_PDF_BYTES, "application/pdf", "doc.pdf"),
],
)
context = _make_context([msg])
state: dict[str, Any] = {}
session = AgentSession()
await provider.before_run(
agent=_make_mock_agent(),
session=session,
context=context,
state=state,
)
# File should be uploaded via backend
mock_backend.upload_file.assert_called_once()
call_args = mock_backend.upload_file.call_args
assert call_args[0][0] == "vs_test123" # vector_store_id
assert call_args[0][1] == "doc.pdf.md" # filename
# file_search tool should be registered on context
assert self._FILE_SEARCH_TOOL in context.tools
async def test_file_search_no_content_injection(
self,
mock_cu_client: AsyncMock,
pdf_analysis_result: AnalysisResult,
) -> None:
"""When file_search is enabled, full content should NOT be injected into context."""
mock_cu_client.begin_analyze_binary = AsyncMock(
return_value=_make_mock_poller(pdf_analysis_result),
)
provider = _make_provider(
mock_client=mock_cu_client,
file_search=self._make_file_search_config(),
)
msg = Message(
role="user",
contents=[
Content.from_text("Analyze this"),
_make_content_from_data(_SAMPLE_PDF_BYTES, "application/pdf", "doc.pdf"),
],
)
context = _make_context([msg])
state: dict[str, Any] = {}
session = AgentSession()
await provider.before_run(
agent=_make_mock_agent(),
session=session,
context=context,
state=state,
)
# Context messages should NOT contain full document content
# (file_search handles retrieval instead)
for msgs in context.context_messages.values():
for m in msgs:
assert "Document Content" not in m.text
async def test_cleanup_deletes_uploaded_files(
self,
mock_cu_client: AsyncMock,
pdf_analysis_result: AnalysisResult,
) -> None:
mock_backend = self._make_mock_backend()
config = self._make_file_search_config(mock_backend)
mock_cu_client.begin_analyze_binary = AsyncMock(
return_value=_make_mock_poller(pdf_analysis_result),
)
provider = _make_provider(
mock_client=mock_cu_client,
file_search=config,
)
msg = Message(
role="user",
contents=[
Content.from_text("Analyze this"),
_make_content_from_data(_SAMPLE_PDF_BYTES, "application/pdf", "doc.pdf"),
],
)
context = _make_context([msg])
state: dict[str, Any] = {}
session = AgentSession()
await provider.before_run(
agent=_make_mock_agent(),
session=session,
context=context,
state=state,
)
# Close should clean up uploaded files (not the vector store itself)
await provider.close()
mock_backend.delete_file.assert_called_once_with("file_test456")
async def test_no_file_search_injects_content(
self,
mock_cu_client: AsyncMock,
pdf_analysis_result: AnalysisResult,
) -> None:
"""Without file_search, full content should be injected (default behavior)."""
mock_cu_client.begin_analyze_binary = AsyncMock(
return_value=_make_mock_poller(pdf_analysis_result),
)
provider = _make_provider(mock_client=mock_cu_client)
msg = Message(
role="user",
contents=[
Content.from_text("Analyze this"),
_make_content_from_data(_SAMPLE_PDF_BYTES, "application/pdf", "doc.pdf"),
],
)
context = _make_context([msg])
state: dict[str, Any] = {}
session = AgentSession()
await provider.before_run(
agent=_make_mock_agent(),
session=session,
context=context,
state=state,
)
# Without file_search, content SHOULD be injected
found_content = False
for msgs in context.context_messages.values():
for m in msgs:
if "Document Content" in m.text or "Contoso" in m.text:
found_content = True
assert found_content
async def test_file_search_multiple_files(
self,
mock_cu_client: AsyncMock,
pdf_analysis_result: AnalysisResult,
audio_analysis_result: AnalysisResult,
) -> None:
"""Multiple files should each be uploaded to the vector store."""
mock_backend = self._make_mock_backend()
# Return different file IDs for each upload
mock_backend.upload_file = AsyncMock(side_effect=["file_001", "file_002"])
config = self._make_file_search_config(mock_backend)
mock_cu_client.begin_analyze_binary = AsyncMock(
side_effect=[
_make_mock_poller(pdf_analysis_result),
_make_mock_poller(audio_analysis_result),
],
)
provider = _make_provider(
mock_client=mock_cu_client,
file_search=config,
)
msg = Message(
role="user",
contents=[
Content.from_text("Compare these"),
_make_content_from_data(_SAMPLE_PDF_BYTES, "application/pdf", "doc.pdf"),
_make_content_from_data(b"\x00audio-fake", "audio/mp3", "call.mp3"),
],
)
context = _make_context([msg])
state: dict[str, Any] = {}
session = AgentSession()
await provider.before_run(agent=_make_mock_agent(), session=session, context=context, state=state)
# Two files uploaded via backend
assert mock_backend.upload_file.call_count == 2
async def test_file_search_skips_empty_markdown(
self,
mock_cu_client: AsyncMock,
) -> None:
"""Upload should be skipped when CU returns no markdown content."""
mock_backend = self._make_mock_backend()
config = self._make_file_search_config(mock_backend)
# Create a result with empty markdown
empty_result = AnalysisResult({"contents": [{"markdown": "", "fields": {}}]})
mock_cu_client.begin_analyze_binary = AsyncMock(
return_value=_make_mock_poller(empty_result),
)
provider = _make_provider(
mock_client=mock_cu_client,
file_search=config,
)
msg = Message(
role="user",
contents=[
Content.from_text("Analyze this"),
_make_content_from_data(_SAMPLE_PDF_BYTES, "application/pdf", "empty.pdf"),
],
)
context = _make_context([msg])
state: dict[str, Any] = {}
session = AgentSession()
await provider.before_run(agent=_make_mock_agent(), session=session, context=context, state=state)
# No file should be uploaded (empty markdown)
mock_backend.upload_file.assert_not_called()
async def test_pending_resolution_uploads_to_vector_store(
self,
mock_cu_client: AsyncMock,
pdf_analysis_result: AnalysisResult,
) -> None:
"""When a background task completes in file_search mode, content should be
uploaded to the vector store — NOT injected into context messages."""
mock_backend = self._make_mock_backend()
config = self._make_file_search_config(mock_backend)
provider = _make_provider(
mock_client=mock_cu_client,
file_search=config,
)
# Simulate a completed background analysis via continuation token
mock_poller = _make_mock_poller(pdf_analysis_result)
mock_poller.done = MagicMock(return_value=True)
mock_cu_client.begin_analyze = AsyncMock(return_value=mock_poller)
state: dict[str, Any] = {
"_pending_tokens": {
"report.pdf": {"continuation_token": "tok_fs", "analyzer_id": "prebuilt-documentSearch"}
},
"documents": {
"report.pdf": {
"status": DocumentStatus.ANALYZING,
"filename": "report.pdf",
"media_type": "application/pdf",
"analyzer_id": "prebuilt-documentSearch",
"analyzed_at": None,
"analysis_duration_s": None,
"upload_duration_s": None,
"result": None,
"error": None,
},
},
}
msg = Message(role="user", contents=[Content.from_text("Is the report ready?")])
context = _make_context([msg])
session = AgentSession()
await provider.before_run(agent=_make_mock_agent(), session=session, context=context, state=state)
# Document should be ready
assert state["documents"]["report.pdf"]["status"] == DocumentStatus.READY
# Content should NOT be injected into context messages
for msgs in context.context_messages.values():
for m in msgs:
assert "Document Content" not in m.text
# Should be uploaded to vector store via backend
mock_backend.upload_file.assert_called_once()
# Context messages should mention file_search, not "provided above"
all_msg_text = " ".join(m.text for msgs in context.context_messages.values() for m in msgs)
assert "file_search" in all_msg_text or any("file_search" in instr for instr in context.instructions)
assert "provided above" not in all_msg_text
class TestCloseCancel:
async def test_close_cleans_up(self) -> None:
"""close() should close the CU client."""
provider = _make_provider(mock_client=AsyncMock())
await provider.close()
# Client should be closed (no tasks to cancel — tokens are just strings)
cast(Any, provider._client.close).assert_called_once()
class TestSessionIsolation:
"""Verify that per-session state (pending tasks, uploads) is isolated between sessions."""
async def test_background_task_isolated_per_session(
self,
mock_cu_client: AsyncMock,
pdf_analysis_result: AnalysisResult,
) -> None:
"""A background task from session A must not leak into session B."""
mock_cu_client.begin_analyze_binary = AsyncMock(return_value=_make_slow_poller(pdf_analysis_result, delay=10.0))
provider = _make_provider(mock_client=mock_cu_client, max_wait=0.1)
# Session A: upload a file that times out → defers to background
msg_a = Message(
role="user",
contents=[
Content.from_text("Analyze this"),
_make_content_from_data(_SAMPLE_PDF_BYTES, "application/pdf", "report.pdf"),
],
)
state_a: dict[str, Any] = {}
context_a = _make_context([msg_a])
await provider.before_run(agent=_make_mock_agent(), session=AgentSession(), context=context_a, state=state_a)
# Session A should have a pending token
assert "report.pdf" in state_a.get("_pending_tokens", {})
# Session B: separate state, no pending tokens
state_b: dict[str, Any] = {}
msg_b = Message(role="user", contents=[Content.from_text("Hello")])
context_b = _make_context([msg_b])
await provider.before_run(agent=_make_mock_agent(), session=AgentSession(), context=context_b, state=state_b)
# Session B must NOT see session A's pending token
assert "_pending_tokens" not in state_b or "report.pdf" not in state_b.get("_pending_tokens", {})
# Session B must NOT have session A's documents
assert "report.pdf" not in state_b.get("documents", {})
async def test_completed_task_resolves_in_correct_session(
self,
mock_cu_client: AsyncMock,
pdf_analysis_result: AnalysisResult,
) -> None:
"""A completed background task should only inject content into its own session."""
provider = _make_provider(mock_client=mock_cu_client)
# Simulate completed analysis in session A via continuation token
mock_poller = _make_mock_poller(pdf_analysis_result)
mock_poller.done = MagicMock(return_value=True)
mock_cu_client.begin_analyze = AsyncMock(return_value=mock_poller)
state_a: dict[str, Any] = {
"_pending_tokens": {
"report.pdf": {"continuation_token": "tok_a", "analyzer_id": "prebuilt-documentSearch"}
},
"documents": {
"report.pdf": {
"status": DocumentStatus.ANALYZING,
"filename": "report.pdf",
"media_type": "application/pdf",
"analyzer_id": "prebuilt-documentSearch",
"analyzed_at": None,
"analysis_duration_s": None,
"upload_duration_s": None,
"result": None,
"error": None,
},
},
}
state_b: dict[str, Any] = {}
# Run session A — should resolve the task
context_a = _make_context([Message(role="user", contents=[Content.from_text("Is it ready?")])])
await provider.before_run(agent=_make_mock_agent(), session=AgentSession(), context=context_a, state=state_a)
assert state_a["documents"]["report.pdf"]["status"] == DocumentStatus.READY
# Run session B — must NOT have any documents or resolved content
context_b = _make_context([Message(role="user", contents=[Content.from_text("Hello")])])
await provider.before_run(agent=_make_mock_agent(), session=AgentSession(), context=context_b, state=state_b)
assert "report.pdf" not in state_b.get("documents", {})
# Session B context should have no document-related messages
assert not any("report.pdf" in m.text for msgs in context_b.context_messages.values() for m in msgs)
class TestAnalyzerAutoDetectionE2E:
"""End-to-end: verify _analyze_file stores the resolved analyzer in DocumentEntry."""
async def test_audio_file_uses_audio_analyzer(
self,
mock_cu_client: AsyncMock,
audio_analysis_result: AnalysisResult,
) -> None:
mock_cu_client.begin_analyze_binary = AsyncMock(
return_value=_make_mock_poller(audio_analysis_result),
)
provider = _make_provider(mock_client=mock_cu_client) # analyzer_id=None
msg = Message(
role="user",
contents=[
Content.from_text("Transcribe this"),
_make_content_from_data(b"\x00audio", "audio/mp3", "call.mp3"),
],
)
context = _make_context([msg])
state: dict[str, Any] = {}
session = AgentSession()
await provider.before_run(agent=_make_mock_agent(), session=session, context=context, state=state)
assert state["documents"]["call.mp3"]["analyzer_id"] == "prebuilt-audioSearch"
# CU client should have been called with the audio analyzer
mock_cu_client.begin_analyze_binary.assert_called_once()
call_args = mock_cu_client.begin_analyze_binary.call_args
assert call_args[0][0] == "prebuilt-audioSearch"
async def test_video_file_uses_video_analyzer(
self,
mock_cu_client: AsyncMock,
video_analysis_result: AnalysisResult,
) -> None:
mock_cu_client.begin_analyze_binary = AsyncMock(
return_value=_make_mock_poller(video_analysis_result),
)
provider = _make_provider(mock_client=mock_cu_client)
msg = Message(
role="user",
contents=[
Content.from_text("Analyze this video"),
_make_content_from_data(b"\x00video", "video/mp4", "demo.mp4"),
],
)
context = _make_context([msg])
state: dict[str, Any] = {}
session = AgentSession()
await provider.before_run(agent=_make_mock_agent(), session=session, context=context, state=state)
assert state["documents"]["demo.mp4"]["analyzer_id"] == "prebuilt-videoSearch"
call_args = mock_cu_client.begin_analyze_binary.call_args
assert call_args[0][0] == "prebuilt-videoSearch"
async def test_pdf_file_uses_document_analyzer(
self,
mock_cu_client: AsyncMock,
pdf_analysis_result: AnalysisResult,
) -> None:
mock_cu_client.begin_analyze_binary = AsyncMock(
return_value=_make_mock_poller(pdf_analysis_result),
)
provider = _make_provider(mock_client=mock_cu_client)
msg = Message(
role="user",
contents=[
Content.from_text("Read this"),
_make_content_from_data(_SAMPLE_PDF_BYTES, "application/pdf", "report.pdf"),
],
)
context = _make_context([msg])
state: dict[str, Any] = {}
session = AgentSession()
await provider.before_run(agent=_make_mock_agent(), session=session, context=context, state=state)
assert state["documents"]["report.pdf"]["analyzer_id"] == "prebuilt-documentSearch"
call_args = mock_cu_client.begin_analyze_binary.call_args
assert call_args[0][0] == "prebuilt-documentSearch"
async def test_explicit_override_ignores_media_type(
self,
mock_cu_client: AsyncMock,
audio_analysis_result: AnalysisResult,
) -> None:
"""Explicit analyzer_id should override auto-detection even for audio."""
mock_cu_client.begin_analyze_binary = AsyncMock(
return_value=_make_mock_poller(audio_analysis_result),
)
provider = _make_provider(mock_client=mock_cu_client, analyzer_id="prebuilt-invoice")
msg = Message(
role="user",
contents=[
Content.from_text("Analyze"),
_make_content_from_data(b"\x00audio", "audio/mp3", "call.mp3"),
],
)
context = _make_context([msg])
state: dict[str, Any] = {}
session = AgentSession()
await provider.before_run(agent=_make_mock_agent(), session=session, context=context, state=state)
assert state["documents"]["call.mp3"]["analyzer_id"] == "prebuilt-invoice"
call_args = mock_cu_client.begin_analyze_binary.call_args
assert call_args[0][0] == "prebuilt-invoice"
async def test_per_file_analyzer_overrides_provider_default(
self,
mock_cu_client: AsyncMock,
pdf_analysis_result: AnalysisResult,
) -> None:
"""Per-file analyzer_id in additional_properties overrides provider-level default."""
mock_cu_client.begin_analyze_binary = AsyncMock(
return_value=_make_mock_poller(pdf_analysis_result),
)
# Provider default is prebuilt-documentSearch
provider = _make_provider(
mock_client=mock_cu_client,
analyzer_id="prebuilt-documentSearch",
)
msg = Message(
role="user",
contents=[
Content.from_text("Process this invoice"),
Content.from_data(
_SAMPLE_PDF_BYTES,
"application/pdf",
# Per-file override to prebuilt-invoice
additional_properties={
"filename": "invoice.pdf",
"analyzer_id": "prebuilt-invoice",
},
),
],
)
context = _make_context([msg])
state: dict[str, Any] = {}
session = AgentSession()
await provider.before_run(agent=_make_mock_agent(), session=session, context=context, state=state)
# Per-file override should win
assert state["documents"]["invoice.pdf"]["analyzer_id"] == "prebuilt-invoice"
call_args = mock_cu_client.begin_analyze_binary.call_args
assert call_args[0][0] == "prebuilt-invoice"
class TestWarningsExtraction:
"""Verify that CU RAI warnings are surfaced via ``to_llm_input`` rendering.
The SDK serializes ``result.warnings`` under the reserved ``rai_warnings``
YAML front-matter key. Telemetry filtering of stray ``LLMStats:`` lines is
handled by the SDK helper (azure-ai-contentunderstanding >= 1.2.0b2).
"""
def test_warnings_included_when_present(self) -> None:
"""Non-empty warnings should appear under ``rai_warnings`` front-matter key."""
provider = _make_provider()
fixture = {
"contents": [
{
"path": "input1",
"markdown": "Some content",
"kind": "document",
}
],
"warnings": [
{
"code": "ContentFiltered",
"message": "Content was filtered due to Responsible AI policy.",
"target": "contents/0/markdown",
},
{
"code": "ContentFiltered",
"message": "Violence content detected and filtered.",
},
],
}
result_obj = AnalysisResult(fixture)
rendered = provider._render_for_llm(result_obj, "doc.pdf")
assert "rai_warnings:" in rendered
assert "ContentFiltered" in rendered
assert "Content was filtered due to Responsible AI policy." in rendered
assert "Violence content detected and filtered." in rendered
def test_warnings_omitted_when_empty(self, pdf_analysis_result: AnalysisResult) -> None:
"""The PDF fixture has no warnings, so ``rai_warnings:`` should not appear."""
provider = _make_provider()
rendered = provider._render_for_llm(pdf_analysis_result, "report.pdf")
assert "rai_warnings:" not in rendered
class TestCategoryExtraction:
"""Verify category metadata (from classifier analyzers) is rendered into output."""
def test_category_included_single_segment(self) -> None:
"""Category from classifier should appear under the ``category`` front-matter key."""
provider = _make_provider()
fixture = {
"contents": [
{
"path": "input1",
"markdown": "Contract text...",
"kind": "document",
"category": "Legal Contract",
}
],
}
result_obj = AnalysisResult(fixture)
rendered = provider._render_for_llm(result_obj, "contract.pdf")
assert "category:" in rendered
assert "Legal Contract" in rendered
def test_category_in_multi_segment_video(self) -> None:
"""Each segment's category should be visible in the rendered text."""
provider = _make_provider()
fixture = {
"contents": [
{
"path": "input1",
"kind": "audioVisual",
"startTimeMs": 0,
"endTimeMs": 30000,
"markdown": "Opening scene with product showcase.",
"category": "ProductDemo",
"fields": {
"Summary": {
"type": "string",
"valueString": "Product demo intro",
}
},
},
{
"path": "input1",
"kind": "audioVisual",
"startTimeMs": 30000,
"endTimeMs": 60000,
"markdown": "Customer testimonial segment.",
"category": "Testimonial",
"fields": {
"Summary": {
"type": "string",
"valueString": "Customer feedback",
}
},
},
],
}
result_obj = AnalysisResult(fixture)
rendered = provider._render_for_llm(result_obj, "promo.mp4")
# Both segments' markdown content visible.
assert "Opening scene with product showcase." in rendered
assert "Customer testimonial segment." in rendered
# Both categories visible.
assert "ProductDemo" in rendered
assert "Testimonial" in rendered
# Segments must be rendered in source order, not arbitrary.
assert rendered.index("Opening scene with product showcase.") < rendered.index("Customer testimonial segment.")
# Category-to-segment mapping must be correct. The SDK separates segments
# with a ``*****`` line, so split on it and verify each block carries the
# right category alongside the right markdown body.
blocks = rendered.split("*****")
assert len(blocks) == 2, f"expected 2 segment blocks, got {len(blocks)}"
assert "Opening scene with product showcase." in blocks[0]
assert "category: ProductDemo" in blocks[0]
assert "Customer testimonial segment." in blocks[1]
assert "category: Testimonial" in blocks[1]
def test_category_omitted_when_none(self, pdf_analysis_result: AnalysisResult) -> None:
"""No category should be in output when the analyzer doesn't classify."""
provider = _make_provider()
rendered = provider._render_for_llm(pdf_analysis_result, "report.pdf")
assert "category:" not in rendered
class TestContentRangeSupport:
"""Verify that content_range from additional_properties is passed to CU."""
async def test_content_range_passed_to_begin_analyze(
self,
mock_cu_client: AsyncMock,
pdf_analysis_result: AnalysisResult,
) -> None:
"""content_range in additional_properties should be forwarded to AnalysisInput."""
from azure.ai.contentunderstanding.models import AnalysisInput
mock_cu_client.begin_analyze = AsyncMock(return_value=_make_mock_poller(pdf_analysis_result))
provider = _make_provider(mock_client=mock_cu_client)
msg = Message(
role="user",
contents=[
Content.from_text("Analyze pages 1-3"),
Content.from_uri(
"https://example.com/report.pdf",
media_type="application/pdf",
additional_properties={"filename": "report.pdf", "content_range": "1-3"},
),
],
)
context = _make_context([msg])
state: dict[str, Any] = {}
session = AgentSession()
await provider.before_run(agent=_make_mock_agent(), session=session, context=context, state=state)
# Verify begin_analyze was called with AnalysisInput containing content_range
mock_cu_client.begin_analyze.assert_called_once()
call_kwargs = mock_cu_client.begin_analyze.call_args
inputs_arg = call_kwargs.kwargs.get("inputs") or call_kwargs[1].get("inputs")
assert inputs_arg is not None
assert len(inputs_arg) == 1
assert isinstance(inputs_arg[0], AnalysisInput)
assert inputs_arg[0].content_range == "1-3"
assert inputs_arg[0].url == "https://example.com/report.pdf"