chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,597 @@
|
||||
"""
|
||||
Test Suite: Deep Crawl Cancellation Tests
|
||||
|
||||
Tests that verify:
|
||||
1. should_cancel callback is called before each URL
|
||||
2. cancel() method immediately stops the crawl
|
||||
3. cancelled property correctly reflects state
|
||||
4. Strategy reuse works after cancellation
|
||||
5. Both sync and async should_cancel callbacks work
|
||||
6. Callback exceptions don't crash the crawl
|
||||
7. State notifications include cancelled flag
|
||||
"""
|
||||
|
||||
import pytest
|
||||
import asyncio
|
||||
from typing import Dict, Any, List
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
from crawl4ai.deep_crawling import (
|
||||
BFSDeepCrawlStrategy,
|
||||
DFSDeepCrawlStrategy,
|
||||
BestFirstCrawlingStrategy,
|
||||
)
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Helper Functions for Mock Crawler
|
||||
# ============================================================================
|
||||
|
||||
def create_mock_config(stream=False):
|
||||
"""Create a mock CrawlerRunConfig."""
|
||||
config = MagicMock()
|
||||
config.stream = stream
|
||||
|
||||
def clone_config(**kwargs):
|
||||
"""Clone returns a new config with overridden values."""
|
||||
new_config = MagicMock()
|
||||
new_config.stream = kwargs.get('stream', stream)
|
||||
new_config.clone = MagicMock(side_effect=clone_config)
|
||||
return new_config
|
||||
|
||||
config.clone = MagicMock(side_effect=clone_config)
|
||||
return config
|
||||
|
||||
|
||||
def create_mock_crawler_with_links(num_links: int = 3):
|
||||
"""Create mock crawler that returns results with links."""
|
||||
call_count = 0
|
||||
|
||||
async def mock_arun_many(urls, config):
|
||||
nonlocal call_count
|
||||
results = []
|
||||
for url in urls:
|
||||
call_count += 1
|
||||
result = MagicMock()
|
||||
result.url = url
|
||||
result.success = True
|
||||
result.metadata = {}
|
||||
|
||||
# Generate child links
|
||||
links = []
|
||||
for i in range(num_links):
|
||||
link_url = f"{url}/child{call_count}_{i}"
|
||||
links.append({"href": link_url})
|
||||
|
||||
result.links = {"internal": links, "external": []}
|
||||
results.append(result)
|
||||
|
||||
# For streaming mode, return async generator
|
||||
if config.stream:
|
||||
async def gen():
|
||||
for r in results:
|
||||
yield r
|
||||
return gen()
|
||||
return results
|
||||
|
||||
crawler = MagicMock()
|
||||
crawler.arun_many = mock_arun_many
|
||||
return crawler
|
||||
|
||||
|
||||
def create_mock_crawler_tracking(crawl_order: List[str], return_no_links: bool = False):
|
||||
"""Create mock crawler that tracks crawl order."""
|
||||
|
||||
async def mock_arun_many(urls, config):
|
||||
results = []
|
||||
for url in urls:
|
||||
crawl_order.append(url)
|
||||
result = MagicMock()
|
||||
result.url = url
|
||||
result.success = True
|
||||
result.metadata = {}
|
||||
result.links = {"internal": [], "external": []} if return_no_links else {"internal": [{"href": f"{url}/child"}], "external": []}
|
||||
results.append(result)
|
||||
|
||||
# For streaming mode, return async generator
|
||||
if config.stream:
|
||||
async def gen():
|
||||
for r in results:
|
||||
yield r
|
||||
return gen()
|
||||
return results
|
||||
|
||||
crawler = MagicMock()
|
||||
crawler.arun_many = mock_arun_many
|
||||
return crawler
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# TEST SUITE: Cancellation via should_cancel Callback
|
||||
# ============================================================================
|
||||
|
||||
class TestBFSCancellation:
|
||||
"""BFS strategy cancellation tests."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_cancel_via_async_callback(self):
|
||||
"""Verify async should_cancel callback stops crawl."""
|
||||
pages_crawled = 0
|
||||
cancel_after = 3
|
||||
|
||||
async def check_cancel():
|
||||
return pages_crawled >= cancel_after
|
||||
|
||||
async def track_pages(state: Dict[str, Any]):
|
||||
nonlocal pages_crawled
|
||||
pages_crawled = state.get("pages_crawled", 0)
|
||||
|
||||
strategy = BFSDeepCrawlStrategy(
|
||||
max_depth=5,
|
||||
max_pages=100,
|
||||
should_cancel=check_cancel,
|
||||
on_state_change=track_pages,
|
||||
)
|
||||
|
||||
mock_crawler = create_mock_crawler_with_links(num_links=5)
|
||||
mock_config = create_mock_config()
|
||||
|
||||
results = await strategy._arun_batch("https://example.com", mock_crawler, mock_config)
|
||||
|
||||
# Should have stopped after cancel_after pages
|
||||
assert strategy.cancelled == True
|
||||
assert strategy._pages_crawled >= cancel_after
|
||||
assert strategy._pages_crawled < 100 # Should not have crawled all pages
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_cancel_via_sync_callback(self):
|
||||
"""Verify sync should_cancel callback works."""
|
||||
cancel_flag = False
|
||||
|
||||
def check_cancel():
|
||||
return cancel_flag
|
||||
|
||||
async def set_cancel_after_3(state: Dict[str, Any]):
|
||||
nonlocal cancel_flag
|
||||
if state.get("pages_crawled", 0) >= 3:
|
||||
cancel_flag = True
|
||||
|
||||
strategy = BFSDeepCrawlStrategy(
|
||||
max_depth=5,
|
||||
max_pages=100,
|
||||
should_cancel=check_cancel,
|
||||
on_state_change=set_cancel_after_3,
|
||||
)
|
||||
|
||||
mock_crawler = create_mock_crawler_with_links(num_links=5)
|
||||
mock_config = create_mock_config()
|
||||
|
||||
await strategy._arun_batch("https://example.com", mock_crawler, mock_config)
|
||||
|
||||
assert strategy.cancelled == True
|
||||
assert strategy._pages_crawled >= 3
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_cancel_method_stops_crawl(self):
|
||||
"""Verify cancel() method immediately stops the crawl."""
|
||||
strategy = BFSDeepCrawlStrategy(
|
||||
max_depth=5,
|
||||
max_pages=100,
|
||||
)
|
||||
|
||||
async def cancel_after_2_pages(state: Dict[str, Any]):
|
||||
if state.get("pages_crawled", 0) >= 2:
|
||||
strategy.cancel()
|
||||
|
||||
strategy._on_state_change = cancel_after_2_pages
|
||||
|
||||
mock_crawler = create_mock_crawler_with_links(num_links=5)
|
||||
mock_config = create_mock_config()
|
||||
|
||||
await strategy._arun_batch("https://example.com", mock_crawler, mock_config)
|
||||
|
||||
assert strategy.cancelled == True
|
||||
assert strategy._pages_crawled >= 2
|
||||
assert strategy._pages_crawled < 100
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_cancelled_property_reflects_state(self):
|
||||
"""Verify cancelled property correctly reflects state."""
|
||||
strategy = BFSDeepCrawlStrategy(max_depth=2, max_pages=10)
|
||||
|
||||
# Before cancel
|
||||
assert strategy.cancelled == False
|
||||
|
||||
# After cancel()
|
||||
strategy.cancel()
|
||||
assert strategy.cancelled == True
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_strategy_reuse_after_cancellation(self):
|
||||
"""Verify strategy can be reused after cancellation."""
|
||||
call_count = 0
|
||||
|
||||
async def cancel_first_time():
|
||||
return call_count == 1
|
||||
|
||||
strategy = BFSDeepCrawlStrategy(
|
||||
max_depth=1,
|
||||
max_pages=5,
|
||||
should_cancel=cancel_first_time,
|
||||
)
|
||||
|
||||
mock_crawler = create_mock_crawler_with_links(num_links=2)
|
||||
mock_config = create_mock_config()
|
||||
|
||||
# First crawl - should be cancelled
|
||||
call_count = 1
|
||||
results1 = await strategy._arun_batch("https://example.com", mock_crawler, mock_config)
|
||||
assert strategy.cancelled == True
|
||||
|
||||
# Second crawl - should work normally (cancel_first_time returns False)
|
||||
call_count = 2
|
||||
results2 = await strategy._arun_batch("https://example.com", mock_crawler, mock_config)
|
||||
assert strategy.cancelled == False
|
||||
assert len(results2) > len(results1)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_callback_exception_continues_crawl(self):
|
||||
"""Verify callback exception doesn't crash crawl (fail-open)."""
|
||||
exception_count = 0
|
||||
|
||||
async def failing_callback():
|
||||
nonlocal exception_count
|
||||
exception_count += 1
|
||||
raise ConnectionError("Redis connection failed")
|
||||
|
||||
strategy = BFSDeepCrawlStrategy(
|
||||
max_depth=1,
|
||||
max_pages=3,
|
||||
should_cancel=failing_callback,
|
||||
)
|
||||
|
||||
mock_crawler = create_mock_crawler_with_links(num_links=2)
|
||||
mock_config = create_mock_config()
|
||||
|
||||
# Should not raise, should complete crawl
|
||||
results = await strategy._arun_batch("https://example.com", mock_crawler, mock_config)
|
||||
|
||||
assert exception_count > 0 # Callback was called
|
||||
assert len(results) > 0 # Crawl completed
|
||||
assert strategy.cancelled == False # Not cancelled due to exception
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_state_includes_cancelled_flag(self):
|
||||
"""Verify state notifications include cancelled flag."""
|
||||
states: List[Dict] = []
|
||||
cancel_at = 3
|
||||
|
||||
async def capture_state(state: Dict[str, Any]):
|
||||
states.append(state)
|
||||
|
||||
async def cancel_after_3():
|
||||
return len(states) >= cancel_at
|
||||
|
||||
strategy = BFSDeepCrawlStrategy(
|
||||
max_depth=5,
|
||||
max_pages=100,
|
||||
should_cancel=cancel_after_3,
|
||||
on_state_change=capture_state,
|
||||
)
|
||||
|
||||
mock_crawler = create_mock_crawler_with_links(num_links=5)
|
||||
mock_config = create_mock_config()
|
||||
|
||||
await strategy._arun_batch("https://example.com", mock_crawler, mock_config)
|
||||
|
||||
# Last state should have cancelled=True
|
||||
assert len(states) > 0
|
||||
assert states[-1].get("cancelled") == True
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_cancel_before_first_url(self):
|
||||
"""Verify cancel before first URL returns empty results."""
|
||||
async def always_cancel():
|
||||
return True
|
||||
|
||||
strategy = BFSDeepCrawlStrategy(
|
||||
max_depth=5,
|
||||
max_pages=100,
|
||||
should_cancel=always_cancel,
|
||||
)
|
||||
|
||||
mock_crawler = create_mock_crawler_with_links(num_links=5)
|
||||
mock_config = create_mock_config()
|
||||
|
||||
results = await strategy._arun_batch("https://example.com", mock_crawler, mock_config)
|
||||
|
||||
assert strategy.cancelled == True
|
||||
assert len(results) == 0
|
||||
|
||||
|
||||
class TestDFSCancellation:
|
||||
"""DFS strategy cancellation tests."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_cancel_via_callback(self):
|
||||
"""Verify DFS respects should_cancel callback."""
|
||||
pages_crawled = 0
|
||||
cancel_after = 3
|
||||
|
||||
async def check_cancel():
|
||||
return pages_crawled >= cancel_after
|
||||
|
||||
async def track_pages(state: Dict[str, Any]):
|
||||
nonlocal pages_crawled
|
||||
pages_crawled = state.get("pages_crawled", 0)
|
||||
|
||||
strategy = DFSDeepCrawlStrategy(
|
||||
max_depth=5,
|
||||
max_pages=100,
|
||||
should_cancel=check_cancel,
|
||||
on_state_change=track_pages,
|
||||
)
|
||||
|
||||
mock_crawler = create_mock_crawler_with_links(num_links=3)
|
||||
mock_config = create_mock_config()
|
||||
|
||||
await strategy._arun_batch("https://example.com", mock_crawler, mock_config)
|
||||
|
||||
assert strategy.cancelled == True
|
||||
assert strategy._pages_crawled >= cancel_after
|
||||
assert strategy._pages_crawled < 100
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_cancel_method_inherited(self):
|
||||
"""Verify DFS inherits cancel() from BFS."""
|
||||
strategy = DFSDeepCrawlStrategy(max_depth=2, max_pages=10)
|
||||
|
||||
assert hasattr(strategy, 'cancel')
|
||||
assert hasattr(strategy, 'cancelled')
|
||||
assert hasattr(strategy, '_check_cancellation')
|
||||
|
||||
strategy.cancel()
|
||||
assert strategy.cancelled == True
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_stream_mode_cancellation(self):
|
||||
"""Verify DFS stream mode respects cancellation."""
|
||||
results_count = 0
|
||||
cancel_after = 2
|
||||
|
||||
async def check_cancel():
|
||||
return results_count >= cancel_after
|
||||
|
||||
strategy = DFSDeepCrawlStrategy(
|
||||
max_depth=5,
|
||||
max_pages=100,
|
||||
should_cancel=check_cancel,
|
||||
)
|
||||
|
||||
mock_crawler = create_mock_crawler_with_links(num_links=3)
|
||||
mock_config = create_mock_config(stream=True)
|
||||
|
||||
async for result in strategy._arun_stream("https://example.com", mock_crawler, mock_config):
|
||||
results_count += 1
|
||||
|
||||
assert strategy.cancelled == True
|
||||
assert results_count >= cancel_after
|
||||
assert results_count < 100
|
||||
|
||||
|
||||
class TestBestFirstCancellation:
|
||||
"""Best-First strategy cancellation tests."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_cancel_via_callback(self):
|
||||
"""Verify Best-First respects should_cancel callback."""
|
||||
pages_crawled = 0
|
||||
cancel_after = 3
|
||||
|
||||
async def check_cancel():
|
||||
return pages_crawled >= cancel_after
|
||||
|
||||
async def track_pages(state: Dict[str, Any]):
|
||||
nonlocal pages_crawled
|
||||
pages_crawled = state.get("pages_crawled", 0)
|
||||
|
||||
strategy = BestFirstCrawlingStrategy(
|
||||
max_depth=5,
|
||||
max_pages=100,
|
||||
should_cancel=check_cancel,
|
||||
on_state_change=track_pages,
|
||||
)
|
||||
|
||||
mock_crawler = create_mock_crawler_with_links(num_links=3)
|
||||
mock_config = create_mock_config(stream=True)
|
||||
|
||||
async for _ in strategy._arun_stream("https://example.com", mock_crawler, mock_config):
|
||||
pass
|
||||
|
||||
assert strategy.cancelled == True
|
||||
assert strategy._pages_crawled >= cancel_after
|
||||
assert strategy._pages_crawled < 100
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_cancel_method_works(self):
|
||||
"""Verify Best-First cancel() method works."""
|
||||
strategy = BestFirstCrawlingStrategy(max_depth=2, max_pages=10)
|
||||
|
||||
assert strategy.cancelled == False
|
||||
strategy.cancel()
|
||||
assert strategy.cancelled == True
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_batch_mode_cancellation(self):
|
||||
"""Verify Best-First batch mode respects cancellation."""
|
||||
pages_crawled = 0
|
||||
cancel_after = 2
|
||||
|
||||
async def check_cancel():
|
||||
return pages_crawled >= cancel_after
|
||||
|
||||
async def track_pages(state: Dict[str, Any]):
|
||||
nonlocal pages_crawled
|
||||
pages_crawled = state.get("pages_crawled", 0)
|
||||
|
||||
strategy = BestFirstCrawlingStrategy(
|
||||
max_depth=5,
|
||||
max_pages=100,
|
||||
should_cancel=check_cancel,
|
||||
on_state_change=track_pages,
|
||||
)
|
||||
|
||||
mock_crawler = create_mock_crawler_with_links(num_links=3)
|
||||
mock_config = create_mock_config(stream=False)
|
||||
|
||||
results = await strategy._arun_batch("https://example.com", mock_crawler, mock_config)
|
||||
|
||||
assert strategy.cancelled == True
|
||||
assert len(results) >= cancel_after
|
||||
assert len(results) < 100
|
||||
|
||||
|
||||
class TestCrossStrategyCancellation:
|
||||
"""Tests that apply to all strategies."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("strategy_class", [
|
||||
BFSDeepCrawlStrategy,
|
||||
DFSDeepCrawlStrategy,
|
||||
BestFirstCrawlingStrategy,
|
||||
])
|
||||
async def test_no_cancel_callback_means_no_cancellation(self, strategy_class):
|
||||
"""Verify crawl completes normally without should_cancel."""
|
||||
strategy = strategy_class(max_depth=1, max_pages=5)
|
||||
|
||||
mock_crawler = create_mock_crawler_with_links(num_links=2)
|
||||
|
||||
if strategy_class == BestFirstCrawlingStrategy:
|
||||
mock_config = create_mock_config(stream=True)
|
||||
results = []
|
||||
async for r in strategy._arun_stream("https://example.com", mock_crawler, mock_config):
|
||||
results.append(r)
|
||||
else:
|
||||
mock_config = create_mock_config()
|
||||
results = await strategy._arun_batch("https://example.com", mock_crawler, mock_config)
|
||||
|
||||
assert strategy.cancelled == False
|
||||
assert len(results) > 0
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("strategy_class", [
|
||||
BFSDeepCrawlStrategy,
|
||||
DFSDeepCrawlStrategy,
|
||||
BestFirstCrawlingStrategy,
|
||||
])
|
||||
async def test_cancel_thread_safety(self, strategy_class):
|
||||
"""Verify cancel() is thread-safe (doesn't raise)."""
|
||||
strategy = strategy_class(max_depth=2, max_pages=10)
|
||||
|
||||
# Call cancel from multiple "threads" (simulated)
|
||||
for _ in range(10):
|
||||
strategy.cancel()
|
||||
|
||||
# Should be cancelled without errors
|
||||
assert strategy.cancelled == True
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("strategy_class", [
|
||||
BFSDeepCrawlStrategy,
|
||||
DFSDeepCrawlStrategy,
|
||||
BestFirstCrawlingStrategy,
|
||||
])
|
||||
async def test_should_cancel_param_accepted(self, strategy_class):
|
||||
"""Verify should_cancel parameter is accepted by constructor."""
|
||||
async def dummy_cancel():
|
||||
return False
|
||||
|
||||
# Should not raise
|
||||
strategy = strategy_class(
|
||||
max_depth=2,
|
||||
max_pages=10,
|
||||
should_cancel=dummy_cancel,
|
||||
)
|
||||
|
||||
assert strategy._should_cancel == dummy_cancel
|
||||
|
||||
|
||||
class TestCancellationEdgeCases:
|
||||
"""Edge case tests for cancellation."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_cancel_during_batch_processing(self):
|
||||
"""Verify cancellation during batch doesn't lose results."""
|
||||
results_count = 0
|
||||
|
||||
async def cancel_mid_batch():
|
||||
# Cancel after receiving first result
|
||||
return results_count >= 1
|
||||
|
||||
strategy = BFSDeepCrawlStrategy(
|
||||
max_depth=2,
|
||||
max_pages=100,
|
||||
should_cancel=cancel_mid_batch,
|
||||
)
|
||||
|
||||
async def track_results(state):
|
||||
nonlocal results_count
|
||||
results_count = state.get("pages_crawled", 0)
|
||||
|
||||
strategy._on_state_change = track_results
|
||||
|
||||
mock_crawler = create_mock_crawler_with_links(num_links=5)
|
||||
mock_config = create_mock_config()
|
||||
|
||||
results = await strategy._arun_batch("https://example.com", mock_crawler, mock_config)
|
||||
|
||||
# Should have at least the first batch of results
|
||||
assert len(results) >= 1
|
||||
assert strategy.cancelled == True
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_partial_results_on_cancel(self):
|
||||
"""Verify partial results are returned on cancellation."""
|
||||
cancel_after = 5
|
||||
|
||||
async def check_cancel():
|
||||
return strategy._pages_crawled >= cancel_after
|
||||
|
||||
strategy = BFSDeepCrawlStrategy(
|
||||
max_depth=10,
|
||||
max_pages=1000,
|
||||
should_cancel=check_cancel,
|
||||
)
|
||||
|
||||
mock_crawler = create_mock_crawler_with_links(num_links=10)
|
||||
mock_config = create_mock_config()
|
||||
|
||||
results = await strategy._arun_batch("https://example.com", mock_crawler, mock_config)
|
||||
|
||||
# Should have results up to cancellation point
|
||||
assert len(results) >= cancel_after
|
||||
assert strategy.cancelled == True
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_cancel_callback_called_once_per_level_bfs(self):
|
||||
"""Verify BFS checks cancellation once per level."""
|
||||
check_count = 0
|
||||
|
||||
async def count_checks():
|
||||
nonlocal check_count
|
||||
check_count += 1
|
||||
return False # Never cancel
|
||||
|
||||
strategy = BFSDeepCrawlStrategy(
|
||||
max_depth=2,
|
||||
max_pages=10,
|
||||
should_cancel=count_checks,
|
||||
)
|
||||
|
||||
mock_crawler = create_mock_crawler_with_links(num_links=2)
|
||||
mock_config = create_mock_config()
|
||||
|
||||
await strategy._arun_batch("https://example.com", mock_crawler, mock_config)
|
||||
|
||||
# Should have checked at least once per level
|
||||
assert check_count >= 1
|
||||
@@ -0,0 +1,340 @@
|
||||
"""
|
||||
Test Suite: Deep Crawl ContextVar Safety (Issue #1917)
|
||||
|
||||
Tests that DeepCrawlDecorator's ContextVar (deep_crawl_active) works correctly
|
||||
when the async generator is consumed in a different asyncio context, as happens
|
||||
with Starlette's StreamingResponse in the Docker API.
|
||||
|
||||
The bug: base_strategy.py used ContextVar.reset(token) in the generator's finally
|
||||
block, but reset() requires the same Context that created the token. When Starlette
|
||||
consumes the generator in a different Task, the Context changes -> ValueError.
|
||||
|
||||
The fix: use ContextVar.set(False) instead of reset(token), which works across
|
||||
context boundaries.
|
||||
"""
|
||||
|
||||
import pytest
|
||||
import asyncio
|
||||
from unittest.mock import MagicMock, AsyncMock
|
||||
|
||||
from crawl4ai.deep_crawling.base_strategy import DeepCrawlDecorator
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Helpers
|
||||
# ============================================================================
|
||||
|
||||
def create_mock_result(url="https://example.com"):
|
||||
result = MagicMock()
|
||||
result.url = url
|
||||
result.success = True
|
||||
result.metadata = {}
|
||||
result.links = {"internal": [], "external": []}
|
||||
return result
|
||||
|
||||
|
||||
def create_streaming_strategy(results):
|
||||
"""Create a mock deep crawl strategy that streams results."""
|
||||
strategy = MagicMock()
|
||||
|
||||
async def mock_arun(start_url, crawler, config):
|
||||
async def gen():
|
||||
for r in results:
|
||||
yield r
|
||||
return gen()
|
||||
|
||||
strategy.arun = mock_arun
|
||||
return strategy
|
||||
|
||||
|
||||
def create_batch_strategy(results):
|
||||
"""Create a mock deep crawl strategy that returns results as a list."""
|
||||
strategy = MagicMock()
|
||||
|
||||
async def mock_arun(start_url, crawler, config):
|
||||
return results
|
||||
|
||||
strategy.arun = mock_arun
|
||||
return strategy
|
||||
|
||||
|
||||
def create_config(stream=False, strategy=None):
|
||||
config = MagicMock()
|
||||
config.stream = stream
|
||||
config.deep_crawl_strategy = strategy
|
||||
return config
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Tests: ContextVar cross-context safety (the core #1917 bug)
|
||||
# ============================================================================
|
||||
|
||||
class TestContextVarCrossContext:
|
||||
"""Tests that deep_crawl_active ContextVar works across task boundaries."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_streaming_generator_consumed_in_different_task(self):
|
||||
"""
|
||||
Core reproduction of issue #1917:
|
||||
Create the generator in one task, consume it in another.
|
||||
Before the fix, this raised ValueError.
|
||||
"""
|
||||
mock_results = [create_mock_result(f"https://example.com/{i}") for i in range(3)]
|
||||
strategy = create_streaming_strategy(mock_results)
|
||||
config = create_config(stream=True, strategy=strategy)
|
||||
|
||||
crawler = MagicMock()
|
||||
decorator = DeepCrawlDecorator(crawler)
|
||||
|
||||
async def original_arun(url, config=None, **kwargs):
|
||||
return create_mock_result(url)
|
||||
|
||||
wrapped = decorator(original_arun)
|
||||
|
||||
# Call wrapped_arun in the current task — sets the token
|
||||
gen = await wrapped("https://example.com", config=config)
|
||||
|
||||
# Consume in a DIFFERENT task (simulates Starlette's StreamingResponse)
|
||||
collected = []
|
||||
|
||||
async def consume_in_new_task():
|
||||
async for result in gen:
|
||||
collected.append(result)
|
||||
|
||||
task = asyncio.create_task(consume_in_new_task())
|
||||
await task
|
||||
|
||||
assert len(collected) == 3
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_batch_mode_in_different_task(self):
|
||||
"""Non-streaming mode should also work across task boundaries."""
|
||||
mock_results = [create_mock_result("https://example.com")]
|
||||
strategy = create_batch_strategy(mock_results)
|
||||
config = create_config(stream=False, strategy=strategy)
|
||||
|
||||
crawler = MagicMock()
|
||||
decorator = DeepCrawlDecorator(crawler)
|
||||
|
||||
async def original_arun(url, config=None, **kwargs):
|
||||
return create_mock_result(url)
|
||||
|
||||
wrapped = decorator(original_arun)
|
||||
result = await wrapped("https://example.com", config=config)
|
||||
|
||||
assert result == mock_results
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Tests: ContextVar state management
|
||||
# ============================================================================
|
||||
|
||||
class TestContextVarState:
|
||||
"""Tests that deep_crawl_active is properly managed."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_flag_is_false_after_streaming_completes(self):
|
||||
"""deep_crawl_active should be False after the generator is exhausted."""
|
||||
strategy = create_streaming_strategy([create_mock_result()])
|
||||
config = create_config(stream=True, strategy=strategy)
|
||||
|
||||
crawler = MagicMock()
|
||||
decorator = DeepCrawlDecorator(crawler)
|
||||
|
||||
async def original_arun(url, config=None, **kwargs):
|
||||
return create_mock_result(url)
|
||||
|
||||
wrapped = decorator(original_arun)
|
||||
gen = await wrapped("https://example.com", config=config)
|
||||
|
||||
async for _ in gen:
|
||||
pass
|
||||
|
||||
assert decorator.deep_crawl_active.get() == False
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_flag_is_false_after_batch_completes(self):
|
||||
"""deep_crawl_active should be False after batch mode completes."""
|
||||
strategy = create_batch_strategy([create_mock_result()])
|
||||
config = create_config(stream=False, strategy=strategy)
|
||||
|
||||
crawler = MagicMock()
|
||||
decorator = DeepCrawlDecorator(crawler)
|
||||
|
||||
async def original_arun(url, config=None, **kwargs):
|
||||
return create_mock_result(url)
|
||||
|
||||
wrapped = decorator(original_arun)
|
||||
await wrapped("https://example.com", config=config)
|
||||
|
||||
assert decorator.deep_crawl_active.get() == False
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_flag_is_true_during_deep_crawl(self):
|
||||
"""deep_crawl_active should be True while the generator is being consumed."""
|
||||
flag_during_yield = None
|
||||
|
||||
async def capturing_arun(start_url, crawler, config):
|
||||
async def gen():
|
||||
nonlocal flag_during_yield
|
||||
flag_during_yield = DeepCrawlDecorator.deep_crawl_active.get()
|
||||
yield create_mock_result(start_url)
|
||||
return gen()
|
||||
|
||||
strategy = MagicMock()
|
||||
strategy.arun = capturing_arun
|
||||
config = create_config(stream=True, strategy=strategy)
|
||||
|
||||
crawler = MagicMock()
|
||||
decorator = DeepCrawlDecorator(crawler)
|
||||
|
||||
async def original_arun(url, config=None, **kwargs):
|
||||
return create_mock_result(url)
|
||||
|
||||
wrapped = decorator(original_arun)
|
||||
gen = await wrapped("https://example.com", config=config)
|
||||
|
||||
async for _ in gen:
|
||||
pass
|
||||
|
||||
assert flag_during_yield == True
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_flag_prevents_recursive_deep_crawl(self):
|
||||
"""When deep_crawl_active is True, nested calls should skip deep crawl."""
|
||||
crawler = MagicMock()
|
||||
decorator = DeepCrawlDecorator(crawler)
|
||||
|
||||
inner_call_hit = False
|
||||
|
||||
async def original_arun(url, config=None, **kwargs):
|
||||
nonlocal inner_call_hit
|
||||
inner_call_hit = True
|
||||
return create_mock_result(url)
|
||||
|
||||
wrapped = decorator(original_arun)
|
||||
|
||||
# Manually set the flag to simulate being inside a deep crawl
|
||||
decorator.deep_crawl_active.set(True)
|
||||
try:
|
||||
strategy = create_batch_strategy([create_mock_result()])
|
||||
config = create_config(stream=False, strategy=strategy)
|
||||
# Should call original_arun directly, NOT go through strategy
|
||||
result = await wrapped("https://example.com", config=config)
|
||||
assert inner_call_hit == True
|
||||
finally:
|
||||
decorator.deep_crawl_active.set(False)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_flag_reset_after_streaming_error(self):
|
||||
"""deep_crawl_active should be reset even if the generator raises."""
|
||||
strategy = MagicMock()
|
||||
|
||||
async def failing_arun(start_url, crawler, config):
|
||||
async def gen():
|
||||
yield create_mock_result("https://example.com")
|
||||
raise RuntimeError("simulated error")
|
||||
return gen()
|
||||
|
||||
strategy.arun = failing_arun
|
||||
config = create_config(stream=True, strategy=strategy)
|
||||
|
||||
crawler = MagicMock()
|
||||
decorator = DeepCrawlDecorator(crawler)
|
||||
|
||||
async def original_arun(url, config=None, **kwargs):
|
||||
return create_mock_result(url)
|
||||
|
||||
wrapped = decorator(original_arun)
|
||||
gen = await wrapped("https://example.com", config=config)
|
||||
|
||||
with pytest.raises(RuntimeError, match="simulated error"):
|
||||
async for _ in gen:
|
||||
pass
|
||||
|
||||
assert decorator.deep_crawl_active.get() == False
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_flag_reset_after_streaming_error_in_different_task(self):
|
||||
"""
|
||||
Combines #1917 fix with error handling: generator raises in a different task.
|
||||
Both the cross-context issue and error cleanup must work together.
|
||||
"""
|
||||
strategy = MagicMock()
|
||||
|
||||
async def failing_arun(start_url, crawler, config):
|
||||
async def gen():
|
||||
yield create_mock_result("https://example.com")
|
||||
raise RuntimeError("simulated error")
|
||||
return gen()
|
||||
|
||||
strategy.arun = failing_arun
|
||||
config = create_config(stream=True, strategy=strategy)
|
||||
|
||||
crawler = MagicMock()
|
||||
decorator = DeepCrawlDecorator(crawler)
|
||||
|
||||
async def original_arun(url, config=None, **kwargs):
|
||||
return create_mock_result(url)
|
||||
|
||||
wrapped = decorator(original_arun)
|
||||
gen = await wrapped("https://example.com", config=config)
|
||||
|
||||
error_caught = False
|
||||
|
||||
async def consume_in_new_task():
|
||||
nonlocal error_caught
|
||||
try:
|
||||
async for _ in gen:
|
||||
pass
|
||||
except RuntimeError:
|
||||
error_caught = True
|
||||
|
||||
task = asyncio.create_task(consume_in_new_task())
|
||||
await task
|
||||
|
||||
assert error_caught == True
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Tests: Concurrent requests
|
||||
# ============================================================================
|
||||
|
||||
class TestConcurrentRequests:
|
||||
"""Tests that multiple concurrent streaming deep crawls don't interfere."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_concurrent_streaming_in_separate_tasks(self):
|
||||
"""
|
||||
Multiple concurrent streaming requests consumed in separate tasks.
|
||||
This simulates multiple clients hitting /crawl/stream simultaneously.
|
||||
"""
|
||||
crawler = MagicMock()
|
||||
decorator = DeepCrawlDecorator(crawler)
|
||||
|
||||
async def original_arun(url, config=None, **kwargs):
|
||||
return create_mock_result(url)
|
||||
|
||||
wrapped = decorator(original_arun)
|
||||
results_per_request = {}
|
||||
|
||||
async def simulate_request(request_id):
|
||||
mock_results = [create_mock_result(f"https://example.com/{request_id}/{i}") for i in range(2)]
|
||||
strategy = create_streaming_strategy(mock_results)
|
||||
config = create_config(stream=True, strategy=strategy)
|
||||
|
||||
gen = await wrapped(f"https://example.com/{request_id}", config=config)
|
||||
results = []
|
||||
async for result in gen:
|
||||
results.append(result)
|
||||
await asyncio.sleep(0.01) # Interleave with other requests
|
||||
results_per_request[request_id] = results
|
||||
|
||||
tasks = [asyncio.create_task(simulate_request(i)) for i in range(3)]
|
||||
await asyncio.gather(*tasks)
|
||||
|
||||
assert len(results_per_request) == 3
|
||||
for request_id, results in results_per_request.items():
|
||||
assert len(results) == 2
|
||||
|
||||
assert decorator.deep_crawl_active.get() == False
|
||||
@@ -0,0 +1,839 @@
|
||||
"""
|
||||
Test Suite: Deep Crawl Resume/Crash Recovery Tests
|
||||
|
||||
Tests that verify:
|
||||
1. State export produces valid JSON-serializable data
|
||||
2. Resume from checkpoint continues without duplicates
|
||||
3. Simulated crash at various points recovers correctly
|
||||
4. State callback fires at expected intervals
|
||||
5. No damage to existing system behavior (regression tests)
|
||||
"""
|
||||
|
||||
import pytest
|
||||
import asyncio
|
||||
import json
|
||||
from typing import Dict, Any, List
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
from crawl4ai.deep_crawling import (
|
||||
BFSDeepCrawlStrategy,
|
||||
DFSDeepCrawlStrategy,
|
||||
BestFirstCrawlingStrategy,
|
||||
FilterChain,
|
||||
URLPatternFilter,
|
||||
DomainFilter,
|
||||
)
|
||||
from crawl4ai.deep_crawling.scorers import KeywordRelevanceScorer
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Helper Functions for Mock Crawler
|
||||
# ============================================================================
|
||||
|
||||
def create_mock_config(stream=False):
|
||||
"""Create a mock CrawlerRunConfig."""
|
||||
config = MagicMock()
|
||||
config.clone = MagicMock(return_value=config)
|
||||
config.stream = stream
|
||||
return config
|
||||
|
||||
|
||||
def create_mock_crawler_with_links(num_links: int = 3, include_keyword: bool = False):
|
||||
"""Create mock crawler that returns results with links."""
|
||||
call_count = 0
|
||||
|
||||
async def mock_arun_many(urls, config):
|
||||
nonlocal call_count
|
||||
results = []
|
||||
for url in urls:
|
||||
call_count += 1
|
||||
result = MagicMock()
|
||||
result.url = url
|
||||
result.success = True
|
||||
result.metadata = {}
|
||||
|
||||
# Generate child links
|
||||
links = []
|
||||
for i in range(num_links):
|
||||
link_url = f"{url}/child{call_count}_{i}"
|
||||
if include_keyword:
|
||||
link_url = f"{url}/important-child{call_count}_{i}"
|
||||
links.append({"href": link_url})
|
||||
|
||||
result.links = {"internal": links, "external": []}
|
||||
results.append(result)
|
||||
|
||||
# For streaming mode, return async generator
|
||||
if config.stream:
|
||||
async def gen():
|
||||
for r in results:
|
||||
yield r
|
||||
return gen()
|
||||
return results
|
||||
|
||||
crawler = MagicMock()
|
||||
crawler.arun_many = mock_arun_many
|
||||
return crawler
|
||||
|
||||
|
||||
def create_mock_crawler_tracking(crawl_order: List[str], return_no_links: bool = False):
|
||||
"""Create mock crawler that tracks crawl order."""
|
||||
|
||||
async def mock_arun_many(urls, config):
|
||||
results = []
|
||||
for url in urls:
|
||||
crawl_order.append(url)
|
||||
result = MagicMock()
|
||||
result.url = url
|
||||
result.success = True
|
||||
result.metadata = {}
|
||||
result.links = {"internal": [], "external": []} if return_no_links else {"internal": [{"href": f"{url}/child"}], "external": []}
|
||||
results.append(result)
|
||||
|
||||
# For streaming mode, return async generator
|
||||
if config.stream:
|
||||
async def gen():
|
||||
for r in results:
|
||||
yield r
|
||||
return gen()
|
||||
return results
|
||||
|
||||
crawler = MagicMock()
|
||||
crawler.arun_many = mock_arun_many
|
||||
return crawler
|
||||
|
||||
|
||||
def create_simple_mock_crawler():
|
||||
"""Basic mock crawler returning 1 result with 2 child links."""
|
||||
call_count = 0
|
||||
|
||||
async def mock_arun_many(urls, config):
|
||||
nonlocal call_count
|
||||
results = []
|
||||
for url in urls:
|
||||
call_count += 1
|
||||
result = MagicMock()
|
||||
result.url = url
|
||||
result.success = True
|
||||
result.metadata = {}
|
||||
result.links = {
|
||||
"internal": [
|
||||
{"href": f"{url}/child1"},
|
||||
{"href": f"{url}/child2"},
|
||||
],
|
||||
"external": []
|
||||
}
|
||||
results.append(result)
|
||||
|
||||
if config.stream:
|
||||
async def gen():
|
||||
for r in results:
|
||||
yield r
|
||||
return gen()
|
||||
return results
|
||||
|
||||
crawler = MagicMock()
|
||||
crawler.arun_many = mock_arun_many
|
||||
return crawler
|
||||
|
||||
|
||||
def create_mock_crawler_unlimited_links():
|
||||
"""Mock crawler that always returns links (for testing limits)."""
|
||||
async def mock_arun_many(urls, config):
|
||||
results = []
|
||||
for url in urls:
|
||||
result = MagicMock()
|
||||
result.url = url
|
||||
result.success = True
|
||||
result.metadata = {}
|
||||
result.links = {
|
||||
"internal": [{"href": f"{url}/link{i}"} for i in range(10)],
|
||||
"external": []
|
||||
}
|
||||
results.append(result)
|
||||
|
||||
if config.stream:
|
||||
async def gen():
|
||||
for r in results:
|
||||
yield r
|
||||
return gen()
|
||||
return results
|
||||
|
||||
crawler = MagicMock()
|
||||
crawler.arun_many = mock_arun_many
|
||||
return crawler
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# TEST SUITE 1: Crash Recovery Tests
|
||||
# ============================================================================
|
||||
|
||||
class TestBFSResume:
|
||||
"""BFS strategy resume tests."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_state_export_json_serializable(self):
|
||||
"""Verify exported state can be JSON serialized."""
|
||||
captured_states: List[Dict] = []
|
||||
|
||||
async def capture_state(state: Dict[str, Any]):
|
||||
# Verify JSON serializable
|
||||
json_str = json.dumps(state)
|
||||
parsed = json.loads(json_str)
|
||||
captured_states.append(parsed)
|
||||
|
||||
strategy = BFSDeepCrawlStrategy(
|
||||
max_depth=2,
|
||||
max_pages=10,
|
||||
on_state_change=capture_state,
|
||||
)
|
||||
|
||||
# Create mock crawler that returns predictable results
|
||||
mock_crawler = create_mock_crawler_with_links(num_links=3)
|
||||
mock_config = create_mock_config()
|
||||
|
||||
results = await strategy._arun_batch("https://example.com", mock_crawler, mock_config)
|
||||
|
||||
# Verify states were captured
|
||||
assert len(captured_states) > 0
|
||||
|
||||
# Verify state structure
|
||||
for state in captured_states:
|
||||
assert state["strategy_type"] == "bfs"
|
||||
assert "visited" in state
|
||||
assert "pending" in state
|
||||
assert "depths" in state
|
||||
assert "pages_crawled" in state
|
||||
assert isinstance(state["visited"], list)
|
||||
assert isinstance(state["pending"], list)
|
||||
assert isinstance(state["depths"], dict)
|
||||
assert isinstance(state["pages_crawled"], int)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_resume_continues_from_checkpoint(self):
|
||||
"""Verify resume starts from saved state, not beginning."""
|
||||
# Simulate state from previous crawl (visited 5 URLs, 3 pending)
|
||||
saved_state = {
|
||||
"strategy_type": "bfs",
|
||||
"visited": [
|
||||
"https://example.com",
|
||||
"https://example.com/page1",
|
||||
"https://example.com/page2",
|
||||
"https://example.com/page3",
|
||||
"https://example.com/page4",
|
||||
],
|
||||
"pending": [
|
||||
{"url": "https://example.com/page5", "parent_url": "https://example.com/page2"},
|
||||
{"url": "https://example.com/page6", "parent_url": "https://example.com/page3"},
|
||||
{"url": "https://example.com/page7", "parent_url": "https://example.com/page3"},
|
||||
],
|
||||
"depths": {
|
||||
"https://example.com": 0,
|
||||
"https://example.com/page1": 1,
|
||||
"https://example.com/page2": 1,
|
||||
"https://example.com/page3": 1,
|
||||
"https://example.com/page4": 1,
|
||||
"https://example.com/page5": 2,
|
||||
"https://example.com/page6": 2,
|
||||
"https://example.com/page7": 2,
|
||||
},
|
||||
"pages_crawled": 5,
|
||||
}
|
||||
|
||||
crawled_urls: List[str] = []
|
||||
|
||||
strategy = BFSDeepCrawlStrategy(
|
||||
max_depth=2,
|
||||
max_pages=20,
|
||||
resume_state=saved_state,
|
||||
)
|
||||
|
||||
# Verify internal state was restored
|
||||
assert strategy._resume_state == saved_state
|
||||
|
||||
mock_crawler = create_mock_crawler_tracking(crawled_urls, return_no_links=True)
|
||||
mock_config = create_mock_config()
|
||||
|
||||
await strategy._arun_batch("https://example.com", mock_crawler, mock_config)
|
||||
|
||||
# Should NOT re-crawl already visited URLs
|
||||
for visited_url in saved_state["visited"]:
|
||||
assert visited_url not in crawled_urls, f"Re-crawled already visited: {visited_url}"
|
||||
|
||||
# Should crawl pending URLs
|
||||
for pending in saved_state["pending"]:
|
||||
assert pending["url"] in crawled_urls, f"Did not crawl pending: {pending['url']}"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_simulated_crash_mid_crawl(self):
|
||||
"""Simulate crash at URL N, verify resume continues from pending URLs."""
|
||||
crash_after = 3
|
||||
states_before_crash: List[Dict] = []
|
||||
|
||||
async def capture_until_crash(state: Dict[str, Any]):
|
||||
states_before_crash.append(state)
|
||||
if state["pages_crawled"] >= crash_after:
|
||||
raise Exception("Simulated crash!")
|
||||
|
||||
strategy1 = BFSDeepCrawlStrategy(
|
||||
max_depth=2,
|
||||
max_pages=10,
|
||||
on_state_change=capture_until_crash,
|
||||
)
|
||||
|
||||
mock_crawler = create_mock_crawler_with_links(num_links=5)
|
||||
mock_config = create_mock_config()
|
||||
|
||||
# First crawl - crashes
|
||||
with pytest.raises(Exception, match="Simulated crash"):
|
||||
await strategy1._arun_batch("https://example.com", mock_crawler, mock_config)
|
||||
|
||||
# Get last state before crash
|
||||
last_state = states_before_crash[-1]
|
||||
assert last_state["pages_crawled"] >= crash_after
|
||||
|
||||
# Calculate which URLs were already crawled vs pending
|
||||
pending_urls = {item["url"] for item in last_state["pending"]}
|
||||
visited_urls = set(last_state["visited"])
|
||||
already_crawled_urls = visited_urls - pending_urls
|
||||
|
||||
# Resume from checkpoint
|
||||
crawled_in_resume: List[str] = []
|
||||
|
||||
strategy2 = BFSDeepCrawlStrategy(
|
||||
max_depth=2,
|
||||
max_pages=10,
|
||||
resume_state=last_state,
|
||||
)
|
||||
|
||||
mock_crawler2 = create_mock_crawler_tracking(crawled_in_resume, return_no_links=True)
|
||||
|
||||
await strategy2._arun_batch("https://example.com", mock_crawler2, mock_config)
|
||||
|
||||
# Verify already-crawled URLs are not re-crawled
|
||||
for crawled_url in already_crawled_urls:
|
||||
assert crawled_url not in crawled_in_resume, f"Re-crawled already visited: {crawled_url}"
|
||||
|
||||
# Verify pending URLs are crawled
|
||||
for pending_url in pending_urls:
|
||||
assert pending_url in crawled_in_resume, f"Did not crawl pending: {pending_url}"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_callback_fires_per_url(self):
|
||||
"""Verify callback fires after each URL for maximum granularity."""
|
||||
callback_count = 0
|
||||
pages_crawled_sequence: List[int] = []
|
||||
|
||||
async def count_callbacks(state: Dict[str, Any]):
|
||||
nonlocal callback_count
|
||||
callback_count += 1
|
||||
pages_crawled_sequence.append(state["pages_crawled"])
|
||||
|
||||
strategy = BFSDeepCrawlStrategy(
|
||||
max_depth=1,
|
||||
max_pages=5,
|
||||
on_state_change=count_callbacks,
|
||||
)
|
||||
|
||||
mock_crawler = create_mock_crawler_with_links(num_links=2)
|
||||
mock_config = create_mock_config()
|
||||
|
||||
await strategy._arun_batch("https://example.com", mock_crawler, mock_config)
|
||||
|
||||
# Callback should fire once per successful URL
|
||||
assert callback_count == strategy._pages_crawled, \
|
||||
f"Callback fired {callback_count} times, expected {strategy._pages_crawled} (per URL)"
|
||||
|
||||
# pages_crawled should increment by 1 each callback
|
||||
for i, count in enumerate(pages_crawled_sequence):
|
||||
assert count == i + 1, f"Expected pages_crawled={i+1} at callback {i}, got {count}"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_export_state_returns_last_captured(self):
|
||||
"""Verify export_state() returns last captured state."""
|
||||
last_state = None
|
||||
|
||||
async def capture(state):
|
||||
nonlocal last_state
|
||||
last_state = state
|
||||
|
||||
strategy = BFSDeepCrawlStrategy(max_depth=2, max_pages=5, on_state_change=capture)
|
||||
|
||||
mock_crawler = create_mock_crawler_with_links(num_links=2)
|
||||
mock_config = create_mock_config()
|
||||
|
||||
await strategy._arun_batch("https://example.com", mock_crawler, mock_config)
|
||||
|
||||
exported = strategy.export_state()
|
||||
assert exported == last_state
|
||||
|
||||
|
||||
class TestDFSResume:
|
||||
"""DFS strategy resume tests."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_state_export_includes_stack_and_dfs_seen(self):
|
||||
"""Verify DFS state includes stack structure and _dfs_seen."""
|
||||
captured_states: List[Dict] = []
|
||||
|
||||
async def capture_state(state: Dict[str, Any]):
|
||||
captured_states.append(state)
|
||||
|
||||
strategy = DFSDeepCrawlStrategy(
|
||||
max_depth=3,
|
||||
max_pages=10,
|
||||
on_state_change=capture_state,
|
||||
)
|
||||
|
||||
mock_crawler = create_mock_crawler_with_links(num_links=2)
|
||||
mock_config = create_mock_config()
|
||||
|
||||
await strategy._arun_batch("https://example.com", mock_crawler, mock_config)
|
||||
|
||||
assert len(captured_states) > 0
|
||||
|
||||
for state in captured_states:
|
||||
assert state["strategy_type"] == "dfs"
|
||||
assert "stack" in state
|
||||
assert "dfs_seen" in state
|
||||
# Stack items should have depth
|
||||
for item in state["stack"]:
|
||||
assert "url" in item
|
||||
assert "parent_url" in item
|
||||
assert "depth" in item
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_resume_restores_stack_order(self):
|
||||
"""Verify DFS stack order is preserved on resume."""
|
||||
saved_state = {
|
||||
"strategy_type": "dfs",
|
||||
"visited": ["https://example.com"],
|
||||
"stack": [
|
||||
{"url": "https://example.com/deep3", "parent_url": "https://example.com/deep2", "depth": 3},
|
||||
{"url": "https://example.com/deep2", "parent_url": "https://example.com/deep1", "depth": 2},
|
||||
{"url": "https://example.com/page1", "parent_url": "https://example.com", "depth": 1},
|
||||
],
|
||||
"depths": {"https://example.com": 0},
|
||||
"pages_crawled": 1,
|
||||
"dfs_seen": ["https://example.com", "https://example.com/deep3", "https://example.com/deep2", "https://example.com/page1"],
|
||||
}
|
||||
|
||||
crawl_order: List[str] = []
|
||||
|
||||
strategy = DFSDeepCrawlStrategy(
|
||||
max_depth=3,
|
||||
max_pages=10,
|
||||
resume_state=saved_state,
|
||||
)
|
||||
|
||||
mock_crawler = create_mock_crawler_tracking(crawl_order, return_no_links=True)
|
||||
mock_config = create_mock_config()
|
||||
|
||||
await strategy._arun_batch("https://example.com", mock_crawler, mock_config)
|
||||
|
||||
# DFS pops from end of stack, so order should be: page1, deep2, deep3
|
||||
assert crawl_order[0] == "https://example.com/page1"
|
||||
assert crawl_order[1] == "https://example.com/deep2"
|
||||
assert crawl_order[2] == "https://example.com/deep3"
|
||||
|
||||
|
||||
class TestBestFirstResume:
|
||||
"""Best-First strategy resume tests."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_state_export_includes_scored_queue(self):
|
||||
"""Verify Best-First state includes queue with scores."""
|
||||
captured_states: List[Dict] = []
|
||||
|
||||
async def capture_state(state: Dict[str, Any]):
|
||||
captured_states.append(state)
|
||||
|
||||
scorer = KeywordRelevanceScorer(keywords=["important"], weight=1.0)
|
||||
|
||||
strategy = BestFirstCrawlingStrategy(
|
||||
max_depth=2,
|
||||
max_pages=10,
|
||||
url_scorer=scorer,
|
||||
on_state_change=capture_state,
|
||||
)
|
||||
|
||||
mock_crawler = create_mock_crawler_with_links(num_links=3, include_keyword=True)
|
||||
mock_config = create_mock_config(stream=True)
|
||||
|
||||
async for _ in strategy._arun_stream("https://example.com", mock_crawler, mock_config):
|
||||
pass
|
||||
|
||||
assert len(captured_states) > 0
|
||||
|
||||
for state in captured_states:
|
||||
assert state["strategy_type"] == "best_first"
|
||||
assert "queue_items" in state
|
||||
for item in state["queue_items"]:
|
||||
assert "score" in item
|
||||
assert "depth" in item
|
||||
assert "url" in item
|
||||
assert "parent_url" in item
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_resume_maintains_priority_order(self):
|
||||
"""Verify priority queue order is maintained on resume."""
|
||||
saved_state = {
|
||||
"strategy_type": "best_first",
|
||||
"visited": ["https://example.com"],
|
||||
"queue_items": [
|
||||
{"score": -0.9, "depth": 1, "url": "https://example.com/high-priority", "parent_url": "https://example.com"},
|
||||
{"score": -0.5, "depth": 1, "url": "https://example.com/medium-priority", "parent_url": "https://example.com"},
|
||||
{"score": -0.1, "depth": 1, "url": "https://example.com/low-priority", "parent_url": "https://example.com"},
|
||||
],
|
||||
"depths": {"https://example.com": 0},
|
||||
"pages_crawled": 1,
|
||||
}
|
||||
|
||||
crawl_order: List[str] = []
|
||||
|
||||
strategy = BestFirstCrawlingStrategy(
|
||||
max_depth=2,
|
||||
max_pages=10,
|
||||
resume_state=saved_state,
|
||||
)
|
||||
|
||||
mock_crawler = create_mock_crawler_tracking(crawl_order, return_no_links=True)
|
||||
mock_config = create_mock_config(stream=True)
|
||||
|
||||
async for _ in strategy._arun_stream("https://example.com", mock_crawler, mock_config):
|
||||
pass
|
||||
|
||||
# Higher negative score = higher priority (min-heap)
|
||||
# So -0.9 should be crawled first
|
||||
assert crawl_order[0] == "https://example.com/high-priority"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_stream_results_follow_priority_order_with_out_of_order_batch(self):
|
||||
saved_state = {
|
||||
"strategy_type": "best_first",
|
||||
"visited": ["https://example.com"],
|
||||
"queue_items": [
|
||||
{"score": -0.9, "depth": 1, "url": "https://example.com/high", "parent_url": "https://example.com"},
|
||||
{"score": -0.1, "depth": 1, "url": "https://example.com/low", "parent_url": "https://example.com"},
|
||||
],
|
||||
"depths": {"https://example.com": 0},
|
||||
"pages_crawled": 1,
|
||||
}
|
||||
|
||||
async def mock_arun_many(urls, config):
|
||||
async def gen():
|
||||
for url in reversed(urls):
|
||||
result = MagicMock(url=url, success=True, metadata={})
|
||||
result.links = {"internal": [], "external": []}
|
||||
yield result
|
||||
return gen()
|
||||
|
||||
strategy = BestFirstCrawlingStrategy(max_depth=2, max_pages=10, resume_state=saved_state)
|
||||
mock_crawler = MagicMock(arun_many=mock_arun_many)
|
||||
results = [
|
||||
result.url
|
||||
async for result in strategy._arun_stream("https://example.com", mock_crawler, create_mock_config(stream=True))
|
||||
]
|
||||
|
||||
assert results[:2] == ["https://example.com/high", "https://example.com/low"]
|
||||
|
||||
|
||||
|
||||
class TestCrossStrategyResume:
|
||||
"""Tests that apply to all strategies."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("strategy_class,strategy_type", [
|
||||
(BFSDeepCrawlStrategy, "bfs"),
|
||||
(DFSDeepCrawlStrategy, "dfs"),
|
||||
(BestFirstCrawlingStrategy, "best_first"),
|
||||
])
|
||||
async def test_no_callback_means_no_overhead(self, strategy_class, strategy_type):
|
||||
"""Verify no state tracking when callback is None."""
|
||||
strategy = strategy_class(max_depth=2, max_pages=5)
|
||||
|
||||
# _queue_shadow should be None for Best-First when no callback
|
||||
if strategy_class == BestFirstCrawlingStrategy:
|
||||
assert strategy._queue_shadow is None
|
||||
|
||||
# _last_state should be None initially
|
||||
assert strategy._last_state is None
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("strategy_class", [
|
||||
BFSDeepCrawlStrategy,
|
||||
DFSDeepCrawlStrategy,
|
||||
BestFirstCrawlingStrategy,
|
||||
])
|
||||
async def test_export_state_returns_last_captured(self, strategy_class):
|
||||
"""Verify export_state() returns last captured state."""
|
||||
last_state = None
|
||||
|
||||
async def capture(state):
|
||||
nonlocal last_state
|
||||
last_state = state
|
||||
|
||||
strategy = strategy_class(max_depth=2, max_pages=5, on_state_change=capture)
|
||||
|
||||
mock_crawler = create_mock_crawler_with_links(num_links=2)
|
||||
|
||||
if strategy_class == BestFirstCrawlingStrategy:
|
||||
mock_config = create_mock_config(stream=True)
|
||||
async for _ in strategy._arun_stream("https://example.com", mock_crawler, mock_config):
|
||||
pass
|
||||
else:
|
||||
mock_config = create_mock_config()
|
||||
await strategy._arun_batch("https://example.com", mock_crawler, mock_config)
|
||||
|
||||
exported = strategy.export_state()
|
||||
assert exported == last_state
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# TEST SUITE 2: Regression Tests (No Damage to Current System)
|
||||
# ============================================================================
|
||||
|
||||
class TestBFSRegressions:
|
||||
"""Ensure BFS works identically when new params not used."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_default_params_unchanged(self):
|
||||
"""Constructor with only original params works."""
|
||||
strategy = BFSDeepCrawlStrategy(
|
||||
max_depth=2,
|
||||
include_external=False,
|
||||
max_pages=10,
|
||||
)
|
||||
|
||||
assert strategy.max_depth == 2
|
||||
assert strategy.include_external == False
|
||||
assert strategy.max_pages == 10
|
||||
assert strategy._resume_state is None
|
||||
assert strategy._on_state_change is None
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_filter_chain_still_works(self):
|
||||
"""FilterChain integration unchanged."""
|
||||
filter_chain = FilterChain([
|
||||
URLPatternFilter(patterns=["*/blog/*"]),
|
||||
DomainFilter(allowed_domains=["example.com"]),
|
||||
])
|
||||
|
||||
strategy = BFSDeepCrawlStrategy(
|
||||
max_depth=2,
|
||||
filter_chain=filter_chain,
|
||||
)
|
||||
|
||||
# Test filter still applies
|
||||
assert await strategy.can_process_url("https://example.com/blog/post1", 1) == True
|
||||
assert await strategy.can_process_url("https://other.com/blog/post1", 1) == False
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_url_scorer_still_works(self):
|
||||
"""URL scoring integration unchanged."""
|
||||
scorer = KeywordRelevanceScorer(keywords=["python", "tutorial"], weight=1.0)
|
||||
|
||||
strategy = BFSDeepCrawlStrategy(
|
||||
max_depth=2,
|
||||
url_scorer=scorer,
|
||||
score_threshold=0.5,
|
||||
)
|
||||
|
||||
assert strategy.url_scorer is not None
|
||||
assert strategy.score_threshold == 0.5
|
||||
|
||||
# Scorer should work
|
||||
score = scorer.score("https://example.com/python-tutorial")
|
||||
assert score > 0
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_batch_mode_returns_list(self):
|
||||
"""Batch mode still returns List[CrawlResult]."""
|
||||
strategy = BFSDeepCrawlStrategy(max_depth=1, max_pages=5)
|
||||
|
||||
mock_crawler = create_simple_mock_crawler()
|
||||
mock_config = create_mock_config(stream=False)
|
||||
|
||||
results = await strategy._arun_batch("https://example.com", mock_crawler, mock_config)
|
||||
|
||||
assert isinstance(results, list)
|
||||
assert len(results) > 0
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_max_pages_limit_respected(self):
|
||||
"""max_pages limit still enforced."""
|
||||
strategy = BFSDeepCrawlStrategy(max_depth=10, max_pages=3)
|
||||
|
||||
mock_crawler = create_mock_crawler_unlimited_links()
|
||||
mock_config = create_mock_config()
|
||||
|
||||
results = await strategy._arun_batch("https://example.com", mock_crawler, mock_config)
|
||||
|
||||
# Should stop at max_pages
|
||||
assert strategy._pages_crawled <= 3
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_max_depth_limit_respected(self):
|
||||
"""max_depth limit still enforced."""
|
||||
strategy = BFSDeepCrawlStrategy(max_depth=2, max_pages=100)
|
||||
|
||||
mock_crawler = create_mock_crawler_unlimited_links()
|
||||
mock_config = create_mock_config()
|
||||
|
||||
results = await strategy._arun_batch("https://example.com", mock_crawler, mock_config)
|
||||
|
||||
# All results should have depth <= max_depth
|
||||
for result in results:
|
||||
assert result.metadata.get("depth", 0) <= 2
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_metadata_depth_still_set(self):
|
||||
"""Result metadata still includes depth."""
|
||||
strategy = BFSDeepCrawlStrategy(max_depth=2, max_pages=5)
|
||||
|
||||
mock_crawler = create_simple_mock_crawler()
|
||||
mock_config = create_mock_config()
|
||||
|
||||
results = await strategy._arun_batch("https://example.com", mock_crawler, mock_config)
|
||||
|
||||
for result in results:
|
||||
assert "depth" in result.metadata
|
||||
assert isinstance(result.metadata["depth"], int)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_metadata_parent_url_still_set(self):
|
||||
"""Result metadata still includes parent_url."""
|
||||
strategy = BFSDeepCrawlStrategy(max_depth=2, max_pages=5)
|
||||
|
||||
mock_crawler = create_simple_mock_crawler()
|
||||
mock_config = create_mock_config()
|
||||
|
||||
results = await strategy._arun_batch("https://example.com", mock_crawler, mock_config)
|
||||
|
||||
# First result (start URL) should have parent_url = None
|
||||
assert results[0].metadata.get("parent_url") is None
|
||||
|
||||
# Child results should have parent_url set
|
||||
for result in results[1:]:
|
||||
assert "parent_url" in result.metadata
|
||||
|
||||
|
||||
class TestDFSRegressions:
|
||||
"""Ensure DFS works identically when new params not used."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_inherits_bfs_params(self):
|
||||
"""DFS still inherits all BFS parameters."""
|
||||
strategy = DFSDeepCrawlStrategy(
|
||||
max_depth=3,
|
||||
include_external=True,
|
||||
max_pages=20,
|
||||
score_threshold=0.5,
|
||||
)
|
||||
|
||||
assert strategy.max_depth == 3
|
||||
assert strategy.include_external == True
|
||||
assert strategy.max_pages == 20
|
||||
assert strategy.score_threshold == 0.5
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_dfs_seen_initialized(self):
|
||||
"""DFS _dfs_seen set still initialized."""
|
||||
strategy = DFSDeepCrawlStrategy(max_depth=2)
|
||||
|
||||
assert hasattr(strategy, '_dfs_seen')
|
||||
assert isinstance(strategy._dfs_seen, set)
|
||||
|
||||
|
||||
class TestBestFirstRegressions:
|
||||
"""Ensure Best-First works identically when new params not used."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_default_params_unchanged(self):
|
||||
"""Constructor with only original params works."""
|
||||
strategy = BestFirstCrawlingStrategy(
|
||||
max_depth=2,
|
||||
include_external=False,
|
||||
max_pages=10,
|
||||
)
|
||||
|
||||
assert strategy.max_depth == 2
|
||||
assert strategy.include_external == False
|
||||
assert strategy.max_pages == 10
|
||||
assert strategy._resume_state is None
|
||||
assert strategy._on_state_change is None
|
||||
assert strategy._queue_shadow is None # Not initialized without callback
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_scorer_integration(self):
|
||||
"""URL scorer still affects crawl priority."""
|
||||
scorer = KeywordRelevanceScorer(keywords=["important"], weight=1.0)
|
||||
|
||||
strategy = BestFirstCrawlingStrategy(
|
||||
max_depth=2,
|
||||
max_pages=10,
|
||||
url_scorer=scorer,
|
||||
)
|
||||
|
||||
assert strategy.url_scorer is scorer
|
||||
|
||||
|
||||
class TestMaxPagesBoundaryYielded:
|
||||
"""best_first must YIELD exactly max_pages results (boundary page kept).
|
||||
|
||||
Regression for #859: best_first used to break BEFORE yielding the page that
|
||||
hit the limit, so it returned max_pages-1 results while BFS already returned
|
||||
max_pages. The downstream scan count (discovered_urls) is derived from the
|
||||
yielded results, so the off-by-one surfaced as "11 requested -> 10 found".
|
||||
BFS is included as the working reference. (DFS batch mode has a separate,
|
||||
unrelated overshoot quirk and is intentionally not asserted here.)
|
||||
"""
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"strategy_class",
|
||||
[BFSDeepCrawlStrategy, BestFirstCrawlingStrategy],
|
||||
)
|
||||
@pytest.mark.asyncio
|
||||
async def test_yields_exactly_max_pages(self, strategy_class):
|
||||
max_pages = 11
|
||||
strategy = strategy_class(max_depth=10, max_pages=max_pages)
|
||||
|
||||
mock_crawler = create_mock_crawler_unlimited_links()
|
||||
# _arun_batch is the path the scan worker takes (crawler.arun with the
|
||||
# default stream=False). best_first's _arun_batch re-clones with
|
||||
# stream=True internally, so clone must honour the requested flag.
|
||||
mock_config = create_mock_config()
|
||||
mock_config.clone = lambda **kw: create_mock_config(stream=kw.get("stream", False))
|
||||
|
||||
results = await strategy._arun_batch("https://example.com", mock_crawler, mock_config)
|
||||
|
||||
assert len(results) == max_pages, (
|
||||
f"{strategy_class.__name__} yielded {len(results)} results, "
|
||||
f"expected exactly max_pages={max_pages}"
|
||||
)
|
||||
|
||||
|
||||
class TestAPICompatibility:
|
||||
"""Ensure API/serialization compatibility."""
|
||||
|
||||
def test_strategy_signature_backward_compatible(self):
|
||||
"""Old code calling with positional/keyword args still works."""
|
||||
# Positional args (old style)
|
||||
s1 = BFSDeepCrawlStrategy(2)
|
||||
assert s1.max_depth == 2
|
||||
|
||||
# Keyword args (old style)
|
||||
s2 = BFSDeepCrawlStrategy(max_depth=3, max_pages=10)
|
||||
assert s2.max_depth == 3
|
||||
|
||||
# Mixed (old style)
|
||||
s3 = BFSDeepCrawlStrategy(2, FilterChain(), None, False, float('-inf'), 100)
|
||||
assert s3.max_depth == 2
|
||||
assert s3.max_pages == 100
|
||||
|
||||
def test_no_required_new_params(self):
|
||||
"""New params are optional, not required."""
|
||||
# Should not raise
|
||||
BFSDeepCrawlStrategy(max_depth=2)
|
||||
DFSDeepCrawlStrategy(max_depth=2)
|
||||
BestFirstCrawlingStrategy(max_depth=2)
|
||||
@@ -0,0 +1,162 @@
|
||||
"""
|
||||
Integration Test: Deep Crawl Resume with Real URLs
|
||||
|
||||
Tests the crash recovery feature using books.toscrape.com - a site
|
||||
designed for scraping practice with a clear hierarchy:
|
||||
- Home page → Category pages → Book detail pages
|
||||
"""
|
||||
|
||||
import pytest
|
||||
import asyncio
|
||||
import json
|
||||
from typing import Dict, Any, List
|
||||
|
||||
from crawl4ai import AsyncWebCrawler, CrawlerRunConfig
|
||||
from crawl4ai.deep_crawling import BFSDeepCrawlStrategy
|
||||
|
||||
|
||||
class TestBFSResumeIntegration:
|
||||
"""Integration tests for BFS resume with real crawling."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_real_crawl_state_capture_and_resume(self):
|
||||
"""
|
||||
Test crash recovery with real URLs from books.toscrape.com.
|
||||
|
||||
Flow:
|
||||
1. Start crawl with state callback
|
||||
2. Stop after N pages (simulated crash)
|
||||
3. Resume from saved state
|
||||
4. Verify no duplicate crawls
|
||||
"""
|
||||
# Phase 1: Initial crawl that "crashes" after 3 pages
|
||||
crash_after = 3
|
||||
captured_states: List[Dict[str, Any]] = []
|
||||
crawled_urls_phase1: List[str] = []
|
||||
|
||||
async def capture_state_until_crash(state: Dict[str, Any]):
|
||||
captured_states.append(state)
|
||||
crawled_urls_phase1.clear()
|
||||
crawled_urls_phase1.extend(state["visited"])
|
||||
|
||||
if state["pages_crawled"] >= crash_after:
|
||||
raise Exception("Simulated crash!")
|
||||
|
||||
strategy1 = BFSDeepCrawlStrategy(
|
||||
max_depth=2,
|
||||
max_pages=10,
|
||||
on_state_change=capture_state_until_crash,
|
||||
)
|
||||
|
||||
config = CrawlerRunConfig(
|
||||
deep_crawl_strategy=strategy1,
|
||||
stream=False,
|
||||
verbose=False,
|
||||
)
|
||||
|
||||
async with AsyncWebCrawler(verbose=False) as crawler:
|
||||
# First crawl - will crash after 3 pages
|
||||
with pytest.raises(Exception, match="Simulated crash"):
|
||||
await crawler.arun("https://books.toscrape.com", config=config)
|
||||
|
||||
# Verify we captured state before crash
|
||||
assert len(captured_states) > 0, "No states captured before crash"
|
||||
last_state = captured_states[-1]
|
||||
|
||||
print(f"\n=== Phase 1: Crashed after {last_state['pages_crawled']} pages ===")
|
||||
print(f"Visited URLs: {len(last_state['visited'])}")
|
||||
print(f"Pending URLs: {len(last_state['pending'])}")
|
||||
|
||||
# Verify state structure
|
||||
assert last_state["strategy_type"] == "bfs"
|
||||
assert last_state["pages_crawled"] >= crash_after
|
||||
assert len(last_state["visited"]) > 0
|
||||
assert "pending" in last_state
|
||||
assert "depths" in last_state
|
||||
|
||||
# Verify state is JSON serializable (important for Redis/DB storage)
|
||||
json_str = json.dumps(last_state)
|
||||
restored_state = json.loads(json_str)
|
||||
assert restored_state == last_state, "State not JSON round-trip safe"
|
||||
|
||||
# Phase 2: Resume from checkpoint
|
||||
crawled_urls_phase2: List[str] = []
|
||||
|
||||
async def track_resumed_crawl(state: Dict[str, Any]):
|
||||
# Track what's being crawled in phase 2
|
||||
new_visited = set(state["visited"]) - set(last_state["visited"])
|
||||
for url in new_visited:
|
||||
if url not in crawled_urls_phase2:
|
||||
crawled_urls_phase2.append(url)
|
||||
|
||||
strategy2 = BFSDeepCrawlStrategy(
|
||||
max_depth=2,
|
||||
max_pages=10,
|
||||
resume_state=restored_state,
|
||||
on_state_change=track_resumed_crawl,
|
||||
)
|
||||
|
||||
config2 = CrawlerRunConfig(
|
||||
deep_crawl_strategy=strategy2,
|
||||
stream=False,
|
||||
verbose=False,
|
||||
)
|
||||
|
||||
async with AsyncWebCrawler(verbose=False) as crawler:
|
||||
results = await crawler.arun("https://books.toscrape.com", config=config2)
|
||||
|
||||
print(f"\n=== Phase 2: Resumed crawl ===")
|
||||
print(f"New URLs crawled: {len(crawled_urls_phase2)}")
|
||||
print(f"Final pages_crawled: {strategy2._pages_crawled}")
|
||||
|
||||
# Verify no duplicates - URLs from phase 1 should not be re-crawled
|
||||
already_crawled = set(last_state["visited"]) - {item["url"] for item in last_state["pending"]}
|
||||
duplicates = set(crawled_urls_phase2) & already_crawled
|
||||
|
||||
assert len(duplicates) == 0, f"Duplicate crawls detected: {duplicates}"
|
||||
|
||||
# Verify we made progress (crawled some of the pending URLs)
|
||||
pending_urls = {item["url"] for item in last_state["pending"]}
|
||||
crawled_pending = set(crawled_urls_phase2) & pending_urls
|
||||
|
||||
print(f"Pending URLs crawled in phase 2: {len(crawled_pending)}")
|
||||
|
||||
# Final state should show more pages crawled than before crash
|
||||
final_state = strategy2.export_state()
|
||||
if final_state:
|
||||
assert final_state["pages_crawled"] >= last_state["pages_crawled"], \
|
||||
"Resume did not make progress"
|
||||
|
||||
print("\n=== Integration test PASSED ===")
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_state_export_method(self):
|
||||
"""Test that export_state() returns valid state during crawl."""
|
||||
states_from_callback: List[Dict] = []
|
||||
|
||||
async def capture(state):
|
||||
states_from_callback.append(state)
|
||||
|
||||
strategy = BFSDeepCrawlStrategy(
|
||||
max_depth=1,
|
||||
max_pages=3,
|
||||
on_state_change=capture,
|
||||
)
|
||||
|
||||
config = CrawlerRunConfig(
|
||||
deep_crawl_strategy=strategy,
|
||||
stream=False,
|
||||
verbose=False,
|
||||
)
|
||||
|
||||
async with AsyncWebCrawler(verbose=False) as crawler:
|
||||
await crawler.arun("https://books.toscrape.com", config=config)
|
||||
|
||||
# export_state should return the last captured state
|
||||
exported = strategy.export_state()
|
||||
|
||||
assert exported is not None, "export_state() returned None"
|
||||
assert exported == states_from_callback[-1], "export_state() doesn't match last callback"
|
||||
|
||||
print(f"\n=== export_state() test PASSED ===")
|
||||
print(f"Final state: {exported['pages_crawled']} pages, {len(exported['visited'])} visited")
|
||||
Reference in New Issue
Block a user