Working on tests

Taylor · Taylor · commit b0a20d85ad43 · 2025-09-01T10:58:22.000+10:00
diff --git a/tests/conftest.py b/tests/conftest.py
@@ -10,6 +10,7 @@
 import msal
 import pytest
 import pytest_asyncio
+from azure.core.credentials import AzureKeyCredential
 from azure.search.documents.agent.aio import KnowledgeAgentRetrievalClient
 from azure.search.documents.aio import SearchClient
 from azure.search.documents.indexes.aio import SearchIndexClient
@@ -1102,7 +1103,7 @@ def mock_user_directory_client(monkeypatch):
 @pytest.fixture
 def chat_approach():
     return ChatReadRetrieveReadApproach(
-        search_client=None,
+        search_client=SearchClient(endpoint="", index_name="", credential=AzureKeyCredential("")),
         search_index_name=None,
         agent_model=None,
         agent_deployment=None,
@@ -1131,3 +1132,38 @@ def chat_approach():
             credential=MockAzureCredential(),
         ),
     )
+
+
+@pytest.fixture
+def chat_approach_with_hydration():
+    return ChatReadRetrieveReadApproach(
+        search_client=SearchClient(endpoint="", index_name="", credential=AzureKeyCredential("")),
+        search_index_name=None,
+        agent_model=None,
+        agent_deployment=None,
+        agent_client=None,
+        auth_helper=None,
+        openai_client=None,
+        chatgpt_model="gpt-4.1-mini",
+        chatgpt_deployment="chat",
+        embedding_deployment="embeddings",
+        embedding_model=MOCK_EMBEDDING_MODEL_NAME,
+        embedding_dimensions=MOCK_EMBEDDING_DIMENSIONS,
+        embedding_field="embedding3",
+        sourcepage_field="",
+        content_field="",
+        query_language="en-us",
+        query_speller="lexicon",
+        prompt_manager=PromptyManager(),
+        hydrate_references=True,
+        user_blob_manager=AdlsBlobManager(
+            endpoint="https://test-userstorage-account.dfs.core.windows.net",
+            container="test-userstorage-container",
+            credential=MockAzureCredential(),
+        ),
+        global_blob_manager=BlobManager(  # on normal Azure storage
+            endpoint="https://test-globalstorage-account.blob.core.windows.net",
+            container="test-globalstorage-container",
+            credential=MockAzureCredential(),
+        ),
+    )
diff --git a/tests/mocks.py b/tests/mocks.py
@@ -269,7 +269,6 @@ def __init__(self, search_text, vector_queries: Optional[list[VectorQuery]]):
                 ]
             ]
         elif search_text == "hydrated":
