update tests - additional coverage

This commit is contained in:
pablodanswer 2024-09-20 18:11:13 -07:00
parent a573ba6fb3
commit ba6ab363f4

View File

@ -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__":