mirror of
https://github.com/danswer-ai/danswer.git
synced 2025-04-08 03:48:14 +02:00
update tests - additional coverage
This commit is contained in:
parent
a573ba6fb3
commit
ba6ab363f4
@ -3,7 +3,9 @@ from unittest.mock import patch
|
||||
|
||||
import pytest
|
||||
|
||||
from danswer.tools.custom.custom_tool import build_custom_tools_from_openapi_schema
|
||||
from danswer.tools.custom.custom_tool import (
|
||||
build_custom_tools_from_openapi_schema_and_headers,
|
||||
)
|
||||
from danswer.tools.custom.custom_tool import CUSTOM_TOOL_RESPONSE_ID
|
||||
from danswer.tools.custom.custom_tool import CustomToolCallSummary
|
||||
from danswer.tools.custom.custom_tool import validate_openapi_schema
|
||||
@ -77,17 +79,27 @@ class TestCustomTool(unittest.TestCase):
|
||||
Test the GET method of a custom tool.
|
||||
Verifies that the tool correctly constructs the URL and makes the GET request.
|
||||
"""
|
||||
tools = build_custom_tools_from_openapi_schema(
|
||||
tools = build_custom_tools_from_openapi_schema_and_headers(
|
||||
self.openapi_schema, dynamic_schema_info=self.dynamic_schema_info
|
||||
)
|
||||
|
||||
result = list(tools[0].run(assistant_id="123"))
|
||||
expected_url = f"http://localhost:8080/{self.dynamic_schema_info.chat_session_id}/test/{self.dynamic_schema_info.message_id}/assistant/123"
|
||||
mock_request.assert_called_once_with("GET", expected_url, json=None)
|
||||
mock_request.assert_called_once_with("GET", expected_url, json=None, headers={})
|
||||
|
||||
self.assertEqual(len(result), 1)
|
||||
self.assertEqual(result[0].id, CUSTOM_TOOL_RESPONSE_ID)
|
||||
self.assertEqual(result[0].response.tool_name, "getAssistant")
|
||||
self.assertEqual(
|
||||
len(result), 1, "Expected exactly one result from the tool run"
|
||||
)
|
||||
self.assertEqual(
|
||||
result[0].id,
|
||||
CUSTOM_TOOL_RESPONSE_ID,
|
||||
"Tool response ID does not match expected value",
|
||||
)
|
||||
self.assertEqual(
|
||||
result[0].response.tool_name,
|
||||
"getAssistant",
|
||||
"Tool name in response does not match expected value",
|
||||
)
|
||||
|
||||
@patch("danswer.tools.custom.custom_tool.requests.request")
|
||||
def test_custom_tool_run_post(self, mock_request):
|
||||
@ -95,24 +107,95 @@ class TestCustomTool(unittest.TestCase):
|
||||
Test the POST method of a custom tool.
|
||||
Verifies that the tool correctly constructs the URL and makes the POST request with the given body.
|
||||
"""
|
||||
tools = build_custom_tools_from_openapi_schema(
|
||||
tools = build_custom_tools_from_openapi_schema_and_headers(
|
||||
self.openapi_schema, dynamic_schema_info=self.dynamic_schema_info
|
||||
)
|
||||
|
||||
result = list(tools[1].run(assistant_id="456"))
|
||||
expected_url = f"http://localhost:8080/{self.dynamic_schema_info.chat_session_id}/test/{self.dynamic_schema_info.message_id}/assistant/456"
|
||||
mock_request.assert_called_once_with("POST", expected_url, json=None)
|
||||
mock_request.assert_called_once_with(
|
||||
"POST", expected_url, json=None, headers={}
|
||||
)
|
||||
|
||||
self.assertEqual(len(result), 1)
|
||||
self.assertEqual(result[0].id, CUSTOM_TOOL_RESPONSE_ID)
|
||||
self.assertEqual(result[0].response.tool_name, "createAssistant")
|
||||
self.assertEqual(
|
||||
len(result), 1, "Expected exactly one result from the tool run"
|
||||
)
|
||||
self.assertEqual(
|
||||
result[0].id,
|
||||
CUSTOM_TOOL_RESPONSE_ID,
|
||||
"Tool response ID does not match expected value",
|
||||
)
|
||||
self.assertEqual(
|
||||
result[0].response.tool_name,
|
||||
"createAssistant",
|
||||
"Tool name in response does not match expected value",
|
||||
)
|
||||
|
||||
@patch("danswer.tools.custom.custom_tool.requests.request")
|
||||
def test_custom_tool_with_headers(self, mock_request):
|
||||
"""
|
||||
Test the custom tool with custom headers.
|
||||
Verifies that the tool correctly includes the custom headers in the request.
|
||||
"""
|
||||
custom_headers = [
|
||||
{"key": "Authorization", "value": "Bearer token123"},
|
||||
{"key": "Custom-Header", "value": "CustomValue"},
|
||||
]
|
||||
tools = build_custom_tools_from_openapi_schema_and_headers(
|
||||
self.openapi_schema,
|
||||
custom_headers=custom_headers,
|
||||
dynamic_schema_info=self.dynamic_schema_info,
|
||||
)
|
||||
|
||||
list(tools[0].run(assistant_id="123"))
|
||||
expected_url = f"http://localhost:8080/{self.dynamic_schema_info.chat_session_id}/test/{self.dynamic_schema_info.message_id}/assistant/123"
|
||||
expected_headers = {
|
||||
"Authorization": "Bearer token123",
|
||||
"Custom-Header": "CustomValue",
|
||||
}
|
||||
mock_request.assert_called_once_with(
|
||||
"GET", expected_url, json=None, headers=expected_headers
|
||||
)
|
||||
|
||||
@patch("danswer.tools.custom.custom_tool.requests.request")
|
||||
def test_custom_tool_with_empty_headers(self, mock_request):
|
||||
"""
|
||||
Test the custom tool with an empty list of custom headers.
|
||||
Verifies that the tool correctly handles an empty list of headers.
|
||||
"""
|
||||
custom_headers = []
|
||||
tools = build_custom_tools_from_openapi_schema_and_headers(
|
||||
self.openapi_schema,
|
||||
custom_headers=custom_headers,
|
||||
dynamic_schema_info=self.dynamic_schema_info,
|
||||
)
|
||||
|
||||
list(tools[0].run(assistant_id="123"))
|
||||
expected_url = f"http://localhost:8080/{self.dynamic_schema_info.chat_session_id}/test/{self.dynamic_schema_info.message_id}/assistant/123"
|
||||
mock_request.assert_called_once_with("GET", expected_url, json=None, headers={})
|
||||
|
||||
def test_invalid_openapi_schema(self):
|
||||
"""
|
||||
Test that an invalid OpenAPI schema raises a ValueError.
|
||||
"""
|
||||
invalid_schema = {
|
||||
"openapi": "3.0.0",
|
||||
"info": {
|
||||
"version": "1.0.0",
|
||||
"title": "Invalid API",
|
||||
},
|
||||
# Missing required 'paths' key
|
||||
}
|
||||
|
||||
with self.assertRaises(ValueError) as _:
|
||||
validate_openapi_schema(invalid_schema)
|
||||
|
||||
def test_custom_tool_final_result(self):
|
||||
"""
|
||||
Test the final_result method of a custom tool.
|
||||
Verifies that the method correctly extracts and returns the tool result.
|
||||
"""
|
||||
tools = build_custom_tools_from_openapi_schema(
|
||||
tools = build_custom_tools_from_openapi_schema_and_headers(
|
||||
self.openapi_schema, dynamic_schema_info=self.dynamic_schema_info
|
||||
)
|
||||
|
||||
@ -125,7 +208,11 @@ class TestCustomTool(unittest.TestCase):
|
||||
)
|
||||
|
||||
final_result = tools[0].final_result(mock_response)
|
||||
self.assertEqual(final_result, {"id": "789", "name": "Final Assistant"})
|
||||
self.assertEqual(
|
||||
final_result,
|
||||
{"id": "789", "name": "Final Assistant"},
|
||||
"Final result does not match expected output",
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
|
Loading…
x
Reference in New Issue
Block a user