Skip to content

Commit 075a54a

Browse files
committed
refactor: Simplify send_request enqueue strategy validation
Hoist origin URL parsing out of the per-call closure so it is computed once per request, not on every send_request call. Drop a redundant type annotation and collapse the three nearly-identical tests into one parametrized case.
1 parent a7ec446 commit 075a54a

2 files changed

Lines changed: 38 additions & 51 deletions

File tree

src/crawlee/crawlers/_basic/_basic_crawler.py

Lines changed: 14 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -444,7 +444,7 @@ def __init__(
444444
self._max_session_rotations = max_session_rotations
445445
self._max_crawl_depth = max_crawl_depth
446446
self._respect_robots_txt_file = respect_robots_txt_file
447-
self._send_request_enqueue_strategy: EnqueueStrategy = send_request_enqueue_strategy
447+
self._send_request_enqueue_strategy = send_request_enqueue_strategy
448448

449449
# Timeouts
450450
self._request_handler_timeout = request_handler_timeout
@@ -1290,24 +1290,26 @@ def _prepare_send_request_function(
12901290
proxy_info: ProxyInfo | None,
12911291
request: Request,
12921292
) -> SendRequestFunction:
1293+
strategy = self._send_request_enqueue_strategy
1294+
origin_url = request.loaded_url or request.url
1295+
origin_parsed = urlparse(origin_url) if strategy != 'all' else None
1296+
12931297
async def send_request(
12941298
url: str,
12951299
*,
12961300
method: HttpMethod = 'GET',
12971301
payload: HttpPayload | None = None,
12981302
headers: HttpHeaders | dict[str, str] | None = None,
12991303
) -> HttpResponse:
1300-
if self._send_request_enqueue_strategy != 'all':
1301-
origin_url = request.loaded_url or request.url
1302-
if not self._check_enqueue_strategy(
1303-
self._send_request_enqueue_strategy,
1304-
target_url=urlparse(url),
1305-
origin_url=urlparse(origin_url),
1306-
):
1307-
raise ValueError(
1308-
f'send_request() refusing to fetch {url!r}: does not match enqueue strategy '
1309-
f'{self._send_request_enqueue_strategy!r} relative to {origin_url!r}.'
1310-
)
1304+
if origin_parsed is not None and not self._check_enqueue_strategy(
1305+
strategy,
1306+
target_url=urlparse(url),
1307+
origin_url=origin_parsed,
1308+
):
1309+
raise ValueError(
1310+
f'send_request() refusing to fetch {url!r}: does not match enqueue strategy '
1311+
f'{strategy!r} relative to {origin_url!r}.'
1312+
)
13111313
return await self._http_client.send_request(
13121314
url=url,
13131315
method=method,

tests/unit/crawlers/_basic/test_basic_crawler.py

Lines changed: 24 additions & 39 deletions
Original file line numberDiff line numberDiff line change
@@ -19,7 +19,7 @@
1919

2020
import pytest
2121

22-
from crawlee import ConcurrencySettings, Glob, service_locator
22+
from crawlee import ConcurrencySettings, EnqueueStrategy, Glob, service_locator
2323
from crawlee._request import Request, RequestState
2424
from crawlee._types import BasicCrawlingContext, EnqueueLinksKwargs, HttpMethod
2525
from crawlee._utils.robots import RobotsTxtFile
@@ -340,56 +340,41 @@ async def failed_request_handler(context: BasicCrawlingContext, error: Exception
340340
await crawler.run(['https://a.placeholder.com', 'https://b.placeholder.com', 'https://c.placeholder.com'])
341341

342342

343-
async def test_send_request_strategy_default_allows_cross_host(server_url: URL) -> None:
344-
"""The default `send_request_enqueue_strategy='all'` permits cross-host `send_request` calls."""
343+
@pytest.mark.parametrize(
344+
('strategy', 'target_path', 'should_succeed'),
345+
[
346+
pytest.param('all', 'get', True, id='default-all-allows-same-host'),
347+
pytest.param('same-hostname', 'get', True, id='same-hostname-allows-same-host'),
348+
pytest.param('same-hostname', 'http://attacker.evil/payload', False, id='same-hostname-rejects-cross-host'),
349+
],
350+
)
351+
async def test_send_request_enqueue_strategy(
352+
server_url: URL, strategy: EnqueueStrategy, target_path: str, *, should_succeed: bool
353+
) -> None:
345354
bodies: list[bytes] = []
346-
347-
crawler = BasicCrawler(max_request_retries=1)
348-
349-
@crawler.router.default_handler
350-
async def handler(context: BasicCrawlingContext) -> None:
351-
response = await context.send_request(str(server_url / 'get'))
352-
bodies.append(await response.read())
353-
354-
await crawler.run([str(server_url / 'a/page')])
355-
356-
assert bodies, 'expected the handler to receive at least one response'
357-
358-
359-
async def test_send_request_strategy_same_hostname_rejects_cross_host(server_url: URL) -> None:
360-
"""`send_request_enqueue_strategy='same-hostname'` raises when target URL is on a different host."""
361355
errors: list[Exception] = []
362356

363-
crawler = BasicCrawler(max_request_retries=1, send_request_enqueue_strategy='same-hostname')
357+
crawler = BasicCrawler(max_request_retries=1, send_request_enqueue_strategy=strategy)
358+
target_url = target_path if target_path.startswith('http') else str(server_url / target_path)
364359

365360
@crawler.router.default_handler
366361
async def handler(context: BasicCrawlingContext) -> None:
367362
try:
368-
await context.send_request('http://attacker.evil/payload')
363+
response = await context.send_request(target_url)
369364
except ValueError as exc:
370365
errors.append(exc)
366+
else:
367+
bodies.append(await response.read())
371368

372369
await crawler.run([str(server_url / 'a/page')])
373370

374-
assert errors, 'expected send_request to refuse the cross-host URL'
375-
assert 'same-hostname' in str(errors[0])
376-
assert 'attacker.evil/payload' in str(errors[0])
377-
378-
379-
async def test_send_request_strategy_same_hostname_allows_same_host(server_url: URL) -> None:
380-
"""`send_request_enqueue_strategy='same-hostname'` lets same-host targets through unchanged."""
381-
bodies: list[bytes] = []
382-
383-
crawler = BasicCrawler(max_request_retries=1, send_request_enqueue_strategy='same-hostname')
384-
385-
@crawler.router.default_handler
386-
async def handler(context: BasicCrawlingContext) -> None:
387-
response = await context.send_request(str(server_url / 'get'))
388-
bodies.append(await response.read())
389-
390-
await crawler.run([str(server_url / 'a/page')])
391-
392-
assert bodies, 'expected the handler to receive a response from the same host'
371+
if should_succeed:
372+
assert bodies, 'expected the handler to receive a response'
373+
assert not errors
374+
else:
375+
assert errors, 'expected send_request to refuse the target URL'
376+
assert strategy in str(errors[0])
377+
assert target_url in str(errors[0])
393378

394379

395380
@pytest.mark.parametrize(

0 commit comments

Comments
 (0)