372 lines
13 KiB
Python
372 lines
13 KiB
Python
import asyncio
|
|
import os
|
|
import unittest
|
|
from unittest.mock import MagicMock, AsyncMock, patch
|
|
|
|
from crawler.config import CrawlerConfig
|
|
from crawler.engine import Crawler, IterationVisitTracker
|
|
|
|
|
|
class TestIterationVisitsConfig(unittest.TestCase):
|
|
|
|
def test_default_config(self):
|
|
config = CrawlerConfig(start_url="https://example.com")
|
|
self.assertEqual(config.iteration_mode, "all")
|
|
self.assertEqual(config.max_visits_per_iteration, 10)
|
|
self.assertEqual(config.visits_per_iteration, 10)
|
|
self.assertTrue(config.randomize_queue)
|
|
|
|
def test_dict_serialization(self):
|
|
config = CrawlerConfig(
|
|
start_url="https://example.com",
|
|
num_iterations=3,
|
|
iteration_mode="visits",
|
|
max_visits_per_iteration=25,
|
|
randomize_queue=False,
|
|
)
|
|
d = config.to_dict()
|
|
self.assertEqual(d["iteration_mode"], "visits")
|
|
self.assertEqual(d["max_visits_per_iteration"], 25)
|
|
self.assertFalse(d["randomize_queue"])
|
|
|
|
# Restore from dict
|
|
restored = CrawlerConfig.from_dict(d)
|
|
self.assertEqual(restored.iteration_mode, "visits")
|
|
self.assertEqual(restored.max_visits_per_iteration, 25)
|
|
self.assertEqual(restored.visits_per_iteration, 25)
|
|
self.assertFalse(restored.randomize_queue)
|
|
|
|
def test_from_dict_alias_and_normalization(self):
|
|
# Using visits_per_iteration alias
|
|
cfg1 = CrawlerConfig.from_dict({
|
|
"start_url": "https://example.com",
|
|
"iteration_mode": "limit",
|
|
"visits_per_iteration": 15,
|
|
"randomize_queue": True,
|
|
})
|
|
self.assertEqual(cfg1.iteration_mode, "visits")
|
|
self.assertEqual(cfg1.max_visits_per_iteration, 15)
|
|
self.assertTrue(cfg1.randomize_queue)
|
|
|
|
# Non-visits mode defaults to "all"
|
|
cfg2 = CrawlerConfig.from_dict({
|
|
"start_url": "https://example.com",
|
|
"iteration_mode": "unknown_mode",
|
|
"max_visits_per_iteration": -5,
|
|
})
|
|
self.assertEqual(cfg2.iteration_mode, "all")
|
|
self.assertEqual(cfg2.max_visits_per_iteration, 10)
|
|
self.assertTrue(cfg2.randomize_queue)
|
|
|
|
def test_env_dict_loading(self):
|
|
with patch.dict(os.environ, {
|
|
"CRAWLER_START_URL": "https://test.com",
|
|
"CRAWLER_ITERATION_MODE": "visits",
|
|
"CRAWLER_MAX_VISITS_PER_ITERATION": "42",
|
|
"CRAWLER_RANDOMIZE_QUEUE": "false",
|
|
}):
|
|
env_data = CrawlerConfig.get_env_dict(env_path="nonexistent_test.env")
|
|
self.assertEqual(env_data["iteration_mode"], "visits")
|
|
self.assertEqual(env_data["max_visits_per_iteration"], 42)
|
|
self.assertFalse(env_data["randomize_queue"])
|
|
|
|
|
|
class TestIterationVisitTracker(unittest.IsolatedAsyncioTestCase):
|
|
|
|
async def test_unlimited_tracker(self):
|
|
tracker = IterationVisitTracker(max_visits=None)
|
|
for _ in range(50):
|
|
can_visit = await tracker.acquire_visit_slot()
|
|
self.assertTrue(can_visit)
|
|
await tracker.record_visit_completed()
|
|
self.assertFalse(tracker.limit_reached_event.is_set())
|
|
|
|
async def test_limited_tracker_slot_reservation_and_completion(self):
|
|
limit = 3
|
|
tracker = IterationVisitTracker(max_visits=limit)
|
|
|
|
# Acquire 3 slots
|
|
self.assertTrue(await tracker.acquire_visit_slot())
|
|
self.assertTrue(await tracker.acquire_visit_slot())
|
|
self.assertTrue(await tracker.acquire_visit_slot())
|
|
|
|
# 4th acquisition attempt fails
|
|
self.assertFalse(await tracker.acquire_visit_slot())
|
|
self.assertEqual(tracker.started_visits, 3)
|
|
self.assertFalse(tracker.limit_reached_event.is_set())
|
|
|
|
# Record completions
|
|
await tracker.record_visit_completed()
|
|
self.assertFalse(tracker.limit_reached_event.is_set())
|
|
await tracker.record_visit_completed()
|
|
self.assertFalse(tracker.limit_reached_event.is_set())
|
|
|
|
# 3rd completion triggers limit_reached_event
|
|
await tracker.record_visit_completed()
|
|
self.assertTrue(tracker.limit_reached_event.is_set())
|
|
self.assertEqual(tracker.completed_visits, 3)
|
|
|
|
|
|
class TestCrawlerIterationVisitCap(unittest.IsolatedAsyncioTestCase):
|
|
|
|
async def test_iteration_stops_at_configured_visits(self):
|
|
config = CrawlerConfig(
|
|
start_url="https://example.com",
|
|
num_tabs=3,
|
|
num_iterations=1,
|
|
iteration_mode="visits",
|
|
max_visits_per_iteration=3,
|
|
)
|
|
|
|
visited_urls = []
|
|
|
|
async def mock_crawl_level(
|
|
page,
|
|
url,
|
|
source_url,
|
|
visited,
|
|
errors,
|
|
start_domain,
|
|
accept_cookies_default=True,
|
|
stats=None,
|
|
timeout_ms=35000,
|
|
event_callback=None,
|
|
worker_name="Worker",
|
|
page_result_callback=None,
|
|
):
|
|
visited_urls.append(url)
|
|
visited.add(url)
|
|
await asyncio.sleep(0.01)
|
|
# Return new links to simulate queue populated with many items
|
|
return {
|
|
f"https://example.com/sub_{i}_{len(visited_urls)}" for i in range(5)
|
|
}
|
|
|
|
crawler = Crawler(config=config)
|
|
|
|
mock_page = MagicMock()
|
|
mock_page.close = AsyncMock()
|
|
mock_page.on = MagicMock()
|
|
|
|
mock_context = MagicMock()
|
|
mock_context.new_page = AsyncMock(return_value=mock_page)
|
|
|
|
mock_browser = MagicMock()
|
|
mock_browser.new_context = AsyncMock(return_value=mock_context)
|
|
mock_browser.close = AsyncMock()
|
|
|
|
mock_playwright = MagicMock()
|
|
mock_playwright.chromium.launch = AsyncMock(return_value=mock_browser)
|
|
|
|
with patch("crawler.engine.crawl_level", side_effect=mock_crawl_level):
|
|
result = await crawler.run_iteration(mock_playwright, iteration=1)
|
|
|
|
# Verified that exactly max_visits_per_iteration (3) were crawled despite many discovered links
|
|
self.assertEqual(len(visited_urls), 3)
|
|
self.assertEqual(result["stats"]["internal"], 3)
|
|
|
|
async def test_multi_iteration_stops_at_configured_visits_each(self):
|
|
config = CrawlerConfig(
|
|
start_url="https://example.com",
|
|
num_tabs=2,
|
|
num_iterations=2,
|
|
iteration_mode="visits",
|
|
max_visits_per_iteration=2,
|
|
)
|
|
|
|
visited_urls_per_iteration = []
|
|
|
|
async def mock_crawl_level(
|
|
page,
|
|
url,
|
|
source_url,
|
|
visited,
|
|
errors,
|
|
start_domain,
|
|
accept_cookies_default=True,
|
|
stats=None,
|
|
timeout_ms=35000,
|
|
event_callback=None,
|
|
worker_name="Worker",
|
|
page_result_callback=None,
|
|
):
|
|
visited_urls_per_iteration.append(url)
|
|
visited.add(url)
|
|
await asyncio.sleep(0.01)
|
|
return {f"https://example.com/link_{i}" for i in range(5)}
|
|
|
|
crawler = Crawler(config=config)
|
|
|
|
mock_page = MagicMock()
|
|
mock_page.close = AsyncMock()
|
|
mock_page.on = MagicMock()
|
|
|
|
mock_context = MagicMock()
|
|
mock_context.new_page = AsyncMock(return_value=mock_page)
|
|
|
|
mock_browser = MagicMock()
|
|
mock_browser.new_context = AsyncMock(return_value=mock_context)
|
|
mock_browser.close = AsyncMock()
|
|
|
|
mock_playwright = MagicMock()
|
|
mock_playwright.chromium.launch = AsyncMock(return_value=mock_browser)
|
|
|
|
with patch("crawler.engine.crawl_level", side_effect=mock_crawl_level):
|
|
# Iteration 1
|
|
res1 = await crawler.run_iteration(mock_playwright, iteration=1)
|
|
self.assertEqual(len(visited_urls_per_iteration), 2)
|
|
self.assertEqual(res1["stats"]["internal"], 2)
|
|
|
|
visited_urls_per_iteration.clear()
|
|
|
|
# Iteration 2
|
|
res2 = await crawler.run_iteration(mock_playwright, iteration=2)
|
|
self.assertEqual(len(visited_urls_per_iteration), 2)
|
|
self.assertEqual(res2["stats"]["internal"], 2)
|
|
|
|
|
|
class TestRandomQueue(unittest.IsolatedAsyncioTestCase):
|
|
|
|
async def test_random_queue_pops_all_items(self):
|
|
from crawler.engine import RandomQueue
|
|
|
|
q = RandomQueue()
|
|
items = [f"https://example.com/page{i}" for i in range(20)]
|
|
for it in items:
|
|
q.put_nowait(it)
|
|
|
|
self.assertEqual(q.qsize(), 20)
|
|
popped = []
|
|
for _ in range(20):
|
|
popped.append(await q.get())
|
|
q.task_done()
|
|
|
|
await q.join()
|
|
self.assertEqual(sorted(popped), sorted(items))
|
|
self.assertEqual(q.qsize(), 0)
|
|
|
|
async def test_crawler_uses_random_queue_when_configured(self):
|
|
from crawler.engine import RandomQueue
|
|
|
|
cfg_random = CrawlerConfig(start_url="https://example.com", randomize_queue=True)
|
|
crawler_random = Crawler(config=cfg_random)
|
|
|
|
mock_page = MagicMock()
|
|
mock_page.close = AsyncMock()
|
|
mock_page.on = MagicMock()
|
|
mock_context = MagicMock()
|
|
mock_context.new_page = AsyncMock(return_value=mock_page)
|
|
mock_browser = MagicMock()
|
|
mock_browser.new_context = AsyncMock(return_value=mock_context)
|
|
mock_browser.close = AsyncMock()
|
|
mock_playwright = MagicMock()
|
|
mock_playwright.chromium.launch = AsyncMock(return_value=mock_browser)
|
|
|
|
# Intercept queue creation by checking instance type
|
|
async def dummy_crawl(*args, **kwargs):
|
|
return set()
|
|
|
|
with patch("crawler.engine.crawl_level", side_effect=dummy_crawl):
|
|
with patch("crawler.engine.RandomQueue", wraps=RandomQueue) as mock_rq:
|
|
await crawler_random.run_iteration(mock_playwright, iteration=1)
|
|
self.assertTrue(mock_rq.called)
|
|
|
|
|
|
class TestCumulativeAllIterationsStatistics(unittest.IsolatedAsyncioTestCase):
|
|
|
|
async def test_cumulative_visited_and_errors_across_iterations(self):
|
|
config = CrawlerConfig(
|
|
start_url="https://example.com",
|
|
num_tabs=2,
|
|
num_iterations=2,
|
|
iteration_mode="visits",
|
|
max_visits_per_iteration=3,
|
|
)
|
|
|
|
events = []
|
|
|
|
async def event_cb(ev):
|
|
events.append(ev)
|
|
|
|
crawler = Crawler(config=config, event_callback=event_cb)
|
|
|
|
visited_per_call = 0
|
|
|
|
async def mock_crawl_level(
|
|
page,
|
|
url,
|
|
source_url,
|
|
visited,
|
|
errors,
|
|
start_domain,
|
|
accept_cookies_default=True,
|
|
stats=None,
|
|
timeout_ms=35000,
|
|
event_callback=None,
|
|
worker_name="Worker",
|
|
page_result_callback=None,
|
|
):
|
|
nonlocal visited_per_call
|
|
visited_per_call += 1
|
|
visited.add(url)
|
|
if visited_per_call == 2:
|
|
errors.append((url, source_url, "Simulated 500 error"))
|
|
await asyncio.sleep(0.01)
|
|
return {f"https://example.com/link_{visited_per_call}_{i}" for i in range(3)}
|
|
|
|
mock_page = MagicMock()
|
|
mock_page.close = AsyncMock()
|
|
mock_page.on = MagicMock()
|
|
mock_context = MagicMock()
|
|
mock_context.new_page = AsyncMock(return_value=mock_page)
|
|
mock_browser = MagicMock()
|
|
mock_browser.new_context = AsyncMock(return_value=mock_context)
|
|
mock_browser.close = AsyncMock()
|
|
mock_playwright = MagicMock()
|
|
mock_playwright.chromium.launch = AsyncMock(return_value=mock_browser)
|
|
|
|
with patch("crawler.engine.crawl_level", side_effect=mock_crawl_level):
|
|
res1 = await crawler.run_iteration(mock_playwright, iteration=1)
|
|
self.assertEqual(res1["visited_count"], 3)
|
|
self.assertEqual(res1["total_visited_all"], 3)
|
|
self.assertEqual(res1["total_errors_all"], 1)
|
|
|
|
res2 = await crawler.run_iteration(mock_playwright, iteration=2)
|
|
self.assertEqual(res2["visited_count"], 3)
|
|
self.assertEqual(res2["total_visited_all"], 6)
|
|
self.assertEqual(crawler.total_visited_all_iterations, 6)
|
|
|
|
# Check emitted stats events
|
|
stats_events = [e for e in events if e.get("type") == "stats"]
|
|
self.assertTrue(len(stats_events) > 0)
|
|
final_stats_event = stats_events[-1]
|
|
self.assertEqual(final_stats_event["total_visited_all"], 6)
|
|
|
|
async def test_session_manager_snapshot_and_broadcast(self):
|
|
from web.manager import CrawlSessionManager
|
|
|
|
manager = CrawlSessionManager()
|
|
snapshot = manager.get_snapshot()
|
|
self.assertIn("total_visited_all", snapshot)
|
|
self.assertIn("total_errors_all", snapshot)
|
|
self.assertEqual(snapshot["total_visited_all"], 0)
|
|
self.assertEqual(snapshot["total_errors_all"], 0)
|
|
|
|
# Simulate receiving crawler stats event
|
|
await manager.handle_crawler_event({
|
|
"type": "stats",
|
|
"visited_count": 5,
|
|
"total_visited_all": 15,
|
|
"total_errors_all": 2,
|
|
"queue_size": 3,
|
|
"stats": {"internal": 5, "external": 0, "cookies_accepted": 0},
|
|
})
|
|
self.assertEqual(manager.total_visited_all, 15)
|
|
self.assertEqual(manager.total_errors_all, 2)
|
|
snapshot_after = manager.get_snapshot()
|
|
self.assertEqual(snapshot_after["total_visited_all"], 15)
|
|
self.assertEqual(snapshot_after["total_errors_all"], 2)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|