|
19 | 19 |
|
20 | 20 | import pytest |
21 | 21 |
|
22 | | -from crawlee import ConcurrencySettings, Glob, service_locator |
| 22 | +from crawlee import ConcurrencySettings, EnqueueStrategy, Glob, service_locator |
23 | 23 | from crawlee._request import Request, RequestState |
24 | 24 | from crawlee._types import BasicCrawlingContext, EnqueueLinksKwargs, HttpMethod |
25 | 25 | from crawlee._utils.robots import RobotsTxtFile |
@@ -340,56 +340,41 @@ async def failed_request_handler(context: BasicCrawlingContext, error: Exception |
340 | 340 | await crawler.run(['https://a.placeholder.com', 'https://b.placeholder.com', 'https://c.placeholder.com']) |
341 | 341 |
|
342 | 342 |
|
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: |
345 | 354 | 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.""" |
361 | 355 | errors: list[Exception] = [] |
362 | 356 |
|
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) |
364 | 359 |
|
365 | 360 | @crawler.router.default_handler |
366 | 361 | async def handler(context: BasicCrawlingContext) -> None: |
367 | 362 | try: |
368 | | - await context.send_request('http://attacker.evil/payload') |
| 363 | + response = await context.send_request(target_url) |
369 | 364 | except ValueError as exc: |
370 | 365 | errors.append(exc) |
| 366 | + else: |
| 367 | + bodies.append(await response.read()) |
371 | 368 |
|
372 | 369 | await crawler.run([str(server_url / 'a/page')]) |
373 | 370 |
|
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]) |
393 | 378 |
|
394 | 379 |
|
395 | 380 | @pytest.mark.parametrize( |
|
0 commit comments