tiennm99/DocsGPT
1
0 Fork 0
mirror of https://github.com/tiennm99/DocsGPT.git synced 2026-10-12 22:14:05 +00:00
Code Issues Packages Projects Releases Wiki Activity

Harden protection with pinned requests and path-param encoding (#2486)

Browse Source 
* 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>
This commit is contained in:
PavelandAlex authored and GitHub committed 2026-05-23 02:33:29 +01:00
1 parent 6cfdc4fbae
commit aaad51f951
20 files changed
+623 -723
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

application/agents/tools/api_tool.py
+22 -65
View File
@@ -2,7 +2,7 @@ import json
import logging
import re
from typing import Any, Dict, Optional
from urllib.parse import urlencode
from urllib.parse import quote, urlencode
import requests
@@ -11,7 +11,7 @@ from application.agents.tools.api_body_serializer import (
RequestBodySerializer,
)
from application.agents.tools.base import Tool
from application.core.url_validation import validate_url, SSRFError
from application.security.safe_url import UnsafeUserUrlError, pinned_request
logger = logging.getLogger(__name__)
@@ -70,18 +70,16 @@ class APITool(Tool):
Returns:
Dict with status_code, data, and message
"""
_VALID_METHODS = {"GET", "POST", "PUT", "DELETE", "PATCH", "HEAD", "OPTIONS"}
request_url = url
request_headers = headers.copy() if headers else {}
response = None
# Validate URL to prevent SSRF attacks
try:
validate_url(request_url)
except SSRFError as e:
logger.error(f"URL validation failed: {e}")
if method.upper() not in _VALID_METHODS:
return {
"status_code": None,
"message": f"URL validation error: {e}",
"message": f"Unsupported HTTP method: {method}",
"data": None,
}
@@ -91,8 +89,9 @@ class APITool(Tool):
for match in re.finditer(r"\{([^}]+)\}", request_url):
param_name = match.group(1)
if param_name in query_params:
safe_value = quote(str(query_params[param_name]), safe="")
request_url = request_url.replace(
f"{{{param_name}}}", str(query_params[param_name])
f"{{{param_name}}}", safe_value
)
path_params_used.add(param_name)
remaining_params = {
@@ -103,19 +102,6 @@ class APITool(Tool):
separator = "&" if "?" in request_url else "?"
request_url = f"{request_url}{separator}{query_string}"
# Re-validate URL after parameter substitution to prevent SSRF via path params
try:
validate_url(request_url)
except SSRFError as e:
logger.error(f"URL validation failed after parameter substitution: {e}")
return {
"status_code": None,
"message": f"URL validation error: {e}",
"data": None,
}
# Serialize body based on content type
if body and body != {}:
try:
serialized_body, body_headers = RequestBodySerializer.serialize(
@@ -141,49 +127,13 @@ class APITool(Tool):
f"API Call: {method} {request_url} | Content-Type: {request_headers.get('Content-Type', 'N/A')}"
)
if method.upper() == "GET":
response = requests.get(
request_url, headers=request_headers, timeout=DEFAULT_TIMEOUT
)
elif method.upper() == "POST":
response = requests.post(
request_url,
data=serialized_body,
headers=request_headers,
timeout=DEFAULT_TIMEOUT,
)
elif method.upper() == "PUT":
response = requests.put(
request_url,
data=serialized_body,
headers=request_headers,
timeout=DEFAULT_TIMEOUT,
)
elif method.upper() == "DELETE":
response = requests.delete(
request_url, headers=request_headers, timeout=DEFAULT_TIMEOUT
)
elif method.upper() == "PATCH":
response = requests.patch(
request_url,
data=serialized_body,
headers=request_headers,
timeout=DEFAULT_TIMEOUT,
)
elif method.upper() == "HEAD":
response = requests.head(
request_url, headers=request_headers, timeout=DEFAULT_TIMEOUT
)
elif method.upper() == "OPTIONS":
response = requests.options(
request_url, headers=request_headers, timeout=DEFAULT_TIMEOUT
)
else:
return {
"status_code": None,
"message": f"Unsupported HTTP method: {method}",
"data": None,
}
response = pinned_request(
method,
request_url,
data=serialized_body,
headers=request_headers,
timeout=DEFAULT_TIMEOUT,
)
response.raise_for_status()
data = self._parse_response(response)
@@ -193,6 +143,13 @@ class APITool(Tool):
"data": data,
"message": "API call successful.",
}
except UnsafeUserUrlError as e:
logger.error(f"URL validation failed: {e}")
return {
"status_code": None,
"message": f"URL validation error: {e}",
"data": None,
}
except requests.exceptions.Timeout:
logger.error(f"Request timeout for {request_url}")
return {
application/agents/tools/ntfy.py
+7 -2
View File
@@ -1,5 +1,5 @@
import requests
from application.agents.tools.base import Tool
from application.security.safe_url import UnsafeUserUrlError, pinned_request
class NtfyTool(Tool):
"""
@@ -71,7 +71,12 @@ class NtfyTool(Tool):
if self.token:
headers["Authorization"] = f"Basic {self.token}"
data = message.encode("utf-8")
response = requests.post(url, headers=headers, data=data, timeout=100)
try:
response = pinned_request(
"POST", url, data=data, headers=headers, timeout=100,
)
except UnsafeUserUrlError as e:
return {"status_code": None, "message": f"URL validation error: {e}"}
return {"status_code": response.status_code, "message": "Message sent"}
def get_actions_metadata(self):
application/agents/tools/read_webpage.py
+12 -17
View File
@@ -1,7 +1,6 @@
import requests
from markdownify import markdownify
from application.agents.tools.base import Tool
from application.core.url_validation import validate_url, SSRFError
from application.security.safe_url import UnsafeUserUrlError, pinned_request
class ReadWebpageTool(Tool):
"""
@@ -31,28 +30,24 @@ class ReadWebpageTool(Tool):
if not url:
return "Error: URL parameter is missing."
# Validate URL to prevent SSRF attacks
try:
url = validate_url(url)
except SSRFError as e:
return f"Error: URL validation failed - {e}"
response = pinned_request(
"GET",
url,
headers={'User-Agent': 'DocsGPT-Agent/1.0'},
timeout=10,
)
response.raise_for_status()
try:
response = requests.get(url, timeout=10, headers={'User-Agent': 'DocsGPT-Agent/1.0'})
response.raise_for_status() # Raise an exception for HTTP errors (4xx or 5xx)
html_content = response.text
#soup = BeautifulSoup(html_content, 'html.parser')
markdown_content = markdownify(html_content, heading_style="ATX", newline_style="BACKSLASH")
return markdown_content
except requests.exceptions.RequestException as e:
return f"Error fetching URL {url}: {e}"
except UnsafeUserUrlError as e:
return f"Error: URL validation failed - {e}"
except Exception as e:
return f"Error processing URL {url}: {e}"
return f"Error fetching URL {url}: {e}"
def get_actions_metadata(self):
"""
application/parser/remote/crawler_loader.py
+14 -25
View File
@@ -1,16 +1,16 @@
import logging
import os
import requests
from urllib.parse import urlparse, urljoin
from bs4 import BeautifulSoup
from urllib.parse import urljoin, urlparse
from application.parser.remote.base import BaseRemote
from application.parser.schema.base import Document
from application.core.url_validation import validate_url, SSRFError
from langchain_community.document_loaders import WebBaseLoader
from application.security.safe_url import pinned_request
class CrawlerLoader(BaseRemote):
def __init__(self, limit=10):
self.loader = WebBaseLoader # Initialize the document loader
self.limit = limit # Set the limit for the number of pages to scrape
def load_data(self, inputs):
@@ -35,34 +35,23 @@ class CrawlerLoader(BaseRemote):
visited_urls.add(current_url)
try:
# Validate each URL before making requests
try:
validate_url(current_url)
except SSRFError as e:
logging.warning(f"Skipping URL due to validation failure: {current_url} - {e}")
continue
response = requests.get(current_url, timeout=30)
response = pinned_request("GET", current_url, timeout=30)
response.raise_for_status()
loader = self.loader([current_url])
docs = loader.load()
# Convert the loaded documents to your Document schema
for doc in docs:
metadata = dict(doc.metadata or {})
source_url = metadata.get("source") or current_url
metadata["file_path"] = self._url_to_virtual_path(source_url)
loaded_content.append(
Document(
doc.page_content,
extra_info=metadata
)
soup = BeautifulSoup(response.text, "html.parser")
loaded_content.append(
Document(
soup.get_text(separator="\n", strip=True),
extra_info={
"source": current_url,
"file_path": self._url_to_virtual_path(current_url),
},
)
)
except Exception as e:
logging.error(f"Error processing URL {current_url}: {e}", exc_info=True)
continue
# Parse the HTML content to extract all links
soup = BeautifulSoup(response.text, 'html.parser')
all_links = [
urljoin(current_url, a['href'])
for a in soup.find_all('a', href=True)
application/parser/remote/crawler_markdown.py
+4 -7
View File
@@ -1,8 +1,8 @@
import requests
from urllib.parse import urlparse, urljoin
from bs4 import BeautifulSoup
from application.parser.remote.base import BaseRemote
from application.core.url_validation import validate_url, SSRFError
from application.security.safe_url import UnsafeUserUrlError, pinned_request
import re
from markdownify import markdownify
from application.parser.schema.base import Document
@@ -20,7 +20,6 @@ class CrawlerLoader(BaseRemote):
"""
self.limit = limit
self.allow_subdomains = allow_subdomains
self.session = requests.Session()
def load_data(self, inputs):
url = inputs
@@ -91,15 +90,13 @@ class CrawlerLoader(BaseRemote):
def _fetch_page(self, url):
try:
# Validate URL before fetching to prevent SSRF
validate_url(url)
response = self.session.get(url, timeout=10)
response = pinned_request("GET", url, timeout=10)
response.raise_for_status()
return response.text
except SSRFError as e:
except UnsafeUserUrlError as e:
print(f"URL validation failed for {url}: {e}")
return None
except requests.exceptions.RequestException as e:
except Exception as e:
print(f"Error fetching URL {url}: {e}")
return None
application/parser/remote/sitemap_loader.py
+21 -19
View File
@@ -1,14 +1,16 @@
import logging
import requests
import re # Import regular expression library
import re
import defusedxml.ElementTree as ET
from bs4 import BeautifulSoup
from application.parser.remote.base import BaseRemote
from application.parser.schema.base import Document
from application.core.url_validation import validate_url, SSRFError
from application.security.safe_url import UnsafeUserUrlError, pinned_request
class SitemapLoader(BaseRemote):
def __init__(self, limit=20):
from langchain_community.document_loaders import WebBaseLoader
self.loader = WebBaseLoader
self.limit = limit # Adding limit to control the number of URLs to process
def load_data(self, inputs):
@@ -42,8 +44,15 @@ class SitemapLoader(BaseRemote):
logging.error(f"URL validation failed for sitemap entry {url}: {e}")
continue
try:
loader = self.loader([url])
documents.extend(loader.load())
response = pinned_request("GET", url, timeout=30)
response.raise_for_status()
soup = BeautifulSoup(response.text, "html.parser")
documents.append(
Document(
soup.get_text(separator="\n", strip=True),
extra_info={"source": url},
)
)
processed_urls += 1 # Increment the counter after processing each URL
except Exception as e:
logging.error(f"Error processing URL {url}: {e}", exc_info=True)
@@ -53,14 +62,12 @@ class SitemapLoader(BaseRemote):
def _extract_urls(self, sitemap_url):
try:
# Validate URL before fetching to prevent SSRF
validate_url(sitemap_url)
response = requests.get(sitemap_url, timeout=30)
response.raise_for_status() # Raise an exception for HTTP errors
except SSRFError as e:
response = pinned_request("GET", sitemap_url, timeout=30)
response.raise_for_status()
except UnsafeUserUrlError as e:
print(f"URL validation failed for sitemap: {sitemap_url}. Error: {e}")
return []
except (requests.exceptions.HTTPError, requests.exceptions.ConnectionError) as e:
except Exception as e:
print(f"Failed to fetch sitemap: {sitemap_url}. Error: {e}")
return []
@@ -90,6 +97,8 @@ class SitemapLoader(BaseRemote):
urls = []
for loc in root.findall('.//url/loc'):
if not loc.text:
continue
urls.append(loc.text)
# Check for nested sitemaps
@@ -97,13 +106,6 @@ class SitemapLoader(BaseRemote):
nested_sitemap_url = sitemap.text
if not nested_sitemap_url:
continue
try:
nested_sitemap_url = validate_url(nested_sitemap_url)
except SSRFError as e:
logging.error(
f"URL validation failed for nested sitemap {nested_sitemap_url}: {e}"
)
continue
urls.extend(self._extract_urls(nested_sitemap_url))
return urls
application/parser/remote/web_loader.py
+20 -12
View File
@@ -1,8 +1,11 @@
import logging
from bs4 import BeautifulSoup
from application.core.url_validation import SSRFError, validate_url
from application.parser.remote.base import BaseRemote
from application.parser.schema.base import Document
from langchain_community.document_loaders import WebBaseLoader
from application.security.safe_url import pinned_request
headers = {
"User-Agent": "Mozilla/5.0",
@@ -17,9 +20,6 @@ headers = {
class WebLoader(BaseRemote):
def __init__(self):
self.loader = WebBaseLoader
def load_data(self, inputs):
urls = inputs
if isinstance(urls, str):
@@ -34,15 +34,23 @@ class WebLoader(BaseRemote):
)
continue
try:
loader = self.loader([url], header_template=headers)
loaded_docs = loader.load()
for doc in loaded_docs:
documents.append(
Document(
doc.page_content,
extra_info=doc.metadata,
)
response = pinned_request("GET", url, headers=headers, timeout=30)
response.raise_for_status()
soup = BeautifulSoup(response.text, "html.parser")
metadata = {"source": url}
if soup.title:
title = soup.title.get_text(strip=True)
if title:
metadata["title"] = title
html_tag = soup.find("html")
if html_tag and html_tag.get("lang"):
metadata["language"] = html_tag.get("lang")
documents.append(
Document(
soup.get_text(separator="\n", strip=True),
extra_info=metadata,
)
)
except Exception as e:
logging.error(f"Error processing URL {url}: {e}", exc_info=True)
continue
application/security/safe_url.py
+64 -25
View File
@@ -291,6 +291,63 @@ def _ip_to_url_host(ip: ipaddress.IPv4Address | ipaddress.IPv6Address) -> str:
return str(ip)
def _url_userinfo_prefix(netloc: str) -> str:
"""Return the exact ``userinfo@`` prefix from a URL netloc, if present."""
if "@" not in netloc:
return ""
return f"{netloc.rsplit('@', 1)[0]}@"
def pinned_request(
method: str,
url: str,
*,
data: Any = None,
json: Any = None,
headers: dict[str, str] | None = None,
timeout: float = 90.0,
allow_redirects: bool = False,
) -> requests.Response:
"""Send an HTTP request with the connection pinned to a validated IP,
closing the DNS-rebinding TOCTOU window left by the naive
validate-then-``requests`` pattern.
Raises:
UnsafeUserUrlError: If the URL fails the SSRF guard.
requests.RequestException: For network-level failures.
"""
host, ip, parts = _validate_and_pick_ip(url)
netloc = f"{_url_userinfo_prefix(parts.netloc)}{_ip_to_url_host(ip)}"
if parts.port is not None:
netloc = f"{netloc}:{parts.port}"
pinned_url = urlunsplit(
(parts.scheme, netloc, parts.path, parts.query, parts.fragment)
)
request_headers = dict(headers or {})
host_header = host if parts.port is None else f"{host}:{parts.port}"
request_headers["Host"] = host_header
session = requests.Session()
if parts.scheme == "https":
session.mount("https://", _PinnedHostAdapter(host))
try:
return session.request(
method=method.upper(),
url=pinned_url,
data=data,
json=json,
headers=request_headers,
timeout=timeout,
allow_redirects=allow_redirects,
)
finally:
session.close()
def pinned_post(
url: str,
*,
@@ -328,33 +385,15 @@ def pinned_post(
requests.RequestException: For network-level failures.
"""
host, ip, parts = _validate_and_pick_ip(url)
netloc = _ip_to_url_host(ip)
if parts.port is not None:
netloc = f"{netloc}:{parts.port}"
pinned_url = urlunsplit(
(parts.scheme, netloc, parts.path, parts.query, parts.fragment)
return pinned_request(
"POST",
url,
json=json,
headers=headers,
timeout=timeout,
allow_redirects=allow_redirects,
)
request_headers = dict(headers or {})
host_header = host if parts.port is None else f"{host}:{parts.port}"
request_headers["Host"] = host_header
session = requests.Session()
if parts.scheme == "https":
session.mount("https://", _PinnedHostAdapter(host))
try:
return session.post(
pinned_url,
json=json,
headers=request_headers,
timeout=timeout,
allow_redirects=allow_redirects,
)
finally:
session.close()
class _PinnedHTTPSTransport(httpx.HTTPTransport):
"""``httpx`` transport pinned to a single validated IP literal.
tests/agents/test_api_tool.py
+39 -51
View File
@@ -45,15 +45,14 @@ class TestAPIToolInit:
@pytest.mark.unit
class TestMakeApiCall:
@patch("application.agents.tools.api_tool.validate_url")
@patch("application.agents.tools.api_tool.requests.get")
def test_successful_get(self, mock_get, mock_validate, tool):
@patch("application.agents.tools.api_tool.pinned_request")
def test_successful_get(self, mock_pinned, tool):
mock_resp = MagicMock()
mock_resp.status_code = 200
mock_resp.headers = {"Content-Type": "application/json"}
mock_resp.json.return_value = {"result": "ok"}
mock_resp.content = b'{"result":"ok"}'
mock_get.return_value = mock_resp
mock_pinned.return_value = mock_resp
result = tool.execute_action("any_action")
@@ -61,54 +60,50 @@ class TestMakeApiCall:
assert result["data"] == {"result": "ok"}
assert result["message"] == "API call successful."
@patch("application.agents.tools.api_tool.validate_url")
@patch("application.agents.tools.api_tool.requests.post")
def test_successful_post(self, mock_post, mock_validate, post_tool):
@patch("application.agents.tools.api_tool.pinned_request")
def test_successful_post(self, mock_pinned, post_tool):
mock_resp = MagicMock()
mock_resp.status_code = 201
mock_resp.headers = {"Content-Type": "application/json"}
mock_resp.json.return_value = {"id": 1}
mock_resp.content = b'{"id":1}'
mock_post.return_value = mock_resp
mock_pinned.return_value = mock_resp
result = post_tool.execute_action("create", name="test")
assert result["status_code"] == 201
@patch("application.agents.tools.api_tool.validate_url")
def test_ssrf_blocked(self, mock_validate, tool):
from application.core.url_validation import SSRFError
@patch("application.agents.tools.api_tool.pinned_request")
def test_ssrf_blocked(self, mock_pinned, tool):
from application.security.safe_url import UnsafeUserUrlError
mock_validate.side_effect = SSRFError("blocked")
mock_pinned.side_effect = UnsafeUserUrlError("blocked")
result = tool.execute_action("any")
assert result["status_code"] is None
assert "URL validation error" in result["message"]
@patch("application.agents.tools.api_tool.validate_url")
@patch("application.agents.tools.api_tool.requests.get")
def test_timeout_error(self, mock_get, mock_validate, tool):
mock_get.side_effect = requests.exceptions.Timeout()
@patch("application.agents.tools.api_tool.pinned_request")
def test_timeout_error(self, mock_pinned, tool):
mock_pinned.side_effect = requests.exceptions.Timeout()
result = tool.execute_action("any")
assert result["status_code"] is None
assert "timeout" in result["message"].lower()
@patch("application.agents.tools.api_tool.validate_url")
@patch("application.agents.tools.api_tool.requests.get")
def test_connection_error(self, mock_get, mock_validate, tool):
mock_get.side_effect = requests.exceptions.ConnectionError("refused")
@patch("application.agents.tools.api_tool.pinned_request")
def test_connection_error(self, mock_pinned, tool):
mock_pinned.side_effect = requests.exceptions.ConnectionError("refused")
result = tool.execute_action("any")
assert result["status_code"] is None
assert "Connection error" in result["message"]
@patch("application.agents.tools.api_tool.validate_url")
@patch("application.agents.tools.api_tool.requests.get")
def test_http_error(self, mock_get, mock_validate, tool):
@patch("application.agents.tools.api_tool.pinned_request")
def test_http_error(self, mock_pinned, tool):
mock_resp = MagicMock()
mock_resp.status_code = 404
mock_resp.text = "Not Found"
@@ -116,15 +111,14 @@ class TestMakeApiCall:
mock_resp.raise_for_status.side_effect = requests.exceptions.HTTPError(
response=mock_resp
)
mock_get.return_value = mock_resp
mock_pinned.return_value = mock_resp
result = tool.execute_action("any")
assert result["status_code"] == 404
assert "HTTP Error" in result["message"]
@patch("application.agents.tools.api_tool.validate_url")
def test_unsupported_method(self, mock_validate):
def test_unsupported_method(self):
tool = APITool(
config={"url": "https://example.com", "method": "CUSTOM"}
)
@@ -132,69 +126,64 @@ class TestMakeApiCall:
assert result["status_code"] is None
assert "Unsupported" in result["message"]
@patch("application.agents.tools.api_tool.validate_url")
@patch("application.agents.tools.api_tool.requests.put")
def test_put_method(self, mock_put, mock_validate):
@patch("application.agents.tools.api_tool.pinned_request")
def test_put_method(self, mock_pinned):
tool = APITool(config={"url": "https://example.com/item/1", "method": "PUT"})
mock_resp = MagicMock()
mock_resp.status_code = 200
mock_resp.headers = {"Content-Type": "application/json"}
mock_resp.json.return_value = {}
mock_resp.content = b'{}'
mock_put.return_value = mock_resp
mock_pinned.return_value = mock_resp
result = tool.execute_action("update", name="new")
assert result["status_code"] == 200
@patch("application.agents.tools.api_tool.validate_url")
@patch("application.agents.tools.api_tool.requests.delete")
def test_delete_method(self, mock_delete, mock_validate):
@patch("application.agents.tools.api_tool.pinned_request")
def test_delete_method(self, mock_pinned):
tool = APITool(config={"url": "https://example.com/item/1", "method": "DELETE"})
mock_resp = MagicMock()
mock_resp.status_code = 204
mock_resp.headers = {"Content-Type": "text/plain"}
mock_resp.content = b''
mock_delete.return_value = mock_resp
mock_pinned.return_value = mock_resp
result = tool.execute_action("delete")
assert result["status_code"] == 204
@patch("application.agents.tools.api_tool.validate_url")
@patch("application.agents.tools.api_tool.requests.patch")
def test_patch_method(self, mock_patch, mock_validate):
@patch("application.agents.tools.api_tool.pinned_request")
def test_patch_method(self, mock_pinned):
tool = APITool(config={"url": "https://example.com/item/1", "method": "PATCH"})
mock_resp = MagicMock()
mock_resp.status_code = 200
mock_resp.headers = {"Content-Type": "application/json"}
mock_resp.json.return_value = {"patched": True}
mock_resp.content = b'{"patched":true}'
mock_patch.return_value = mock_resp
mock_pinned.return_value = mock_resp
result = tool.execute_action("patch", field="val")
assert result["status_code"] == 200
@patch("application.agents.tools.api_tool.validate_url")
@patch("application.agents.tools.api_tool.requests.head")
def test_head_method(self, mock_head, mock_validate):
@patch("application.agents.tools.api_tool.pinned_request")
def test_head_method(self, mock_pinned):
tool = APITool(config={"url": "https://example.com", "method": "HEAD"})
mock_resp = MagicMock()
mock_resp.status_code = 200
mock_resp.headers = {"Content-Type": "text/html"}
mock_resp.content = b''
mock_head.return_value = mock_resp
mock_pinned.return_value = mock_resp
result = tool.execute_action("check")
assert result["status_code"] == 200
@patch("application.agents.tools.api_tool.validate_url")
@patch("application.agents.tools.api_tool.requests.options")
def test_options_method(self, mock_options, mock_validate):
@patch("application.agents.tools.api_tool.pinned_request")
def test_options_method(self, mock_pinned):
tool = APITool(config={"url": "https://example.com", "method": "OPTIONS"})
mock_resp = MagicMock()
mock_resp.status_code = 200
mock_resp.headers = {"Content-Type": "text/plain"}
mock_resp.content = b''
mock_options.return_value = mock_resp
mock_pinned.return_value = mock_resp
result = tool.execute_action("options")
assert result["status_code"] == 200
@@ -202,9 +191,8 @@ class TestMakeApiCall:
@pytest.mark.unit
class TestPathParamSubstitution:
@patch("application.agents.tools.api_tool.validate_url")
@patch("application.agents.tools.api_tool.requests.get")
def test_path_params_substituted(self, mock_get, mock_validate):
@patch("application.agents.tools.api_tool.pinned_request")
def test_path_params_substituted(self, mock_pinned):
tool = APITool(
config={
"url": "https://api.example.com/users/{user_id}/posts/{post_id}",
@@ -217,11 +205,11 @@ class TestPathParamSubstitution:
mock_resp.headers = {"Content-Type": "application/json"}
mock_resp.json.return_value = []
mock_resp.content = b'[]'
mock_get.return_value = mock_resp
mock_pinned.return_value = mock_resp
tool.execute_action("get")
called_url = mock_get.call_args[0][0]
called_url = mock_pinned.call_args[0][1]
assert "/users/42/posts/7" in called_url
assert "{user_id}" not in called_url
tests/agents/test_ntfy_tool.py
+24 -23
View File
@@ -23,11 +23,11 @@ class TestNtfyExecuteAction:
with pytest.raises(ValueError, match="Unknown action"):
tool.execute_action("bad_action")
@patch("application.agents.tools.ntfy.requests.post")
def test_send_message_basic(self, mock_post, tool):
@patch("application.agents.tools.ntfy.pinned_request")
def test_send_message_basic(self, mock_pinned_request, tool):
mock_resp = MagicMock()
mock_resp.status_code = 200
mock_post.return_value = mock_resp
mock_pinned_request.return_value = mock_resp
result = tool.execute_action(
"ntfy_send_message",
@@ -38,16 +38,17 @@ class TestNtfyExecuteAction:
assert result["status_code"] == 200
assert result["message"] == "Message sent"
mock_post.assert_called_once()
call_args = mock_post.call_args
assert call_args[0][0] == "https://ntfy.sh/test"
assert call_args[1]["data"] == b"Hello"
mock_pinned_request.assert_called_once()
args, kwargs = mock_pinned_request.call_args
assert args[0] == "POST"
assert args[1] == "https://ntfy.sh/test"
assert kwargs["data"] == b"Hello"
@patch("application.agents.tools.ntfy.requests.post")
def test_send_with_title_and_priority(self, mock_post, tool):
@patch("application.agents.tools.ntfy.pinned_request")
def test_send_with_title_and_priority(self, mock_pinned_request, tool):
mock_resp = MagicMock()
mock_resp.status_code = 200
mock_post.return_value = mock_resp
mock_pinned_request.return_value = mock_resp
tool.execute_action(
"ntfy_send_message",
@@ -58,15 +59,15 @@ class TestNtfyExecuteAction:
priority=5,
)
headers = mock_post.call_args[1]["headers"]
headers = mock_pinned_request.call_args[1]["headers"]
assert headers["X-Title"] == "Warning"
assert headers["X-Priority"] == "5"
@patch("application.agents.tools.ntfy.requests.post")
def test_auth_header_with_token(self, mock_post, tool):
@patch("application.agents.tools.ntfy.pinned_request")
def test_auth_header_with_token(self, mock_pinned_request, tool):
mock_resp = MagicMock()
mock_resp.status_code = 200
mock_post.return_value = mock_resp
mock_pinned_request.return_value = mock_resp
tool.execute_action(
"ntfy_send_message",
@@ -75,14 +76,14 @@ class TestNtfyExecuteAction:
topic="t",
)
headers = mock_post.call_args[1]["headers"]
headers = mock_pinned_request.call_args[1]["headers"]
assert headers["Authorization"] == "Basic test_token"
@patch("application.agents.tools.ntfy.requests.post")
def test_no_auth_without_token(self, mock_post, tool_no_token):
@patch("application.agents.tools.ntfy.pinned_request")
def test_no_auth_without_token(self, mock_pinned_request, tool_no_token):
mock_resp = MagicMock()
mock_resp.status_code = 200
mock_post.return_value = mock_resp
mock_pinned_request.return_value = mock_resp
tool_no_token.execute_action(
"ntfy_send_message",
@@ -91,7 +92,7 @@ class TestNtfyExecuteAction:
topic="t",
)
headers = mock_post.call_args[1]["headers"]
headers = mock_pinned_request.call_args[1]["headers"]
assert "Authorization" not in headers
def test_invalid_priority_raises(self, tool):
@@ -114,11 +115,11 @@ class TestNtfyExecuteAction:
priority="abc",
)
@patch("application.agents.tools.ntfy.requests.post")
def test_trailing_slash_stripped(self, mock_post, tool):
@patch("application.agents.tools.ntfy.pinned_request")
def test_trailing_slash_stripped(self, mock_pinned_request, tool):
mock_resp = MagicMock()
mock_resp.status_code = 200
mock_post.return_value = mock_resp
mock_pinned_request.return_value = mock_resp
tool.execute_action(
"ntfy_send_message",
@@ -127,7 +128,7 @@ class TestNtfyExecuteAction:
topic="test",
)
assert mock_post.call_args[0][0] == "https://ntfy.sh/test"
assert mock_pinned_request.call_args[0][1] == "https://ntfy.sh/test"
@pytest.mark.unit
tests/agents/test_read_webpage_tool.py
+13 -19
View File
@@ -25,48 +25,42 @@ class TestReadWebpageExecuteAction:
assert "Error" in result
assert "URL parameter is missing" in result
@patch("application.agents.tools.read_webpage.validate_url")
@patch("application.agents.tools.read_webpage.requests.get")
def test_successful_fetch(self, mock_get, mock_validate, tool):
mock_validate.return_value = "https://example.com"
@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_get.return_value = mock_resp
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.validate_url")
@patch("application.agents.tools.read_webpage.requests.get")
def test_request_error(self, mock_get, mock_validate, tool):
mock_validate.return_value = "https://example.com"
mock_get.side_effect = requests.exceptions.ConnectionError("refused")
@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.validate_url")
def test_ssrf_blocked(self, mock_validate, tool):
from application.core.url_validation import SSRFError
@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_validate.side_effect = SSRFError("blocked")
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.validate_url")
@patch("application.agents.tools.read_webpage.requests.get")
def test_http_error(self, mock_get, mock_validate, tool):
mock_validate.return_value = "https://example.com/404"
@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_get.return_value = mock_resp
mock_pinned_request.return_value = mock_resp
result = tool.execute_action("read_webpage", url="https://example.com/404")
tests/agents/tools/test_api_tool.py
+115 -101
View File
@@ -81,104 +81,103 @@ class TestAPIToolInit:
@pytest.mark.unit
class TestMakeApiCall:
@patch("application.agents.tools.api_tool.validate_url")
@patch("application.agents.tools.api_tool.requests.get")
def test_successful_get(self, mock_get, mock_validate, get_tool):
@patch("application.agents.tools.api_tool.pinned_request")
def test_successful_get(self, mock_pinned, get_tool):
mock_resp = MagicMock()
mock_resp.status_code = 200
mock_resp.headers = {"Content-Type": "application/json"}
mock_resp.json.return_value = {"result": "ok"}
mock_resp.content = b'{"result":"ok"}'
mock_get.return_value = mock_resp
mock_pinned.return_value = mock_resp
result = get_tool.execute_action("any_action")
assert result["status_code"] == 200
assert result["data"] == {"result": "ok"}
assert result["message"] == "API call successful."
assert mock_pinned.call_args[0][0] == "GET"
@patch("application.agents.tools.api_tool.validate_url")
@patch("application.agents.tools.api_tool.requests.post")
def test_successful_post(self, mock_post, mock_validate, post_tool):
@patch("application.agents.tools.api_tool.pinned_request")
def test_successful_post(self, mock_pinned, post_tool):
mock_resp = MagicMock()
mock_resp.status_code = 201
mock_resp.headers = {"Content-Type": "application/json"}
mock_resp.json.return_value = {"id": 1}
mock_resp.content = b'{"id":1}'
mock_post.return_value = mock_resp
mock_pinned.return_value = mock_resp
result = post_tool.execute_action("create", name="test")
assert result["status_code"] == 201
assert mock_pinned.call_args[0][0] == "POST"
@patch("application.agents.tools.api_tool.validate_url")
@patch("application.agents.tools.api_tool.requests.put")
def test_put_method(self, mock_put, mock_validate):
@patch("application.agents.tools.api_tool.pinned_request")
def test_put_method(self, mock_pinned):
tool = APITool(config={"url": "https://example.com/item/1", "method": "PUT"})
mock_resp = MagicMock()
mock_resp.status_code = 200
mock_resp.headers = {"Content-Type": "application/json"}
mock_resp.json.return_value = {}
mock_resp.content = b'{}'
mock_put.return_value = mock_resp
mock_pinned.return_value = mock_resp
result = tool.execute_action("update", name="new")
assert result["status_code"] == 200
assert mock_pinned.call_args[0][0] == "PUT"
@patch("application.agents.tools.api_tool.validate_url")
@patch("application.agents.tools.api_tool.requests.delete")
def test_delete_method(self, mock_delete, mock_validate):
@patch("application.agents.tools.api_tool.pinned_request")
def test_delete_method(self, mock_pinned):
tool = APITool(config={"url": "https://example.com/item/1", "method": "DELETE"})
mock_resp = MagicMock()
mock_resp.status_code = 204
mock_resp.headers = {"Content-Type": "text/plain"}
mock_resp.content = b''
mock_delete.return_value = mock_resp
mock_pinned.return_value = mock_resp
result = tool.execute_action("delete")
assert result["status_code"] == 204
assert mock_pinned.call_args[0][0] == "DELETE"
@patch("application.agents.tools.api_tool.validate_url")
@patch("application.agents.tools.api_tool.requests.patch")
def test_patch_method(self, mock_patch, mock_validate):
@patch("application.agents.tools.api_tool.pinned_request")
def test_patch_method(self, mock_pinned):
tool = APITool(config={"url": "https://example.com/item/1", "method": "PATCH"})
mock_resp = MagicMock()
mock_resp.status_code = 200
mock_resp.headers = {"Content-Type": "application/json"}
mock_resp.json.return_value = {"patched": True}
mock_resp.content = b'{"patched":true}'
mock_patch.return_value = mock_resp
mock_pinned.return_value = mock_resp
result = tool.execute_action("patch", field="val")
assert result["status_code"] == 200
assert mock_pinned.call_args[0][0] == "PATCH"
@patch("application.agents.tools.api_tool.validate_url")
@patch("application.agents.tools.api_tool.requests.head")
def test_head_method(self, mock_head, mock_validate):
@patch("application.agents.tools.api_tool.pinned_request")
def test_head_method(self, mock_pinned):
tool = APITool(config={"url": "https://example.com", "method": "HEAD"})
mock_resp = MagicMock()
mock_resp.status_code = 200
mock_resp.headers = {"Content-Type": "text/html"}
mock_resp.content = b''
mock_head.return_value = mock_resp
mock_pinned.return_value = mock_resp
result = tool.execute_action("check")
assert result["status_code"] == 200
assert mock_pinned.call_args[0][0] == "HEAD"
@patch("application.agents.tools.api_tool.validate_url")
@patch("application.agents.tools.api_tool.requests.options")
def test_options_method(self, mock_options, mock_validate):
@patch("application.agents.tools.api_tool.pinned_request")
def test_options_method(self, mock_pinned):
tool = APITool(config={"url": "https://example.com", "method": "OPTIONS"})
mock_resp = MagicMock()
mock_resp.status_code = 200
mock_resp.headers = {"Content-Type": "text/plain"}
mock_resp.content = b''
mock_options.return_value = mock_resp
mock_pinned.return_value = mock_resp
result = tool.execute_action("options")
assert result["status_code"] == 200
assert mock_pinned.call_args[0][0] == "OPTIONS"
@patch("application.agents.tools.api_tool.validate_url")
def test_unsupported_method(self, mock_validate):
def test_unsupported_method(self):
tool = APITool(config={"url": "https://example.com", "method": "CUSTOM"})
result = tool.execute_action("any")
assert result["status_code"] is None
@@ -193,19 +192,18 @@ class TestMakeApiCall:
@pytest.mark.unit
class TestSSRFValidation:
@patch("application.agents.tools.api_tool.validate_url")
def test_ssrf_blocked_initial_url(self, mock_validate, get_tool):
from application.core.url_validation import SSRFError
@patch("application.agents.tools.api_tool.pinned_request")
def test_ssrf_blocked(self, mock_pinned, get_tool):
from application.security.safe_url import UnsafeUserUrlError
mock_validate.side_effect = SSRFError("blocked")
mock_pinned.side_effect = UnsafeUserUrlError("blocked")
result = get_tool.execute_action("any")
assert result["status_code"] is None
assert "URL validation error" in result["message"]
@patch("application.agents.tools.api_tool.validate_url")
@patch("application.agents.tools.api_tool.requests.get")
def test_ssrf_blocked_after_param_substitution(self, mock_get, mock_validate):
from application.core.url_validation import SSRFError
@patch("application.agents.tools.api_tool.pinned_request")
def test_ssrf_blocked_with_path_params(self, mock_pinned):
from application.security.safe_url import UnsafeUserUrlError
tool = APITool(config={
"url": "https://api.example.com/{host}/data",
@@ -213,14 +211,7 @@ class TestSSRFValidation:
"query_params": {"host": "169.254.169.254"},
})
call_count = [0]
def side_effect(url):
call_count[0] += 1
if call_count[0] == 2:
raise SSRFError("blocked after substitution")
mock_validate.side_effect = side_effect
mock_pinned.side_effect = UnsafeUserUrlError("blocked")
result = tool.execute_action("any")
assert result["status_code"] is None
assert "URL validation error" in result["message"]
@@ -234,40 +225,36 @@ class TestSSRFValidation:
@pytest.mark.unit
class TestErrorHandling:
@patch("application.agents.tools.api_tool.validate_url")
@patch("application.agents.tools.api_tool.requests.get")
def test_timeout_error(self, mock_get, mock_validate, get_tool):
mock_get.side_effect = requests.exceptions.Timeout()
@patch("application.agents.tools.api_tool.pinned_request")
def test_timeout_error(self, mock_pinned, get_tool):
mock_pinned.side_effect = requests.exceptions.Timeout()
result = get_tool.execute_action("any")
assert result["status_code"] is None
assert "timeout" in result["message"].lower()
@patch("application.agents.tools.api_tool.validate_url")
@patch("application.agents.tools.api_tool.requests.get")
def test_connection_error(self, mock_get, mock_validate, get_tool):
mock_get.side_effect = requests.exceptions.ConnectionError("refused")
@patch("application.agents.tools.api_tool.pinned_request")
def test_connection_error(self, mock_pinned, get_tool):
mock_pinned.side_effect = requests.exceptions.ConnectionError("refused")
result = get_tool.execute_action("any")
assert result["status_code"] is None
assert "Connection error" in result["message"]
@patch("application.agents.tools.api_tool.validate_url")
@patch("application.agents.tools.api_tool.requests.get")
def test_http_error_with_json(self, mock_get, mock_validate, get_tool):
@patch("application.agents.tools.api_tool.pinned_request")
def test_http_error_with_json(self, mock_pinned, get_tool):
mock_resp = MagicMock()
mock_resp.status_code = 422
mock_resp.json.return_value = {"error": "invalid_field"}
mock_resp.raise_for_status.side_effect = requests.exceptions.HTTPError(
response=mock_resp
)
mock_get.return_value = mock_resp
mock_pinned.return_value = mock_resp
result = get_tool.execute_action("any")
assert result["status_code"] == 422
assert result["data"] == {"error": "invalid_field"}
@patch("application.agents.tools.api_tool.validate_url")
@patch("application.agents.tools.api_tool.requests.get")
def test_http_error_non_json_body(self, mock_get, mock_validate, get_tool):
@patch("application.agents.tools.api_tool.pinned_request")
def test_http_error_non_json_body(self, mock_pinned, get_tool):
mock_resp = MagicMock()
mock_resp.status_code = 404
mock_resp.text = "Not Found"
@@ -275,29 +262,26 @@ class TestErrorHandling:
mock_resp.raise_for_status.side_effect = requests.exceptions.HTTPError(
response=mock_resp
)
mock_get.return_value = mock_resp
mock_pinned.return_value = mock_resp
result = get_tool.execute_action("any")
assert result["status_code"] == 404
assert result["data"] == "Not Found"
@patch("application.agents.tools.api_tool.validate_url")
@patch("application.agents.tools.api_tool.requests.get")
def test_request_exception(self, mock_get, mock_validate, get_tool):
mock_get.side_effect = requests.exceptions.RequestException("something")
@patch("application.agents.tools.api_tool.pinned_request")
def test_request_exception(self, mock_pinned, get_tool):
mock_pinned.side_effect = requests.exceptions.RequestException("something")
result = get_tool.execute_action("any")
assert "API call failed" in result["message"]
@patch("application.agents.tools.api_tool.validate_url")
@patch("application.agents.tools.api_tool.requests.get")
def test_unexpected_exception(self, mock_get, mock_validate, get_tool):
mock_get.side_effect = RuntimeError("unexpected")
@patch("application.agents.tools.api_tool.pinned_request")
def test_unexpected_exception(self, mock_pinned, get_tool):
mock_pinned.side_effect = RuntimeError("unexpected")
result = get_tool.execute_action("any")
assert "Unexpected error" in result["message"]
@patch("application.agents.tools.api_tool.validate_url")
@patch("application.agents.tools.api_tool.requests.post")
def test_body_serialization_error(self, mock_post, mock_validate):
@patch("application.agents.tools.api_tool.pinned_request")
def test_body_serialization_error(self, mock_pinned):
tool = APITool(config={
"url": "https://example.com",
"method": "POST",
@@ -320,9 +304,8 @@ class TestErrorHandling:
@pytest.mark.unit
class TestPathParamSubstitution:
@patch("application.agents.tools.api_tool.validate_url")
@patch("application.agents.tools.api_tool.requests.get")
def test_path_params_substituted(self, mock_get, mock_validate):
@patch("application.agents.tools.api_tool.pinned_request")
def test_path_params_substituted(self, mock_pinned):
tool = APITool(config={
"url": "https://api.example.com/users/{user_id}/posts/{post_id}",
"method": "GET",
@@ -333,17 +316,16 @@ class TestPathParamSubstitution:
mock_resp.headers = {"Content-Type": "application/json"}
mock_resp.json.return_value = []
mock_resp.content = b'[]'
mock_get.return_value = mock_resp
mock_pinned.return_value = mock_resp
tool.execute_action("get")
called_url = mock_get.call_args[0][0]
called_url = mock_pinned.call_args[0][1]
assert "/users/42/posts/7" in called_url
assert "{user_id}" not in called_url
@patch("application.agents.tools.api_tool.validate_url")
@patch("application.agents.tools.api_tool.requests.get")
def test_remaining_query_params_appended(self, mock_get, mock_validate):
@patch("application.agents.tools.api_tool.pinned_request")
def test_remaining_query_params_appended(self, mock_pinned):
tool = APITool(config={
"url": "https://api.example.com/items",
"method": "GET",
@@ -354,19 +336,16 @@ class TestPathParamSubstitution:
mock_resp.headers = {"Content-Type": "application/json"}
mock_resp.json.return_value = []
mock_resp.content = b'[]'
mock_get.return_value = mock_resp
mock_pinned.return_value = mock_resp
tool.execute_action("get")
called_url = mock_get.call_args[0][0]
called_url = mock_pinned.call_args[0][1]
assert "page=2" in called_url
assert "limit=10" in called_url
@patch("application.agents.tools.api_tool.validate_url")
@patch("application.agents.tools.api_tool.requests.get")
def test_query_params_append_with_existing_query_string(
self, mock_get, mock_validate
):
@patch("application.agents.tools.api_tool.pinned_request")
def test_query_params_append_with_existing_query_string(self, mock_pinned):
tool = APITool(config={
"url": "https://api.example.com/items?existing=true",
"method": "GET",
@@ -377,27 +356,65 @@ class TestPathParamSubstitution:
mock_resp.headers = {"Content-Type": "application/json"}
mock_resp.json.return_value = []
mock_resp.content = b'[]'
mock_get.return_value = mock_resp
mock_pinned.return_value = mock_resp
tool.execute_action("get")
called_url = mock_get.call_args[0][0]
called_url = mock_pinned.call_args[0][1]
assert "&page=1" in called_url
@patch("application.agents.tools.api_tool.validate_url")
@patch("application.agents.tools.api_tool.requests.post")
def test_empty_body_no_serialization(self, mock_post, mock_validate):
@patch("application.agents.tools.api_tool.pinned_request")
def test_empty_body_no_serialization(self, mock_pinned):
tool = APITool(config={"url": "https://example.com", "method": "POST"})
mock_resp = MagicMock()
mock_resp.status_code = 200
mock_resp.headers = {"Content-Type": "application/json"}
mock_resp.json.return_value = {}
mock_resp.content = b'{}'
mock_post.return_value = mock_resp
mock_pinned.return_value = mock_resp
result = tool.execute_action("create")
assert result["status_code"] == 200
@patch("application.agents.tools.api_tool.pinned_request")
def test_path_params_are_url_encoded(self, mock_pinned):
tool = APITool(config={
"url": "https://api.example.com/users/{user_id}/profile",
"method": "GET",
"query_params": {"user_id": "../../admin"},
})
mock_resp = MagicMock()
mock_resp.status_code = 200
mock_resp.headers = {"Content-Type": "application/json"}
mock_resp.json.return_value = {}
mock_resp.content = b'{}'
mock_pinned.return_value = mock_resp
tool.execute_action("get")
called_url = mock_pinned.call_args[0][1]
assert "../../admin" not in called_url
assert "%2F" in called_url or "%2f" in called_url
@patch("application.agents.tools.api_tool.pinned_request")
def test_path_params_query_injection_encoded(self, mock_pinned):
tool = APITool(config={
"url": "https://api.example.com/items/{item_id}",
"method": "GET",
"query_params": {"item_id": "x?admin=true"},
})
mock_resp = MagicMock()
mock_resp.status_code = 200
mock_resp.headers = {"Content-Type": "application/json"}
mock_resp.json.return_value = {}
mock_resp.content = b'{}'
mock_pinned.return_value = mock_resp
tool.execute_action("get")
called_url = mock_pinned.call_args[0][1]
assert "x?admin=true" not in called_url
# =====================================================================
# Parse Response
@@ -494,11 +511,8 @@ class TestAPIToolMetadata:
def test_config_requirements_empty(self, get_tool):
assert get_tool.get_config_requirements() == {}
@patch("application.agents.tools.api_tool.validate_url")
@patch("application.agents.tools.api_tool.requests.post")
def test_content_type_set_for_post_with_no_headers(
self, mock_post, mock_validate
):
@patch("application.agents.tools.api_tool.pinned_request")
def test_content_type_set_for_post_with_no_headers(self, mock_pinned):
tool = APITool(config={
"url": "https://example.com",
"method": "POST",
@@ -509,8 +523,8 @@ class TestAPIToolMetadata:
mock_resp.headers = {"Content-Type": "application/json"}
mock_resp.json.return_value = {}
mock_resp.content = b'{}'
mock_post.return_value = mock_resp
mock_pinned.return_value = mock_resp
tool.execute_action("create")
call_headers = mock_post.call_args[1]["headers"]
call_headers = mock_pinned.call_args.kwargs["headers"]
assert "Content-Type" in call_headers
tests/agents/tools/test_read_webpage_errors.py
+6 -12
View File
@@ -11,10 +11,7 @@ class TestReadWebpageErrors:
tool = ReadWebpageTool(config={})
with patch(
"application.agents.tools.read_webpage.validate_url",
return_value=None,
), patch(
"application.agents.tools.read_webpage.requests.get",
"application.agents.tools.read_webpage.pinned_request",
side_effect=requests.exceptions.RequestException("bad url"),
):
got = tool.execute_action(
@@ -27,20 +24,17 @@ class TestReadWebpageErrors:
tool = ReadWebpageTool(config={})
with patch(
"application.agents.tools.read_webpage.validate_url",
return_value=None,
), patch(
"application.agents.tools.read_webpage.markdownify",
side_effect=RuntimeError("boom"),
), patch(
"application.agents.tools.read_webpage.requests.get",
) as mock_get:
mock_get.return_value.text = "<h1>hi</h1>"
mock_get.return_value.raise_for_status.return_value = None
"application.agents.tools.read_webpage.pinned_request",
) as mock_pinned_request:
mock_pinned_request.return_value.text = "<h1>hi</h1>"
mock_pinned_request.return_value.raise_for_status.return_value = None
got = tool.execute_action(
"read_webpage", url="https://example.com/",
)
assert "Error processing URL" in got
assert "Error fetching URL" in got
class TestBaseAgentMinorBranches:
tests/parser/remote/test_crawler_loader.py
+28 -97
View File
@@ -4,7 +4,6 @@ import pytest
from application.parser.remote.crawler_loader import CrawlerLoader
from application.parser.schema.base import Document
from langchain_core.documents import Document as LCDocument
class DummyResponse:
@@ -24,8 +23,8 @@ def _mock_validate_url(url):
@patch("application.parser.remote.crawler_loader.validate_url", side_effect=_mock_validate_url)
@patch("application.parser.remote.crawler_loader.requests.get")
def test_load_data_crawls_same_domain_links(mock_requests_get, mock_validate_url):
@patch("application.parser.remote.crawler_loader.pinned_request")
def test_load_data_crawls_same_domain_links(mock_pinned_request, mock_validate_url):
responses = {
"http://example.com": DummyResponse(
"""
@@ -40,37 +39,14 @@ def test_load_data_crawls_same_domain_links(mock_requests_get, mock_validate_url
"http://example.com/about": DummyResponse("<html><body>About page</body></html>"),
}
def response_side_effect(url: str, timeout=30):
def response_side_effect(_method: str, url: str, timeout=30):
if url not in responses:
raise AssertionError(f"Unexpected request for URL: {url}")
return responses[url]
mock_requests_get.side_effect = response_side_effect
root_doc = MagicMock(spec=LCDocument)
root_doc.page_content = "Root content"
root_doc.metadata = {"source": "http://example.com"}
about_doc = MagicMock(spec=LCDocument)
about_doc.page_content = "About content"
about_doc.metadata = {"source": "http://example.com/about"}
loader_instances = {
"http://example.com": MagicMock(),
"http://example.com/about": MagicMock(),
}
loader_instances["http://example.com"].load.return_value = [root_doc]
loader_instances["http://example.com/about"].load.return_value = [about_doc]
loader_call_order = []
def loader_factory(url_list):
url = url_list[0]
loader_call_order.append(url)
return loader_instances[url]
mock_pinned_request.side_effect = response_side_effect
crawler = CrawlerLoader(limit=5)
crawler.loader = MagicMock(side_effect=loader_factory)
result = crawler.load_data("http://example.com")
@@ -84,34 +60,23 @@ def test_load_data_crawls_same_domain_links(mock_requests_get, mock_validate_url
assert paths == {"index.md", "about.md"}
texts = {doc.text for doc in result}
assert texts == {"Root content", "About content"}
assert texts == {"About\nExternal", "About page"}
assert mock_requests_get.call_count == 2
assert loader_call_order == ["http://example.com", "http://example.com/about"]
assert mock_pinned_request.call_count == 2
@patch("application.parser.remote.crawler_loader.validate_url", side_effect=_mock_validate_url)
@patch("application.parser.remote.crawler_loader.requests.get")
def test_load_data_accepts_list_input_and_adds_scheme(mock_requests_get, mock_validate_url):
mock_requests_get.return_value = DummyResponse("<html><body>No links here</body></html>")
doc = MagicMock(spec=LCDocument)
doc.page_content = "Homepage"
doc.metadata = {"source": "http://example.com"}
loader_instance = MagicMock()
loader_instance.load.return_value = [doc]
@patch("application.parser.remote.crawler_loader.pinned_request")
def test_load_data_accepts_list_input_and_adds_scheme(mock_pinned_request, mock_validate_url):
mock_pinned_request.return_value = DummyResponse("<html><body>No links here</body></html>")
crawler = CrawlerLoader()
crawler.loader = MagicMock(return_value=loader_instance)
result = crawler.load_data(["example.com", "unused.com"])
mock_requests_get.assert_called_once_with("http://example.com", timeout=30)
crawler.loader.assert_called_once_with(["http://example.com"])
mock_pinned_request.assert_called_once_with("GET", "http://example.com", timeout=30)
assert len(result) == 1
assert result[0].text == "Homepage"
assert result[0].text == "No links here"
assert result[0].extra_info == {
"source": "http://example.com",
"file_path": "index.md",
@@ -119,8 +84,8 @@ def test_load_data_accepts_list_input_and_adds_scheme(mock_requests_get, mock_va
@patch("application.parser.remote.crawler_loader.validate_url", side_effect=_mock_validate_url)
@patch("application.parser.remote.crawler_loader.requests.get")
def test_load_data_respects_limit(mock_requests_get, mock_validate_url):
@patch("application.parser.remote.crawler_loader.pinned_request")
def test_load_data_respects_limit(mock_pinned_request, mock_validate_url):
responses = {
"http://example.com": DummyResponse(
"""
@@ -134,51 +99,28 @@ def test_load_data_respects_limit(mock_requests_get, mock_validate_url):
"http://example.com/about": DummyResponse("<html><body>About</body></html>"),
}
mock_requests_get.side_effect = lambda url, timeout=30: responses[url]
root_doc = MagicMock(spec=LCDocument)
root_doc.page_content = "Root content"
root_doc.metadata = {"source": "http://example.com"}
about_doc = MagicMock(spec=LCDocument)
about_doc.page_content = "About content"
about_doc.metadata = {"source": "http://example.com/about"}
loader_instances = {
"http://example.com": MagicMock(),
"http://example.com/about": MagicMock(),
}
loader_instances["http://example.com"].load.return_value = [root_doc]
loader_instances["http://example.com/about"].load.return_value = [about_doc]
mock_pinned_request.side_effect = lambda _method, url, timeout=30: responses[url]
crawler = CrawlerLoader(limit=1)
crawler.loader = MagicMock(side_effect=lambda url_list: loader_instances[url_list[0]])
result = crawler.load_data("http://example.com")
assert len(result) == 1
assert result[0].text == "Root content"
assert mock_requests_get.call_count == 1
assert crawler.loader.call_count == 1
assert result[0].text == "About"
assert mock_pinned_request.call_count == 1
@patch("application.parser.remote.crawler_loader.validate_url", side_effect=_mock_validate_url)
@patch("application.parser.remote.crawler_loader.logging")
@patch("application.parser.remote.crawler_loader.requests.get")
def test_load_data_logs_and_skips_on_loader_error(mock_requests_get, mock_logging, mock_validate_url):
mock_requests_get.return_value = DummyResponse("<html><body>Error route</body></html>")
failing_loader_instance = MagicMock()
failing_loader_instance.load.side_effect = Exception("load failure")
@patch("application.parser.remote.crawler_loader.pinned_request")
def test_load_data_logs_and_skips_on_request_error(mock_pinned_request, mock_logging, mock_validate_url):
mock_pinned_request.side_effect = Exception("load failure")
crawler = CrawlerLoader()
crawler.loader = MagicMock(return_value=failing_loader_instance)
result = crawler.load_data("http://example.com")
assert result == []
mock_requests_get.assert_called_once_with("http://example.com", timeout=30)
failing_loader_instance.load.assert_called_once()
mock_pinned_request.assert_called_once_with("GET", "http://example.com", timeout=30)
mock_logging.error.assert_called_once()
message, = mock_logging.error.call_args.args
@@ -221,34 +163,23 @@ def test_url_to_virtual_path_variants():
@pytest.mark.unit
class TestCrawlerLoaderGaps:
def test_ssrf_validation_skips_invalid_url(self):
"""Cover lines 41-43: SSRF validation failure skips URL."""
def test_pinned_fetch_builds_document_without_webbase_loader(self):
"""The crawler should index the response body it already fetched."""
from application.parser.remote.crawler_loader import CrawlerLoader
from application.core.url_validation import SSRFError
loader = CrawlerLoader(limit=5)
with patch(
"application.parser.remote.crawler_loader.validate_url",
side_effect=[
"https://example.com",
SSRFError("blocked"),
],
return_value="https://example.com",
):
with patch(
"application.parser.remote.crawler_loader.requests.get"
) as mock_get:
# First URL succeeds validation but response has no links
"application.parser.remote.crawler_loader.pinned_request"
) as mock_pinned_request:
mock_response = MagicMock()
mock_response.status_code = 200
mock_response.text = "<html><body>test</body></html>"
mock_response.raise_for_status.return_value = None
mock_get.return_value = mock_response
mock_pinned_request.return_value = mock_response
with patch.object(loader, "loader") as mock_loader_cls:
mock_doc = MagicMock()
mock_doc.page_content = "test content"
mock_doc.metadata = {}
mock_loader_cls.return_value.load.return_value = [mock_doc]
result = loader.load_data("https://example.com")
assert isinstance(result, list)
result = loader.load_data("https://example.com")
assert result[0].text == "test"
tests/parser/remote/test_crawler_markdown.py
+18 -10
View File
@@ -67,13 +67,18 @@ def _patch_markdownify(monkeypatch):
return outputs
def _setup_session(mock_get_side_effect):
session = MagicMock()
session.get.side_effect = mock_get_side_effect
return session
def _patch_pinned_request(monkeypatch, side_effect):
"""Replace pinned_request with a stub that maps URL -> response or raises."""
def fake_pinned_request(method, url, **_kwargs):
return side_effect(url)
monkeypatch.setattr(
"application.parser.remote.crawler_markdown.pinned_request",
fake_pinned_request,
)
def test_load_data_filters_external_links(_patch_markdownify):
def test_load_data_filters_external_links(monkeypatch, _patch_markdownify):
root_html = """
<html><head><title>Home</title></head>
<body><a href="/about">About</a><a href="https://other.com">Other</a><p>Welcome</p></body>
@@ -89,8 +94,9 @@ def test_load_data_filters_external_links(_patch_markdownify):
"http://example.com/about": DummyResponse(about_html),
}
_patch_pinned_request(monkeypatch, lambda url: responses[url])
loader = CrawlerLoader(limit=5)
loader.session = _setup_session(lambda url, timeout=10: responses[url])
docs = loader.load_data("http://example.com")
@@ -102,7 +108,7 @@ def test_load_data_filters_external_links(_patch_markdownify):
assert texts == {"Home Markdown", "About Markdown"}
def test_load_data_allows_subdomains(_patch_markdownify):
def test_load_data_allows_subdomains(monkeypatch, _patch_markdownify):
root_html = """
<html><head><title>Home</title></head>
<body><a href="http://blog.example.com/post">Blog</a></body>
@@ -118,8 +124,9 @@ def test_load_data_allows_subdomains(_patch_markdownify):
"http://blog.example.com/post": DummyResponse(blog_html),
}
_patch_pinned_request(monkeypatch, lambda url: responses[url])
loader = CrawlerLoader(limit=5, allow_subdomains=True)
loader.session = _setup_session(lambda url, timeout=10: responses[url])
docs = loader.load_data("http://example.com")
@@ -137,13 +144,14 @@ def test_load_data_handles_fetch_errors(monkeypatch, _patch_markdownify, _patch_
_patch_markdownify[root_html] = "Home Markdown"
def side_effect(url, timeout=10):
def side_effect(url):
if url == "http://example.com":
return DummyResponse(root_html)
raise requests.exceptions.RequestException("boom")
_patch_pinned_request(monkeypatch, side_effect)
loader = CrawlerLoader(limit=5)
loader.session = _setup_session(side_effect)
mock_print = MagicMock()
monkeypatch.setattr("builtins.print", mock_print)
tests/parser/remote/test_sitemap_loader.py
+69 -43
View File
@@ -10,6 +10,7 @@ import pytest
import requests
from application.parser.remote.sitemap_loader import SitemapLoader
from application.parser.schema.base import Document
# =====================================================================
@@ -30,7 +31,7 @@ class TestSitemapLoaderInit:
def test_has_loader_class(self):
loader = SitemapLoader()
assert loader.loader is not None
assert not hasattr(loader, "loader")
# =====================================================================
@@ -136,6 +137,19 @@ class TestParseSitemap:
urls = loader._parse_sitemap(sitemap_xml)
assert urls == []
def test_parse_skips_empty_loc_entries(self):
"""Empty or self-closing <loc> tags must not yield None URLs."""
loader = SitemapLoader()
sitemap_xml = b"""<?xml version="1.0" encoding="UTF-8"?>
<urlset xmlns="http://www.sitemaps.org/schemas/sitemap/0.9">
<url><loc></loc></url>
<url><loc>https://example.com/p</loc></url>
<url><loc/></url>
</urlset>"""
urls = loader._parse_sitemap(sitemap_xml)
assert urls == ["https://example.com/p"]
# =====================================================================
# _extract_urls
@@ -145,9 +159,8 @@ class TestParseSitemap:
@pytest.mark.unit
class TestExtractUrls:
@patch("application.parser.remote.sitemap_loader.validate_url")
@patch("application.parser.remote.sitemap_loader.requests.get")
def test_extract_urls_from_sitemap(self, mock_get, mock_validate):
@patch("application.parser.remote.sitemap_loader.pinned_request")
def test_extract_urls_from_sitemap(self, mock_pinned_request):
loader = SitemapLoader()
response = MagicMock()
@@ -158,51 +171,50 @@ class TestExtractUrls:
<urlset xmlns="http://www.sitemaps.org/schemas/sitemap/0.9">
<url><loc>https://example.com/p</loc></url>
</urlset>"""
mock_get.return_value = response
response.raise_for_status.return_value = None
mock_pinned_request.return_value = response
urls = loader._extract_urls("https://example.com/sitemap.xml")
assert "https://example.com/p" in urls
@patch("application.parser.remote.sitemap_loader.validate_url")
@patch("application.parser.remote.sitemap_loader.requests.get")
def test_extract_urls_not_sitemap(self, mock_get, mock_validate):
@patch("application.parser.remote.sitemap_loader.pinned_request")
def test_extract_urls_not_sitemap(self, mock_pinned_request):
loader = SitemapLoader()
response = MagicMock()
response.headers = {"Content-Type": "text/html"}
response.url = "https://example.com/page"
response.text = "<html>Normal page</html>"
mock_get.return_value = response
response.raise_for_status.return_value = None
mock_pinned_request.return_value = response
urls = loader._extract_urls("https://example.com/page")
assert urls == ["https://example.com/page"]
@patch("application.parser.remote.sitemap_loader.validate_url")
@patch("application.parser.remote.sitemap_loader.requests.get")
def test_extract_urls_http_error(self, mock_get, mock_validate):
@patch("application.parser.remote.sitemap_loader.pinned_request")
def test_extract_urls_http_error(self, mock_pinned_request):
loader = SitemapLoader()
mock_get.side_effect = requests.exceptions.HTTPError("404")
mock_pinned_request.side_effect = requests.exceptions.HTTPError("404")
urls = loader._extract_urls("https://example.com/missing")
assert urls == []
@patch("application.parser.remote.sitemap_loader.validate_url")
@patch("application.parser.remote.sitemap_loader.requests.get")
def test_extract_urls_connection_error(self, mock_get, mock_validate):
@patch("application.parser.remote.sitemap_loader.pinned_request")
def test_extract_urls_connection_error(self, mock_pinned_request):
loader = SitemapLoader()
mock_get.side_effect = requests.exceptions.ConnectionError()
mock_pinned_request.side_effect = requests.exceptions.ConnectionError()
urls = loader._extract_urls("https://example.com/bad")
assert urls == []
def test_extract_urls_ssrf_blocked(self):
from application.core.url_validation import SSRFError
from application.security.safe_url import UnsafeUserUrlError
loader = SitemapLoader()
with patch(
"application.parser.remote.sitemap_loader.validate_url",
side_effect=SSRFError("blocked"),
"application.parser.remote.sitemap_loader.pinned_request",
side_effect=UnsafeUserUrlError("blocked"),
):
urls = loader._extract_urls("http://169.254.169.254/")
assert urls == []
@@ -217,13 +229,14 @@ class TestExtractUrls:
class TestSitemapLoaderLoadData:
@patch("application.parser.remote.sitemap_loader.validate_url")
def test_load_data_success(self, mock_validate):
@patch("application.parser.remote.sitemap_loader.pinned_request")
def test_load_data_success(self, mock_pinned_request, mock_validate):
loader = SitemapLoader(limit=10)
mock_doc = MagicMock()
mock_loader_instance = MagicMock()
mock_loader_instance.load.return_value = [mock_doc]
loader.loader = MagicMock(return_value=mock_loader_instance)
mock_validate.side_effect = lambda url: url
response = MagicMock()
response.text = "Page body"
response.raise_for_status.return_value = None
mock_pinned_request.return_value = response
with patch.object(
loader, "_extract_urls",
@@ -231,6 +244,9 @@ class TestSitemapLoaderLoadData:
):
docs = loader.load_data("https://example.com/sitemap.xml")
assert len(docs) == 1
assert isinstance(docs[0], Document)
assert docs[0].text == "Page body"
assert docs[0].extra_info == {"source": "https://example.com/page1"}
@patch("application.parser.remote.sitemap_loader.validate_url")
def test_load_data_no_urls(self, mock_validate):
@@ -241,12 +257,14 @@ class TestSitemapLoaderLoadData:
assert docs == []
@patch("application.parser.remote.sitemap_loader.validate_url")
def test_load_data_list_input(self, mock_validate):
@patch("application.parser.remote.sitemap_loader.pinned_request")
def test_load_data_list_input(self, mock_pinned_request, mock_validate):
loader = SitemapLoader()
mock_loader_instance = MagicMock()
mock_loader_instance.load.return_value = [MagicMock()]
loader.loader = MagicMock(return_value=mock_loader_instance)
mock_validate.side_effect = lambda url: url
response = MagicMock()
response.text = "List body"
response.raise_for_status.return_value = None
mock_pinned_request.return_value = response
with patch.object(
loader, "_extract_urls",
@@ -254,24 +272,30 @@ class TestSitemapLoaderLoadData:
):
docs = loader.load_data(["https://example.com/sitemap.xml"])
assert len(docs) == 1
assert docs[0].text == "List body"
@patch("application.parser.remote.sitemap_loader.validate_url")
def test_load_data_respects_limit(self, mock_validate):
@patch("application.parser.remote.sitemap_loader.pinned_request")
def test_load_data_respects_limit(self, mock_pinned_request, mock_validate):
loader = SitemapLoader(limit=2)
mock_loader_instance = MagicMock()
mock_loader_instance.load.return_value = [MagicMock()]
loader.loader = MagicMock(return_value=mock_loader_instance)
mock_validate.side_effect = lambda url: url
response = MagicMock()
response.text = "Limited body"
response.raise_for_status.return_value = None
mock_pinned_request.return_value = response
urls = [f"https://example.com/page{i}" for i in range(10)]
with patch.object(loader, "_extract_urls", return_value=urls):
docs = loader.load_data("https://example.com/sitemap.xml")
assert len(docs) == 2
assert mock_pinned_request.call_count == 2
@patch("application.parser.remote.sitemap_loader.validate_url")
def test_load_data_handles_url_error(self, mock_validate):
@patch("application.parser.remote.sitemap_loader.pinned_request")
def test_load_data_handles_url_error(self, mock_pinned_request, mock_validate):
loader = SitemapLoader()
loader.loader = MagicMock(side_effect=Exception("Load failed"))
mock_validate.side_effect = lambda url: url
mock_pinned_request.side_effect = Exception("Load failed")
with patch.object(
loader, "_extract_urls",
@@ -293,12 +317,14 @@ class TestSitemapLoaderLoadData:
assert docs == []
@patch("application.parser.remote.sitemap_loader.validate_url")
def test_load_data_no_limit(self, mock_validate):
@patch("application.parser.remote.sitemap_loader.pinned_request")
def test_load_data_no_limit(self, mock_pinned_request, mock_validate):
loader = SitemapLoader(limit=None)
mock_loader_instance = MagicMock()
mock_loader_instance.load.return_value = [MagicMock()]
loader.loader = MagicMock(return_value=mock_loader_instance)
mock_validate.side_effect = lambda url: url
response = MagicMock()
response.text = "No limit body"
response.raise_for_status.return_value = None
mock_pinned_request.return_value = response
urls = [f"https://example.com/page{i}" for i in range(5)]
with patch.object(loader, "_extract_urls", return_value=urls):
tests/parser/remote/test_web_loader.py
+125 -188
View File
@@ -15,44 +15,31 @@ def _mock_validate_url(url):
return url
def _fake_response(html: str) -> MagicMock:
"""Build a MagicMock that quacks like a requests.Response for our purposes."""
response = MagicMock()
response.text = html
response.raise_for_status.return_value = None
return response
@pytest.fixture
def web_loader():
return WebLoader()
@pytest.fixture
def mock_langchain_document():
"""Create a mock LangChain document."""
doc = MagicMock(spec=LCDocument)
doc.page_content = "Test web page content"
doc.metadata = {"source": "https://example.com", "title": "Test Page"}
return doc
@pytest.fixture
def mock_web_base_loader():
"""Create a mock WebBaseLoader class."""
mock_loader_class = MagicMock()
mock_loader_instance = MagicMock()
mock_loader_class.return_value = mock_loader_instance
return mock_loader_class, mock_loader_instance
class TestWebLoaderInitialization:
"""Test WebLoader initialization."""
def test_init(self, web_loader):
"""Test WebLoader initialization."""
assert web_loader.loader is not None
from langchain_community.document_loaders import WebBaseLoader
assert web_loader.loader == WebBaseLoader
def test_init_has_no_loader_attribute(self, web_loader):
"""Post-pinned-request migration, WebLoader no longer keeps a langchain loader class."""
assert not hasattr(web_loader, "loader")
class TestWebLoaderHeaders:
"""Test WebLoader headers configuration."""
def test_headers_defined(self):
"""Test that headers are properly defined."""
assert isinstance(headers, dict)
assert "User-Agent" in headers
assert "Accept" in headers
@@ -63,7 +50,6 @@ class TestWebLoaderHeaders:
assert "Upgrade-Insecure-Requests" in headers
def test_headers_values(self):
"""Test header values are reasonable."""
assert headers["User-Agent"] == "Mozilla/5.0"
assert "text/html" in headers["Accept"]
assert headers["Referer"] == "https://www.google.com/"
@@ -75,49 +61,34 @@ class TestWebLoaderLoadData:
"""Test WebLoader load_data method."""
@patch("application.parser.remote.web_loader.validate_url", side_effect=_mock_validate_url)
def test_load_data_single_url_string(self, mock_validate, web_loader, mock_langchain_document):
"""Test loading data from a single URL passed as string."""
mock_loader_instance = MagicMock()
mock_loader_instance.load.return_value = [mock_langchain_document]
mock_web_base_loader_class = MagicMock()
mock_web_base_loader_class.return_value = mock_loader_instance
web_loader.loader = mock_web_base_loader_class
@patch("application.parser.remote.web_loader.pinned_request")
def test_load_data_single_url_string(self, mock_pinned_request, mock_validate, web_loader):
mock_pinned_request.return_value = _fake_response(
"<html lang='en'><head><title>Test Page</title></head>"
"<body><p>Test web page content</p></body></html>"
)
result = web_loader.load_data("https://example.com")
assert len(result) == 1
assert isinstance(result[0], Document)
assert result[0].text == "Test web page content"
assert result[0].extra_info == {"source": "https://example.com", "title": "Test Page"}
mock_web_base_loader_class.assert_called_once_with(["https://example.com"], header_template=headers)
mock_loader_instance.load.assert_called_once()
assert result[0].text == "Test Page\nTest web page content"
assert result[0].extra_info == {
"source": "https://example.com",
"title": "Test Page",
"language": "en",
}
mock_pinned_request.assert_called_once_with(
"GET", "https://example.com", headers=headers, timeout=30
)
@patch("application.parser.remote.web_loader.validate_url", side_effect=_mock_validate_url)
def test_load_data_multiple_urls_list(self, mock_validate, web_loader):
"""Test loading data from multiple URLs passed as list."""
doc1 = MagicMock(spec=LCDocument)
doc1.page_content = "Content from site 1"
doc1.metadata = {"source": "https://site1.com"}
doc2 = MagicMock(spec=LCDocument)
doc2.page_content = "Content from site 2"
doc2.metadata = {"source": "https://site2.com"}
mock_loader_instance1 = MagicMock()
mock_loader_instance1.load.return_value = [doc1]
mock_loader_instance2 = MagicMock()
mock_loader_instance2.load.return_value = [doc2]
mock_web_base_loader_class = MagicMock()
mock_web_base_loader_class.side_effect = [mock_loader_instance1, mock_loader_instance2]
web_loader.loader = mock_web_base_loader_class
@patch("application.parser.remote.web_loader.pinned_request")
def test_load_data_multiple_urls_list(self, mock_pinned_request, mock_validate, web_loader):
mock_pinned_request.side_effect = [
_fake_response("<html><body>Content from site 1</body></html>"),
_fake_response("<html><body>Content from site 2</body></html>"),
]
urls = ["https://site1.com", "https://site2.com"]
result = web_loader.load_data(urls)
@@ -129,121 +100,108 @@ class TestWebLoaderLoadData:
assert result[0].extra_info == {"source": "https://site1.com"}
assert result[1].extra_info == {"source": "https://site2.com"}
assert mock_web_base_loader_class.call_count == 2
mock_web_base_loader_class.assert_any_call(["https://site1.com"], header_template=headers)
mock_web_base_loader_class.assert_any_call(["https://site2.com"], header_template=headers)
assert mock_pinned_request.call_count == 2
mock_pinned_request.assert_any_call(
"GET", "https://site1.com", headers=headers, timeout=30
)
mock_pinned_request.assert_any_call(
"GET", "https://site2.com", headers=headers, timeout=30
)
@patch("application.parser.remote.web_loader.validate_url", side_effect=_mock_validate_url)
def test_load_data_url_without_scheme(self, mock_validate, web_loader, mock_langchain_document):
"""Test loading data from URL without scheme (should add http://)."""
mock_loader_instance = MagicMock()
mock_loader_instance.load.return_value = [mock_langchain_document]
mock_web_base_loader_class = MagicMock()
mock_web_base_loader_class.return_value = mock_loader_instance
web_loader.loader = mock_web_base_loader_class
@patch("application.parser.remote.web_loader.pinned_request")
def test_load_data_url_without_scheme(self, mock_pinned_request, mock_validate, web_loader):
mock_pinned_request.return_value = _fake_response(
"<html><body>Schemeless</body></html>"
)
result = web_loader.load_data("example.com")
assert len(result) == 1
assert isinstance(result[0], Document)
# Verify WebBaseLoader was called with http:// prefix
mock_web_base_loader_class.assert_called_once_with(["http://example.com"], header_template=headers)
mock_pinned_request.assert_called_once_with(
"GET", "http://example.com", headers=headers, timeout=30
)
@patch("application.parser.remote.web_loader.validate_url", side_effect=_mock_validate_url)
def test_load_data_url_with_scheme(self, mock_validate, web_loader, mock_langchain_document):
"""Test loading data from URL with scheme (should not modify)."""
mock_loader_instance = MagicMock()
mock_loader_instance.load.return_value = [mock_langchain_document]
mock_web_base_loader_class = MagicMock()
mock_web_base_loader_class.return_value = mock_loader_instance
web_loader.loader = mock_web_base_loader_class
@patch("application.parser.remote.web_loader.pinned_request")
def test_load_data_url_with_scheme(self, mock_pinned_request, mock_validate, web_loader):
mock_pinned_request.return_value = _fake_response(
"<html><body>Schemed</body></html>"
)
result = web_loader.load_data("https://example.com")
assert len(result) == 1
# Verify WebBaseLoader was called with original URL
mock_web_base_loader_class.assert_called_once_with(["https://example.com"], header_template=headers)
mock_pinned_request.assert_called_once_with(
"GET", "https://example.com", headers=headers, timeout=30
)
@patch("application.parser.remote.web_loader.validate_url", side_effect=_mock_validate_url)
def test_load_data_multiple_documents_per_url(self, mock_validate, web_loader):
"""Test loading multiple documents from a single URL."""
doc1 = MagicMock(spec=LCDocument)
doc1.page_content = "First document content"
doc1.metadata = {"source": "https://example.com", "section": "intro"}
doc2 = MagicMock(spec=LCDocument)
doc2.page_content = "Second document content"
doc2.metadata = {"source": "https://example.com", "section": "main"}
mock_loader_instance = MagicMock()
mock_loader_instance.load.return_value = [doc1, doc2]
mock_web_base_loader_class = MagicMock()
mock_web_base_loader_class.return_value = mock_loader_instance
web_loader.loader = mock_web_base_loader_class
@patch("application.parser.remote.web_loader.pinned_request")
def test_load_data_skips_pages_without_title(self, mock_pinned_request, mock_validate, web_loader):
"""A page without <title> or <html lang=...> still loads, just with bare metadata."""
mock_pinned_request.return_value = _fake_response("<p>Bare body</p>")
result = web_loader.load_data("https://example.com")
assert len(result) == 2
assert result[0].text == "First document content"
assert result[1].text == "Second document content"
assert result[0].extra_info == {"source": "https://example.com", "section": "intro"}
assert result[1].extra_info == {"source": "https://example.com", "section": "main"}
assert len(result) == 1
assert result[0].text == "Bare body"
assert result[0].extra_info == {"source": "https://example.com"}
@patch("application.parser.remote.web_loader.validate_url", side_effect=_mock_validate_url)
@patch("application.parser.remote.web_loader.pinned_request")
def test_load_data_extracts_title_with_nested_markup(self, mock_pinned_request, mock_validate, web_loader):
"""<title> with nested-looking content must still produce a non-empty title.
BS4's `html.parser` differs across Python/bs4 versions: some treat
`<title>` as RAWTEXT (per HTML5 spec) and return the literal string
``"Hello <span>World</span>"``; others parse the inner tag and return
``"HelloWorld"``. The regression guarded here is that ``soup.title.string``
returns ``None`` in *either* case, dropping the title entirely — using
``get_text`` keeps it. Assert only that the title survives.
"""
mock_pinned_request.return_value = _fake_response(
"<html><head><title>Hello <span>World</span></title></head>"
"<body><p>x</p></body></html>"
)
result = web_loader.load_data("https://example.com")
title = result[0].extra_info.get("title")
assert title, "title was dropped — soup.title.string regression?"
assert "Hello" in title
class TestWebLoaderErrorHandling:
"""Test WebLoader error handling."""
@patch("application.parser.remote.web_loader.validate_url", side_effect=_mock_validate_url)
@patch('application.parser.remote.web_loader.logging')
def test_load_data_single_url_error(self, mock_logging, mock_validate, web_loader):
"""Test error handling for single URL that fails to load."""
mock_loader_instance = MagicMock()
mock_loader_instance.load.side_effect = Exception("Network error")
mock_web_base_loader_class = MagicMock()
mock_web_base_loader_class.return_value = mock_loader_instance
web_loader.loader = mock_web_base_loader_class
@patch("application.parser.remote.web_loader.pinned_request")
@patch("application.parser.remote.web_loader.logging")
def test_load_data_single_url_error(self, mock_logging, mock_pinned_request, mock_validate, web_loader):
mock_pinned_request.side_effect = Exception("Network error")
result = web_loader.load_data("https://invalid-url.com")
assert result == [] # Should return empty list on error
assert result == []
mock_logging.error.assert_called_once()
error_call = mock_logging.error.call_args
assert "Error processing URL https://invalid-url.com" in error_call[0][0]
assert error_call[1]["exc_info"] is True
@patch("application.parser.remote.web_loader.validate_url", side_effect=_mock_validate_url)
@patch('application.parser.remote.web_loader.logging')
def test_load_data_partial_failure(self, mock_logging, mock_validate, web_loader):
"""Test partial failure - some URLs succeed, some fail."""
doc1 = MagicMock(spec=LCDocument)
doc1.page_content = "Success content"
doc1.metadata = {"source": "https://good-url.com"}
mock_loader_instance1 = MagicMock()
mock_loader_instance1.load.return_value = [doc1]
mock_loader_instance2 = MagicMock()
mock_loader_instance2.load.side_effect = Exception("Network error")
mock_web_base_loader_class = MagicMock()
mock_web_base_loader_class.side_effect = [mock_loader_instance1, mock_loader_instance2]
web_loader.loader = mock_web_base_loader_class
@patch("application.parser.remote.web_loader.pinned_request")
@patch("application.parser.remote.web_loader.logging")
def test_load_data_partial_failure(self, mock_logging, mock_pinned_request, mock_validate, web_loader):
mock_pinned_request.side_effect = [
_fake_response("<html><body>Success content</body></html>"),
Exception("Network error"),
]
urls = ["https://good-url.com", "https://bad-url.com"]
result = web_loader.load_data(urls)
assert len(result) == 1 # Only successful URL should be in results
assert len(result) == 1
assert result[0].text == "Success content"
assert result[0].extra_info == {"source": "https://good-url.com"}
@@ -256,26 +214,22 @@ class TestWebLoaderSSRF:
"""Test WebLoader SSRF protection."""
@patch("application.parser.remote.web_loader.validate_url")
def test_load_data_skips_url_failing_ssrf_validation(self, mock_validate, web_loader):
@patch("application.parser.remote.web_loader.pinned_request")
def test_load_data_skips_url_failing_ssrf_validation(self, mock_pinned_request, mock_validate, web_loader):
"""A URL that fails SSRF validation must be skipped, never reaching the fetcher."""
mock_validate.side_effect = SSRFError("Access to private/internal IP addresses is not allowed.")
mock_web_base_loader_class = MagicMock()
web_loader.loader = mock_web_base_loader_class
result = web_loader.load_data("http://169.254.169.254/latest/meta-data/")
assert result == []
mock_validate.assert_called_once_with("http://169.254.169.254/latest/meta-data/")
mock_web_base_loader_class.assert_not_called()
mock_pinned_request.assert_not_called()
@patch("application.parser.remote.web_loader.validate_url")
@patch("application.parser.remote.web_loader.logging")
def test_load_data_logs_warning_on_ssrf_failure(self, mock_logging, mock_validate, web_loader):
"""SSRF rejections should be logged as warnings, not errors."""
mock_validate.side_effect = SSRFError("blocked")
web_loader.loader = MagicMock()
web_loader.load_data("http://127.0.0.1")
mock_logging.warning.assert_called_once()
@@ -284,65 +238,52 @@ class TestWebLoaderSSRF:
assert "127.0.0.1" in warning_msg
@patch("application.parser.remote.web_loader.validate_url")
def test_load_data_mixed_safe_and_unsafe_urls(self, mock_validate, web_loader):
"""In a list of URLs, unsafe ones are skipped while safe ones still load."""
@patch("application.parser.remote.web_loader.pinned_request")
def test_load_data_mixed_safe_and_unsafe_urls(self, mock_pinned_request, mock_validate, web_loader):
def validate(url):
if "169.254.169.254" in url:
raise SSRFError("metadata blocked")
return url
mock_validate.side_effect = validate
safe_doc = MagicMock(spec=LCDocument)
safe_doc.page_content = "Public page"
safe_doc.metadata = {"source": "https://example.com"}
safe_loader_instance = MagicMock()
safe_loader_instance.load.return_value = [safe_doc]
mock_web_base_loader_class = MagicMock(return_value=safe_loader_instance)
web_loader.loader = mock_web_base_loader_class
mock_pinned_request.return_value = _fake_response(
"<html><body>Public page</body></html>"
)
urls = ["https://example.com", "http://169.254.169.254/latest/meta-data/"]
result = web_loader.load_data(urls)
assert len(result) == 1
assert result[0].text == "Public page"
# Only the safe URL should have been fetched
mock_web_base_loader_class.assert_called_once_with(["https://example.com"], header_template=headers)
mock_pinned_request.assert_called_once_with(
"GET", "https://example.com", headers=headers, timeout=30
)
class TestWebLoaderEdgeCases:
"""Test WebLoader edge cases."""
def test_load_data_empty_list(self, web_loader):
"""Test loading data with empty URL list."""
result = web_loader.load_data([])
assert result == []
@patch("application.parser.remote.web_loader.validate_url", side_effect=_mock_validate_url)
def test_load_data_empty_response(self, mock_validate, web_loader):
"""Test loading data when WebBaseLoader returns empty list."""
mock_loader_instance = MagicMock()
mock_loader_instance.load.return_value = []
mock_web_base_loader_class = MagicMock()
mock_web_base_loader_class.return_value = mock_loader_instance
web_loader.loader = mock_web_base_loader_class
@patch("application.parser.remote.web_loader.pinned_request")
def test_load_data_empty_response_body(self, mock_pinned_request, mock_validate, web_loader):
mock_pinned_request.return_value = _fake_response("")
result = web_loader.load_data("https://empty-page.com")
assert result == []
assert len(result) == 1
assert result[0].text == ""
assert result[0].extra_info == {"source": "https://empty-page.com"}
def test_url_scheme_detection(self):
"""Test URL scheme detection logic."""
# Test URLs with schemes
assert urlparse("https://example.com").scheme == "https"
assert urlparse("http://example.com").scheme == "http"
assert urlparse("ftp://example.com").scheme == "ftp"
# Test URLs without schemes
assert urlparse("example.com").scheme == ""
assert urlparse("www.example.com").scheme == ""
@@ -351,29 +292,25 @@ class TestWebLoaderIntegration:
"""Test WebLoader integration with base class."""
def test_inherits_from_base_remote(self, web_loader):
"""Test that WebLoader inherits from BaseRemote."""
from application.parser.remote.base import BaseRemote
assert isinstance(web_loader, BaseRemote)
def test_implements_load_data_method(self, web_loader):
"""Test that WebLoader implements required load_data method."""
assert hasattr(web_loader, 'load_data')
assert hasattr(web_loader, "load_data")
assert callable(web_loader.load_data)
@patch("application.parser.remote.web_loader.validate_url", side_effect=_mock_validate_url)
def test_load_langchain_documents_method(self, mock_validate, web_loader, mock_langchain_document):
"""Test inherited load_langchain_documents method."""
mock_loader_instance = MagicMock()
mock_loader_instance.load.return_value = [mock_langchain_document]
mock_web_base_loader_class = MagicMock()
mock_web_base_loader_class.return_value = mock_loader_instance
web_loader.loader = mock_web_base_loader_class
@patch("application.parser.remote.web_loader.pinned_request")
def test_load_langchain_documents_method(self, mock_pinned_request, mock_validate, web_loader):
mock_pinned_request.return_value = _fake_response(
"<html><head><title>Test Page</title></head>"
"<body><p>Test web page content</p></body></html>"
)
result = web_loader.load_langchain_documents(inputs="https://example.com")
assert len(result) == 1
assert isinstance(result[0], LCDocument)
assert result[0].page_content == "Test web page content"
assert result[0].metadata == {"source": "https://example.com", "title": "Test Page"}
assert "Test web page content" in result[0].page_content
assert result[0].metadata["source"] == "https://example.com"
assert result[0].metadata["title"] == "Test Page"
tests/security/test_safe_url.py
+17
View File
@@ -354,6 +354,23 @@ def test_pinned_post_preserves_explicit_port(monkeypatch):
assert captured["prepared"].headers["Host"] == "api.example.com:8443"
@pytest.mark.unit
def test_pinned_post_preserves_url_userinfo(monkeypatch):
captured = _capture_send(monkeypatch)
with mock.patch("socket.getaddrinfo", return_value=_addrinfo("93.184.216.34")):
pinned_post(
"https://user:pass@api.example.com:8443/v1/test",
json={},
headers={},
timeout=5,
allow_redirects=False,
)
prepared = captured["prepared"]
assert prepared.url == "https://user:pass@93.184.216.34:8443/v1/test"
assert prepared.headers["Host"] == "api.example.com:8443"
assert prepared.headers["Authorization"] == "Basic dXNlcjpwYXNz"
@pytest.mark.unit
def test_pinned_post_handles_ip_literal_url_without_dns(monkeypatch):
"""If the URL already has an IP literal, no DNS lookup happens."""
tests/test_coverage_gaps.py
+1 -5
View File
@@ -1276,11 +1276,7 @@ class TestCrawlerMarkdownEdge:
loader = CrawlerLoader()
with patch(
"application.parser.remote.crawler_markdown.validate_url",
side_effect=lambda u: u,
), patch.object(
loader.session,
"get",
"application.parser.remote.crawler_markdown.pinned_request",
side_effect=requests.exceptions.ConnectionError("fail"),
):
result = loader._fetch_page("http://fail.com")
tests/test_remaining_coverage.py
+4 -2
View File
@@ -581,11 +581,13 @@ class TestReadWebpageErrors:
from application.agents.tools.read_webpage import ReadWebpageTool
tool = ReadWebpageTool({})
with patch("application.agents.tools.read_webpage.requests.get") as mock_get:
with patch(
"application.agents.tools.read_webpage.pinned_request"
) as mock_pinned_request:
mock_response = MagicMock()
mock_response.raise_for_status.return_value = None
mock_response.text = "<html><body>test</body></html>"
mock_get.return_value = mock_response
mock_pinned_request.return_value = mock_response
with patch(
"application.agents.tools.read_webpage.markdownify",
side_effect=Exception("parse error"),