Skip to content

Commit cb40a58

Browse files
Pavan-MicrosoftPavan-Microsoft
committed
fix: streamline citation handling in test cases by using RunStepToolCallDetails directly
1 parent 5cd6b97 commit cb40a58

File tree

1 file changed

+77
-88
lines changed

1 file changed

+77
-88
lines changed

‎src/tests/api/plugins/test_chat_with_data_plugin.py‎

Lines changed: 77 additions & 88 deletions
Original file line numberDiff line numberDiff line change
@@ -304,26 +304,23 @@ async def test_chat_with_multiple_citations_out_of_order(self, mock_get_agent, c
304304

305305
# Mock run steps with multiple citations (indices 0, 1, 2, 3)
306306
mock_step = MagicMock()
307-
mock_step.type = "tool_calls"
308-
mock_step.step_details = MagicMock(spec=RunStepToolCallDetails)
309-
310-
# Create 4 citations
311-
mock_tool_call = {
312-
'azure_ai_search': {
313-
'output': str({
314-
"metadata": {
315-
"get_urls": [
316-
"https://example.com/doc0",
317-
"https://example.com/doc1",
318-
"https://example.com/doc2",
319-
"https://example.com/doc3"
320-
],
321-
"titles": ["Doc 0", "Doc 1", "Doc 2", "Doc 3"]
322-
}
323-
})
307+
mock_step.step_details = RunStepToolCallDetails(tool_calls=[
308+
{
309+
'azure_ai_search': {
310+
'output': str({
311+
"metadata": {
312+
"get_urls": [
313+
"https://example.com/doc0",
314+
"https://example.com/doc1",
315+
"https://example.com/doc2",
316+
"https://example.com/doc3"
317+
],
318+
"titles": ["Doc 0", "Doc 1", "Doc 2", "Doc 3"]
319+
}
320+
})
321+
}
324322
}
325-
}
326-
mock_step.step_details.tool_calls = [mock_tool_call]
323+
])
327324
mock_client.agents.run_steps.list.return_value = [mock_step]
328325

329326
# Mock messages with out-of-order citation markers: [2], [0], [3], [1]
@@ -377,23 +374,21 @@ async def test_chat_with_out_of_range_citation_markers(self, mock_get_agent, cha
377374

378375
# Mock run steps with only 2 citations (indices 0, 1)
379376
mock_step = MagicMock()
380-
mock_step.type = "tool_calls"
381-
mock_step.step_details = MagicMock(spec=RunStepToolCallDetails)
382-
383-
mock_tool_call = {
384-
'azure_ai_search': {
385-
'output': str({
386-
"metadata": {
387-
"get_urls": [
388-
"https://example.com/valid1",
389-
"https://example.com/valid2"
390-
],
391-
"titles": ["Valid Doc 1", "Valid Doc 2"]
392-
}
393-
})
377+
mock_step.step_details = RunStepToolCallDetails(tool_calls=[
378+
{
379+
'azure_ai_search': {
380+
'output': str({
381+
"metadata": {
382+
"get_urls": [
383+
"https://example.com/valid1",
384+
"https://example.com/valid2"
385+
],
386+
"titles": ["Valid Doc 1", "Valid Doc 2"]
387+
}
388+
})
389+
}
394390
}
395-
}
396-
mock_step.step_details.tool_calls = [mock_tool_call]
391+
])
397392
mock_client.agents.run_steps.list.return_value = [mock_step]
398393

399394
# Mock messages with valid and out-of-range markers: [1] (valid), [5] (invalid), [0] (valid), [10] (invalid)
@@ -440,28 +435,26 @@ async def test_chat_with_unused_citations(self, mock_get_agent, chat_plugin):
440435
mock_run.status = "succeeded"
441436
mock_client.agents.runs.create_and_process.return_value = mock_run
442437

443-
# Mock run steps with 5 citations but only 2 will be used
438+
# Mock run steps with 5 citations but only 3 will be used
444439
mock_step = MagicMock()
445-
mock_step.type = "tool_calls"
446-
mock_step.step_details = MagicMock(spec=RunStepToolCallDetails)
447-
448-
mock_tool_call = {
449-
'azure_ai_search': {
450-
'output': str({
451-
"metadata": {
452-
"get_urls": [
453-
"https://example.com/doc0",
454-
"https://example.com/doc1",
455-
"https://example.com/doc2", # unused
456-
"https://example.com/doc3",
457-
"https://example.com/doc4" # unused
458-
],
459-
"titles": ["Doc 0", "Doc 1", "Doc 2", "Doc 3", "Doc 4"]
460-
}
461-
})
440+
mock_step.step_details = RunStepToolCallDetails(tool_calls=[
441+
{
442+
'azure_ai_search': {
443+
'output': str({
444+
"metadata": {
445+
"get_urls": [
446+
"https://example.com/doc0",
447+
"https://example.com/doc1",
448+
"https://example.com/doc2", # unused
449+
"https://example.com/doc3",
450+
"https://example.com/doc4" # unused
451+
],
452+
"titles": ["Doc 0", "Doc 1", "Doc 2", "Doc 3", "Doc 4"]
453+
}
454+
})
455+
}
462456
}
463-
}
464-
mock_step.step_details.tool_calls = [mock_tool_call]
457+
])
465458
mock_client.agents.run_steps.list.return_value = [mock_step]
466459

467460
# Mock messages with only citations to indices 1, 3, 0
@@ -509,23 +502,21 @@ async def test_chat_with_all_out_of_range_citations_clears_list(self, mock_get_a
509502

510503
# Mock run steps with 2 citations
511504
mock_step = MagicMock()
512-
mock_step.type = "tool_calls"
513-
mock_step.step_details = MagicMock(spec=RunStepToolCallDetails)
514-
515-
mock_tool_call = {
516-
'azure_ai_search': {
517-
'output': str({
518-
"metadata": {
519-
"get_urls": [
520-
"https://example.com/doc0",
521-
"https://example.com/doc1"
522-
],
523-
"titles": ["Doc 0", "Doc 1"]
524-
}
525-
})
505+
mock_step.step_details = RunStepToolCallDetails(tool_calls=[
506+
{
507+
'azure_ai_search': {
508+
'output': str({
509+
"metadata": {
510+
"get_urls": [
511+
"https://example.com/doc0",
512+
"https://example.com/doc1"
513+
],
514+
"titles": ["Doc 0", "Doc 1"]
515+
}
516+
})
517+
}
526518
}
527-
}
528-
mock_step.step_details.tool_calls = [mock_tool_call]
519+
])
529520
mock_client.agents.run_steps.list.return_value = [mock_step]
530521

531522
# Mock messages with only out-of-range markers
@@ -569,24 +560,22 @@ async def test_chat_with_repeated_citation_markers(self, mock_get_agent, chat_pl
569560

570561
# Mock run steps with 3 citations
571562
mock_step = MagicMock()
572-
mock_step.type = "tool_calls"
573-
mock_step.step_details = MagicMock(spec=RunStepToolCallDetails)
574-
575-
mock_tool_call = {
576-
'azure_ai_search': {
577-
'output': str({
578-
"metadata": {
579-
"get_urls": [
580-
"https://example.com/doc0",
581-
"https://example.com/doc1",
582-
"https://example.com/doc2"
583-
],
584-
"titles": ["Doc 0", "Doc 1", "Doc 2"]
585-
}
586-
})
563+
mock_step.step_details = RunStepToolCallDetails(tool_calls=[
564+
{
565+
'azure_ai_search': {
566+
'output': str({
567+
"metadata": {
568+
"get_urls": [
569+
"https://example.com/doc0",
570+
"https://example.com/doc1",
571+
"https://example.com/doc2"
572+
],
573+
"titles": ["Doc 0", "Doc 1", "Doc 2"]
574+
}
575+
})
576+
}
587577
}
588-
}
589-
mock_step.step_details.tool_calls = [mock_tool_call]
578+
])
590579
mock_client.agents.run_steps.list.return_value = [mock_step]
591580

592581
# Mock messages with repeated markers: [1], [2], [1] again, [0], [1] again

0 commit comments

Comments
 (0)