baee4949d3
Rename LlmConfig to LLMConfig across the codebase to follow consistent naming conventions. Update all imports and usages to use the new name. Update documentation and examples to reflect the change. BREAKING CHANGE: LlmConfig has been renamed to LLMConfig. Users need to update their imports and usage.
150 lines
5.3 KiB
Python
150 lines
5.3 KiB
Python
import unittest, os
|
|
from crawl4ai.types import LLMConfig
|
|
from crawl4ai.web_crawler import WebCrawler
|
|
from crawl4ai.chunking_strategy import (
|
|
RegexChunking,
|
|
FixedLengthWordChunking,
|
|
SlidingWindowChunking,
|
|
)
|
|
from crawl4ai.extraction_strategy import (
|
|
CosineStrategy,
|
|
LLMExtractionStrategy,
|
|
TopicExtractionStrategy,
|
|
NoExtractionStrategy,
|
|
)
|
|
|
|
|
|
class TestWebCrawler(unittest.TestCase):
|
|
def setUp(self):
|
|
self.crawler = WebCrawler()
|
|
|
|
def test_warmup(self):
|
|
self.crawler.warmup()
|
|
self.assertTrue(self.crawler.ready, "WebCrawler failed to warm up")
|
|
|
|
def test_run_default_strategies(self):
|
|
result = self.crawler.run(
|
|
url="https://www.nbcnews.com/business",
|
|
word_count_threshold=5,
|
|
chunking_strategy=RegexChunking(),
|
|
extraction_strategy=CosineStrategy(),
|
|
bypass_cache=True,
|
|
)
|
|
self.assertTrue(
|
|
result.success, "Failed to crawl and extract using default strategies"
|
|
)
|
|
|
|
def test_run_different_strategies(self):
|
|
url = "https://www.nbcnews.com/business"
|
|
|
|
# Test with FixedLengthWordChunking and LLMExtractionStrategy
|
|
result = self.crawler.run(
|
|
url=url,
|
|
word_count_threshold=5,
|
|
chunking_strategy=FixedLengthWordChunking(chunk_size=100),
|
|
extraction_strategy=LLMExtractionStrategy(
|
|
llm_config=LLMConfig(provider="openai/gpt-3.5-turbo", api_token=os.getenv("OPENAI_API_KEY"))
|
|
),
|
|
bypass_cache=True,
|
|
)
|
|
self.assertTrue(
|
|
result.success,
|
|
"Failed to crawl and extract with FixedLengthWordChunking and LLMExtractionStrategy",
|
|
)
|
|
|
|
# Test with SlidingWindowChunking and TopicExtractionStrategy
|
|
result = self.crawler.run(
|
|
url=url,
|
|
word_count_threshold=5,
|
|
chunking_strategy=SlidingWindowChunking(window_size=100, step=50),
|
|
extraction_strategy=TopicExtractionStrategy(num_keywords=5),
|
|
bypass_cache=True,
|
|
)
|
|
self.assertTrue(
|
|
result.success,
|
|
"Failed to crawl and extract with SlidingWindowChunking and TopicExtractionStrategy",
|
|
)
|
|
|
|
def test_invalid_url(self):
|
|
with self.assertRaises(Exception) as context:
|
|
self.crawler.run(url="invalid_url", bypass_cache=True)
|
|
self.assertIn("Invalid URL", str(context.exception))
|
|
|
|
def test_unsupported_extraction_strategy(self):
|
|
with self.assertRaises(Exception) as context:
|
|
self.crawler.run(
|
|
url="https://www.nbcnews.com/business",
|
|
extraction_strategy="UnsupportedStrategy",
|
|
bypass_cache=True,
|
|
)
|
|
self.assertIn("Unsupported extraction strategy", str(context.exception))
|
|
|
|
def test_invalid_css_selector(self):
|
|
with self.assertRaises(ValueError) as context:
|
|
self.crawler.run(
|
|
url="https://www.nbcnews.com/business",
|
|
css_selector="invalid_selector",
|
|
bypass_cache=True,
|
|
)
|
|
self.assertIn("Invalid CSS selector", str(context.exception))
|
|
|
|
def test_crawl_with_cache_and_bypass_cache(self):
|
|
url = "https://www.nbcnews.com/business"
|
|
|
|
# First crawl with cache enabled
|
|
result = self.crawler.run(url=url, bypass_cache=False)
|
|
self.assertTrue(result.success, "Failed to crawl and cache the result")
|
|
|
|
# Second crawl with bypass_cache=True
|
|
result = self.crawler.run(url=url, bypass_cache=True)
|
|
self.assertTrue(result.success, "Failed to bypass cache and fetch fresh data")
|
|
|
|
def test_fetch_multiple_pages(self):
|
|
urls = ["https://www.nbcnews.com/business", "https://www.bbc.com/news"]
|
|
results = []
|
|
for url in urls:
|
|
result = self.crawler.run(
|
|
url=url,
|
|
word_count_threshold=5,
|
|
chunking_strategy=RegexChunking(),
|
|
extraction_strategy=CosineStrategy(),
|
|
bypass_cache=True,
|
|
)
|
|
results.append(result)
|
|
|
|
self.assertEqual(len(results), 2, "Failed to crawl and extract multiple pages")
|
|
for result in results:
|
|
self.assertTrue(
|
|
result.success, "Failed to crawl and extract a page in the list"
|
|
)
|
|
|
|
def test_run_fixed_length_word_chunking_and_no_extraction(self):
|
|
result = self.crawler.run(
|
|
url="https://www.nbcnews.com/business",
|
|
word_count_threshold=5,
|
|
chunking_strategy=FixedLengthWordChunking(chunk_size=100),
|
|
extraction_strategy=NoExtractionStrategy(),
|
|
bypass_cache=True,
|
|
)
|
|
self.assertTrue(
|
|
result.success,
|
|
"Failed to crawl and extract with FixedLengthWordChunking and NoExtractionStrategy",
|
|
)
|
|
|
|
def test_run_sliding_window_and_no_extraction(self):
|
|
result = self.crawler.run(
|
|
url="https://www.nbcnews.com/business",
|
|
word_count_threshold=5,
|
|
chunking_strategy=SlidingWindowChunking(window_size=100, step=50),
|
|
extraction_strategy=NoExtractionStrategy(),
|
|
bypass_cache=True,
|
|
)
|
|
self.assertTrue(
|
|
result.success,
|
|
"Failed to crawl and extract with SlidingWindowChunking and NoExtractionStrategy",
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|