248 lines
8.2 KiB
Python
248 lines
8.2 KiB
Python
"""Tests for method web.py"""
|
|
|
|
import json
|
|
import logging
|
|
|
|
import pytest
|
|
from unittest.mock import patch
|
|
from modules.methods.method_web import MethodWeb
|
|
from tests.fixtures.tavily_responses import (
|
|
RESPONSE_SEARCH_HOW_OLD_IS_EARTH_NO_ANSWER,
|
|
RESPONSE_EXTRACT_HOW_OLD_IS_EARTH_NO_ANSWER,
|
|
)
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.expensive
|
|
async def test_method_web_search_live():
|
|
"""Tests method web search with live API calls."""
|
|
|
|
logger.info("=" * 50)
|
|
logger.info("==> Test: Method Web Search Live")
|
|
|
|
method_web = MethodWeb(serviceCenter=None)
|
|
|
|
# Actual request
|
|
action_result = await method_web.search(
|
|
{"query": "How old is the earth", "maxResults": 5}
|
|
)
|
|
|
|
# Evaluate results
|
|
assert action_result.success
|
|
assert len(action_result.documents) > 0
|
|
|
|
logger.info(f"Action result success status: {action_result.success}")
|
|
logger.info(f"Action result error: {action_result.error}")
|
|
logger.info(f"Action result label: {action_result.resultLabel}")
|
|
|
|
logger.info("Documents:")
|
|
for doc in action_result.documents:
|
|
logger.info(f" - Document Name: {doc.documentName}")
|
|
logger.info(f" --> Document Mime Type: {doc.mimeType}")
|
|
logger.info(f" --> Document Data: {doc.documentData}")
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_method_web_search_dummy():
|
|
"""Tests method web search with dummy response data - no external API calls."""
|
|
|
|
logger.info("=" * 50)
|
|
logger.info("==> Test: Method Web Search Dummy")
|
|
|
|
method_web = MethodWeb(serviceCenter=None)
|
|
|
|
# Mock the Tavily API response
|
|
with patch(
|
|
"tavily.AsyncTavilyClient.search",
|
|
return_value=RESPONSE_SEARCH_HOW_OLD_IS_EARTH_NO_ANSWER,
|
|
) as mock_client:
|
|
action_result = await method_web.search(
|
|
{"query": "How old is the earth", "maxResults": 5}
|
|
)
|
|
mock_client.assert_called_once()
|
|
|
|
# Evaluate results
|
|
assert action_result.success
|
|
assert len(action_result.documents) > 0
|
|
|
|
logger.info(f"Action result success status: {action_result.success}")
|
|
logger.info(f"Action result error: {action_result.error}")
|
|
logger.info(f"Action result label: {action_result.resultLabel}")
|
|
|
|
logger.info("Documents:")
|
|
for doc in action_result.documents:
|
|
logger.info(f" - Document Name: {doc.documentName}")
|
|
logger.info(f" --> Document Mime Type: {doc.mimeType}")
|
|
logger.info(f" --> Document Data: {doc.documentData}")
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.expensive
|
|
async def test_method_web_crawl_live():
|
|
"""Tests method web crawl with live API calls."""
|
|
|
|
logger.info("=" * 50)
|
|
logger.info("==> Test: Method Web Crawl Live")
|
|
|
|
method_web = MethodWeb(serviceCenter=None)
|
|
|
|
# Create mock document data with URLs from search results
|
|
search_results_json = {
|
|
"documentData": {
|
|
"results": [
|
|
{"url": "https://en.wikipedia.org/wiki/Age_of_Earth"},
|
|
{"url": "https://www.planetary.org/articles/how-old-is-the-earth"},
|
|
]
|
|
}
|
|
}
|
|
|
|
# Mock the service center methods
|
|
with patch.object(method_web, "service") as mock_service:
|
|
mock_service.getChatDocumentsFromDocumentList.return_value = [
|
|
type("MockDoc", (), {"fileId": "test-file-id"})()
|
|
]
|
|
mock_service.getFileData.return_value = json.dumps(search_results_json).encode(
|
|
"utf-8"
|
|
)
|
|
|
|
# Actual request
|
|
action_result = await method_web.crawl({"document": "test-document-ref"})
|
|
|
|
# Evaluate results
|
|
assert action_result.success
|
|
assert len(action_result.documents) > 0
|
|
|
|
logger.info(f"Action result success status: {action_result.success}")
|
|
logger.info(f"Action result error: {action_result.error}")
|
|
logger.info(f"Action result label: {action_result.resultLabel}")
|
|
|
|
logger.info("Documents:")
|
|
for doc in action_result.documents:
|
|
logger.info(f" - Document Name: {doc.documentName}")
|
|
logger.info(f" --> Document Mime Type: {doc.mimeType}")
|
|
logger.info(f" --> Document Data: {doc.documentData}")
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_method_web_crawl_dummy():
|
|
"""Tests method web crawl with dummy response data - no external API calls."""
|
|
|
|
logger.info("=" * 50)
|
|
logger.info("==> Test: Method Web Crawl Dummy")
|
|
|
|
method_web = MethodWeb(serviceCenter=None)
|
|
|
|
# Create mock document data with URLs from search results
|
|
search_results_json = {
|
|
"documentData": {
|
|
"results": [
|
|
{"url": "https://en.wikipedia.org/wiki/Age_of_Earth"},
|
|
{"url": "https://www.planetary.org/articles/how-old-is-the-earth"},
|
|
]
|
|
}
|
|
}
|
|
|
|
# Mock both the service center and Tavily API
|
|
with (
|
|
patch.object(method_web, "service") as mock_service,
|
|
patch(
|
|
"tavily.AsyncTavilyClient.extract",
|
|
return_value=RESPONSE_EXTRACT_HOW_OLD_IS_EARTH_NO_ANSWER,
|
|
) as mock_client,
|
|
):
|
|
mock_service.getChatDocumentsFromDocumentList.return_value = [
|
|
type("MockDoc", (), {"fileId": "test-file-id"})()
|
|
]
|
|
mock_service.getFileData.return_value = json.dumps(search_results_json).encode(
|
|
"utf-8"
|
|
)
|
|
|
|
action_result = await method_web.crawl({"document": "test-document-ref"})
|
|
mock_client.assert_called_once()
|
|
|
|
# Evaluate results
|
|
assert action_result.success
|
|
assert len(action_result.documents) > 0
|
|
|
|
logger.info(f"Action result success status: {action_result.success}")
|
|
logger.info(f"Action result error: {action_result.error}")
|
|
logger.info(f"Action result label: {action_result.resultLabel}")
|
|
|
|
logger.info("Documents:")
|
|
for doc in action_result.documents:
|
|
logger.info(f" - Document Name: {doc.documentName}")
|
|
logger.info(f" --> Document Mime Type: {doc.mimeType}")
|
|
logger.info(f" --> Document Data: {doc.documentData}")
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.expensive
|
|
async def test_method_web_scrape_live():
|
|
"""Tests method web scrape with live API calls."""
|
|
|
|
logger.info("=" * 50)
|
|
logger.info("==> Test: Method Web Scrape Live")
|
|
|
|
method_web = MethodWeb(serviceCenter=None)
|
|
|
|
# Actual request
|
|
action_result = await method_web.scrape(
|
|
{"query": "How old is the earth", "maxResults": 3}
|
|
)
|
|
|
|
# Evaluate results
|
|
assert action_result.success
|
|
assert len(action_result.documents) > 0
|
|
|
|
logger.info(f"Action result success status: {action_result.success}")
|
|
logger.info(f"Action result error: {action_result.error}")
|
|
logger.info(f"Action result label: {action_result.resultLabel}")
|
|
|
|
logger.info("Documents:")
|
|
for doc in action_result.documents:
|
|
logger.info(f" - Document Name: {doc.documentName}")
|
|
logger.info(f" --> Document Mime Type: {doc.mimeType}")
|
|
logger.info(f" --> Document Data: {doc.documentData}")
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_method_web_scrape_dummy():
|
|
"""Tests method web scrape with dummy response data - no external API calls."""
|
|
|
|
logger.info("=" * 50)
|
|
logger.info("==> Test: Method Web Scrape Dummy")
|
|
|
|
method_web = MethodWeb(serviceCenter=None)
|
|
|
|
# Mock both Tavily API responses (search + extract)
|
|
with (
|
|
patch(
|
|
"tavily.AsyncTavilyClient.search",
|
|
return_value=RESPONSE_SEARCH_HOW_OLD_IS_EARTH_NO_ANSWER,
|
|
) as mock_search,
|
|
patch(
|
|
"tavily.AsyncTavilyClient.extract",
|
|
return_value=RESPONSE_EXTRACT_HOW_OLD_IS_EARTH_NO_ANSWER,
|
|
) as mock_extract,
|
|
):
|
|
action_result = await method_web.scrape(
|
|
{"query": "How old is the earth", "maxResults": 3}
|
|
)
|
|
mock_search.assert_called_once()
|
|
mock_extract.assert_called_once()
|
|
|
|
# Evaluate results
|
|
assert action_result.success
|
|
assert len(action_result.documents) > 0
|
|
|
|
logger.info(f"Action result success status: {action_result.success}")
|
|
logger.info(f"Action result error: {action_result.error}")
|
|
logger.info(f"Action result label: {action_result.resultLabel}")
|
|
|
|
logger.info("Documents:")
|
|
for doc in action_result.documents:
|
|
logger.info(f" - Document Name: {doc.documentName}")
|
|
logger.info(f" --> Document Mime Type: {doc.mimeType}")
|
|
logger.info(f" --> Document Data: {doc.documentData}")
|