chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -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
|
||||
Reference in New Issue
Block a user