mirror of
https://github.com/tiennm99/DocsGPT.git
synced 2026-10-04 16:13:23 +00:00
PavelandAlex
aaad51f951
* Harden protection with pinned requests and path-param encoding * fix: domain pinning * fix: tests * fix: html test 2 --------- Co-authored-by: Alex <a@tushynski.me>
81 lines
2.7 KiB
Python
81 lines
2.7 KiB
Python
"""Tests for application/agents/tools/read_webpage.py"""
|
|
|
|
from unittest.mock import MagicMock, patch
|
|
|
|
import pytest
|
|
import requests
|
|
|
|
from application.agents.tools.read_webpage import ReadWebpageTool
|
|
|
|
|
|
@pytest.fixture
|
|
def tool():
|
|
return ReadWebpageTool()
|
|
|
|
|
|
@pytest.mark.unit
|
|
class TestReadWebpageExecuteAction:
|
|
def test_unknown_action(self, tool):
|
|
result = tool.execute_action("unknown_action")
|
|
assert "Error" in result
|
|
assert "Unknown action" in result
|
|
|
|
def test_missing_url(self, tool):
|
|
result = tool.execute_action("read_webpage")
|
|
assert "Error" in result
|
|
assert "URL parameter is missing" in result
|
|
|
|
@patch("application.agents.tools.read_webpage.pinned_request")
|
|
def test_successful_fetch(self, mock_pinned_request, tool):
|
|
mock_resp = MagicMock()
|
|
mock_resp.status_code = 200
|
|
mock_resp.text = "<html><body><h1>Title</h1><p>Content</p></body></html>"
|
|
mock_pinned_request.return_value = mock_resp
|
|
|
|
result = tool.execute_action("read_webpage", url="https://example.com")
|
|
|
|
assert "Title" in result
|
|
assert "Content" in result
|
|
|
|
@patch("application.agents.tools.read_webpage.pinned_request")
|
|
def test_request_error(self, mock_pinned_request, tool):
|
|
mock_pinned_request.side_effect = requests.exceptions.ConnectionError("refused")
|
|
|
|
result = tool.execute_action("read_webpage", url="https://example.com")
|
|
|
|
assert "Error fetching URL" in result
|
|
|
|
@patch("application.agents.tools.read_webpage.pinned_request")
|
|
def test_ssrf_blocked(self, mock_pinned_request, tool):
|
|
from application.security.safe_url import UnsafeUserUrlError
|
|
|
|
mock_pinned_request.side_effect = UnsafeUserUrlError("blocked")
|
|
|
|
result = tool.execute_action("read_webpage", url="http://169.254.169.254/")
|
|
|
|
assert "Error" in result
|
|
assert "validation failed" in result
|
|
|
|
@patch("application.agents.tools.read_webpage.pinned_request")
|
|
def test_http_error(self, mock_pinned_request, tool):
|
|
mock_resp = MagicMock()
|
|
mock_resp.raise_for_status.side_effect = requests.exceptions.HTTPError("404")
|
|
mock_pinned_request.return_value = mock_resp
|
|
|
|
result = tool.execute_action("read_webpage", url="https://example.com/404")
|
|
|
|
assert "Error fetching URL" in result
|
|
|
|
|
|
@pytest.mark.unit
|
|
class TestReadWebpageMetadata:
|
|
def test_actions_metadata(self, tool):
|
|
meta = tool.get_actions_metadata()
|
|
assert len(meta) == 1
|
|
assert meta[0]["name"] == "read_webpage"
|
|
assert "url" in meta[0]["parameters"]["properties"]
|
|
assert "url" in meta[0]["parameters"]["required"]
|
|
|
|
def test_config_requirements(self, tool):
|
|
assert tool.get_config_requirements() == {}
|