-            # Mock search results for hydration testing with complete data
             self.data = [
                 [
                     {
@@ -287,7 +286,6 @@ def __init__(self, search_text, vector_queries: Optional[list[VectorQuery]]):
                 ]
             ]
         elif search_text == "hydrated_multi":
-            # Mock search results for multiple document hydration
             self.data = [
                 [
                     {
@@ -315,7 +313,6 @@ def __init__(self, search_text, vector_queries: Optional[list[VectorQuery]]):
                 ]
             ]
         elif search_text == "hydrated_single":
-            # Mock search results for single document hydration
             self.data = [
                 [
                     {
@@ -458,6 +455,186 @@ def mock_retrieval_response():
     )
 
 
+def mock_retrieval_response_with_sorting():
+    """Mock response with multiple references for testing sorting"""
+    return KnowledgeAgentRetrievalResponse(
+        response=[
+            KnowledgeAgentMessage(
+                role="assistant",
+                content=[KnowledgeAgentMessageTextContent(text="Test response")],
+            )
+        ],
+        activity=[
+            KnowledgeAgentSearchActivityRecord(
+                id=1,
+                target_index="index",
+                query=KnowledgeAgentSearchActivityRecordQuery(search="first query"),
+                count=10,
+                elapsed_ms=50,
+            ),
+            KnowledgeAgentSearchActivityRecord(
+                id=2,
+                target_index="index",
+                query=KnowledgeAgentSearchActivityRecordQuery(search="second query"),
+                count=10,
+                elapsed_ms=50,
+            ),
+        ],
+        references=[
+            KnowledgeAgentAzureSearchDocReference(
+                id="2",  # Higher ID for testing interleaved sorting
+                activity_source=2,
+                doc_key="doc2",
+                source_data={"content": "Content 2", "sourcepage": "page2.pdf"},
+            ),
+            KnowledgeAgentAzureSearchDocReference(
+                id="1",  # Lower ID for testing interleaved sorting
+                activity_source=1,
+                doc_key="doc1",
+                source_data={"content": "Content 1", "sourcepage": "page1.pdf"},
+            ),
+        ],
+    )
+
+
+def mock_retrieval_response_with_duplicates():
+    """Mock response with duplicate doc_keys for testing deduplication"""
+    return KnowledgeAgentRetrievalResponse(
+        response=[
+            KnowledgeAgentMessage(
+                role="assistant",
+                content=[KnowledgeAgentMessageTextContent(text="Test response")],
+            )
+        ],
+        activity=[
+            KnowledgeAgentSearchActivityRecord(
+                id=1,
+                target_index="index",
+                query=KnowledgeAgentSearchActivityRecordQuery(search="query for doc1"),
+                count=10,
+                elapsed_ms=50,
+            ),
+            KnowledgeAgentSearchActivityRecord(
+                id=2,
+                target_index="index",
+                query=KnowledgeAgentSearchActivityRecordQuery(search="another query for doc1"),
+                count=10,
+                elapsed_ms=50,
+            ),
+        ],
+        references=[
+            KnowledgeAgentAzureSearchDocReference(
+                id="1",
+                activity_source=1,
+                doc_key="doc1",  # Same doc_key
+                source_data={"content": "Content 1", "sourcepage": "page1.pdf"},
+            ),
+            KnowledgeAgentAzureSearchDocReference(
+                id="2",
+                activity_source=2,
+                doc_key="doc1",  # Duplicate doc_key
+                source_data={"content": "Content 1", "sourcepage": "page1.pdf"},
+            ),
+            KnowledgeAgentAzureSearchDocReference(
+                id="3",
+                activity_source=1,
+                doc_key="doc2",  # Different doc_key
+                source_data={"content": "Content 2", "sourcepage": "page2.pdf"},
+            ),
+        ],
+    )
+
+
+async def mock_search_for_hydration(*args, **kwargs):
+    """Mock search that returns documents matching the filter"""
+    filter_param = kwargs.get("filter", "")
+
+    # Create documents based on filter - use search_text to distinguish different calls
+    search_text = ""
+    if "doc1" in filter_param and "doc2" in filter_param:
+        search_text = "hydrated_multi"
+    elif "doc1" in filter_param:
+        search_text = "hydrated_single"
+    else:
+        search_text = "hydrated_empty"
+
+    return MockAsyncSearchResultsIterator(search_text, None)
+
+
+def mock_retrieval_response_with_missing_doc_key():
+    """Mock response with missing doc_key to test continue condition"""
+    return KnowledgeAgentRetrievalResponse(
+        response=[
+            KnowledgeAgentMessage(
+                role="assistant",
+                content=[KnowledgeAgentMessageTextContent(text="Test response")],
+            )
+        ],
+        activity=[
+            KnowledgeAgentSearchActivityRecord(
+                id=1,
+                target_index="index",
+                query=KnowledgeAgentSearchActivityRecordQuery(search="query"),
+                count=10,
+                elapsed_ms=50,
+            ),
+        ],
+        references=[
+            KnowledgeAgentAzureSearchDocReference(
+                id="1",
+                activity_source=1,
+                doc_key=None,  # Missing doc_key
+                source_data={"content": "Content 1", "sourcepage": "page1.pdf"},
+            ),
+            KnowledgeAgentAzureSearchDocReference(
+                id="2",
+                activity_source=1,
+                doc_key="",  # Empty doc_key
+                source_data={"content": "Content 2", "sourcepage": "page2.pdf"},
+            ),
+            KnowledgeAgentAzureSearchDocReference(
+                id="3",
+                activity_source=1,
+                doc_key="doc3",  # Valid doc_key
+                source_data={"content": "Content 3", "sourcepage": "page3.pdf"},
+            ),
+        ],
+    )
+
+
+def mock_retrieval_response_with_top_limit():
+    """Mock response with many references to test top limit during document building"""
+    references = []
+    for i in range(15):  # More than any reasonable top limit
+        references.append(
+            KnowledgeAgentAzureSearchDocReference(
+                id=str(i),
+                activity_source=1,
+                doc_key=f"doc{i}",
+                source_data={"content": f"Content {i}", "sourcepage": f"page{i}.pdf"},
+            )
+        )
+
+    return KnowledgeAgentRetrievalResponse(
+        response=[
+            KnowledgeAgentMessage(
+                role="assistant",
+                content=[KnowledgeAgentMessageTextContent(text="Test response")],
+            )
+        ],
+        activity=[
+            KnowledgeAgentSearchActivityRecord(
+                id=1,
+                target_index="index",
+                query=KnowledgeAgentSearchActivityRecordQuery(search="query"),
+                count=10,
+                elapsed_ms=50,
+            ),
+        ],
+        references=references,
+    )
+
+
 class MockAudio:
     def __init__(self, audio_data):
         self.audio_data = audio_data
diff --git a/tests/snapshots/test_app/test_ask_rtr_text_agent/agent_client0/result.json b/tests/snapshots/test_app/test_ask_rtr_text_agent/agent_client0/result.json
@@ -38,7 +38,7 @@
                         "category": null,
                         "content": "There is a whistleblower policy.",
                         "groups": null,
-                        "id": "Benefit_Options-2.pdf",
+                        "id": "file-Benefit_Options_pdf-42656E656669745F4F7074696F6E732E706466-page-2",
                         "images": null,
                         "oids": null,
                         "reranker_score": 3.4577205181121826,
diff --git a/tests/snapshots/test_app/test_ask_rtr_text_agent_filter/agent_auth_client0/result.json b/tests/snapshots/test_app/test_ask_rtr_text_agent_filter/agent_auth_client0/result.json
@@ -38,7 +38,7 @@
                         "category": null,
                         "content": "There is a whistleblower policy.",
                         "groups": null,
-                        "id": "Benefit_Options-2.pdf",
+                        "id": "file-Benefit_Options_pdf-42656E656669745F4F7074696F6E732E706466-page-2",
                         "images": null,
                         "oids": null,
                         "reranker_score": 3.4577205181121826,
diff --git a/tests/snapshots/test_app/test_chat_text_agent/agent_client0/result.json b/tests/snapshots/test_app/test_chat_text_agent/agent_client0/result.json
@@ -39,7 +39,7 @@
                         "category": null,
                         "content": "There is a whistleblower policy.",
                         "groups": null,
-                        "id": "Benefit_Options-2.pdf",
+                        "id": "file-Benefit_Options_pdf-42656E656669745F4F7074696F6E732E706466-page-2",
                         "images": null,
                         "oids": null,
                         "reranker_score": 3.4577205181121826,
diff --git a/tests/snapshots/test_app/test_chat_text_filter_agent/agent_auth_client0/result.json b/tests/snapshots/test_app/test_chat_text_filter_agent/agent_auth_client0/result.json
@@ -39,7 +39,7 @@
                         "category": null,
                         "content": "There is a whistleblower policy.",
                         "groups": null,
-                        "id": "Benefit_Options-2.pdf",
+                        "id": "file-Benefit_Options_pdf-42656E656669745F4F7074696F6E732E706466-page-2",
                         "images": null,
                         "oids": null,
                         "reranker_score": 3.4577205181121826,
diff --git a/tests/test_agentic_retrieval.py b/tests/test_agentic_retrieval.py
diff --git a/tests/test_chatapproach.py b/tests/test_chatapproach.py
