diff --git a/src/crawlee/crawlers/_adaptive_playwright/_adaptive_playwright_crawler.py b/src/crawlee/crawlers/_adaptive_playwright/_adaptive_playwright_crawler.py index 84780fd525..a3139d3da1 100644 --- a/src/crawlee/crawlers/_adaptive_playwright/_adaptive_playwright_crawler.py +++ b/src/crawlee/crawlers/_adaptive_playwright/_adaptive_playwright_crawler.py @@ -315,7 +315,7 @@ async def get_input_state( ), logger=self._logger, ) - return SubCrawlerRun(result=result) + return SubCrawlerRun(result=result, run_context=context_linked_to_result) except Exception as e: return SubCrawlerRun(exception=e) @@ -371,7 +371,8 @@ async def _run_request_handler(self, context: BasicCrawlingContext) -> None: self.track_http_only_request_handler_runs() static_run = await self._crawl_one(rendering_type='static', context=context) - if static_run.result and self.result_checker(static_run.result): + if static_run.result and static_run.run_context and self.result_checker(static_run.result): + self._update_context_from_copy(context, static_run.run_context) self._context_result_map[context] = static_run.result return if static_run.exception: @@ -402,13 +403,10 @@ async def _run_request_handler(self, context: BasicCrawlingContext) -> None: if pw_run.exception is not None: raise pw_run.exception - if pw_run.result: - self._context_result_map[context] = pw_run.result - + if pw_run.result and pw_run.run_context: if should_detect_rendering_type: detection_result: RenderingType static_run = await self._crawl_one('static', context=context, state=old_state_copy) - if static_run.result and self.result_comparator(static_run.result, pw_run.result): detection_result = 'static' else: @@ -417,6 +415,9 @@ async def _run_request_handler(self, context: BasicCrawlingContext) -> None: context.log.debug(f'Detected rendering type {detection_result} for {context.request.url}') self.rendering_type_predictor.store_result(context.request, detection_result) + self._update_context_from_copy(context, pw_run.run_context) + self._context_result_map[context] = pw_run.result + def pre_navigation_hook( self, hook: Callable[[AdaptivePlaywrightPreNavCrawlingContext], Awaitable[None]] | None = None, @@ -451,8 +452,32 @@ def track_browser_request_handler_runs(self) -> None: def track_rendering_type_mispredictions(self) -> None: self.statistics.state.rendering_type_mispredictions += 1 + def _update_context_from_copy(self, context: BasicCrawlingContext, context_copy: BasicCrawlingContext) -> None: + """Update mutable fields of `context` from `context_copy`. + + Uses object.__setattr__ to bypass frozen dataclass restrictions, + allowing state synchronization after isolated crawler execution. + """ + updating_attributes = { + 'request': ('headers', 'user_data'), + 'session': ('_user_data', '_usage_count', '_error_score', '_cookies'), + } + + for attr, sub_attrs in updating_attributes.items(): + original_sub_obj = getattr(context, attr) + copy_sub_obj = getattr(context_copy, attr) + + # Check that both sub objects are not None + if original_sub_obj is None or copy_sub_obj is None: + continue + + for sub_attr in sub_attrs: + new_value = getattr(copy_sub_obj, sub_attr) + object.__setattr__(original_sub_obj, sub_attr, new_value) + @dataclass(frozen=True) class SubCrawlerRun: result: RequestHandlerRunResult | None = None exception: Exception | None = None + run_context: BasicCrawlingContext | None = None diff --git a/tests/unit/crawlers/_adaptive_playwright/test_adaptive_playwright_crawler.py b/tests/unit/crawlers/_adaptive_playwright/test_adaptive_playwright_crawler.py index 8e77e2dfeb..96b5dce3c6 100644 --- a/tests/unit/crawlers/_adaptive_playwright/test_adaptive_playwright_crawler.py +++ b/tests/unit/crawlers/_adaptive_playwright/test_adaptive_playwright_crawler.py @@ -29,9 +29,10 @@ from crawlee.crawlers._adaptive_playwright._adaptive_playwright_crawling_context import ( AdaptiveContextError, ) +from crawlee.sessions import SessionPool from crawlee.statistics import Statistics from crawlee.storage_clients import SqlStorageClient -from crawlee.storages import KeyValueStore +from crawlee.storages import KeyValueStore, RequestQueue if TYPE_CHECKING: from collections.abc import AsyncGenerator, Iterator @@ -730,6 +731,84 @@ async def request_handler(context: AdaptivePlaywrightCrawlingContext) -> None: mocked_h3_handler.assert_called_once_with(None) +@pytest.mark.parametrize( + 'test_input', + [ + pytest.param( + TestInput( + expected_pw_count=0, + expected_static_count=2, + rendering_types=cycle(['static']), + detection_probability_recommendation=cycle([0]), + ), + id='Static only', + ), + pytest.param( + TestInput( + expected_pw_count=2, + expected_static_count=0, + rendering_types=cycle(['client only']), + detection_probability_recommendation=cycle([0]), + ), + id='Client only', + ), + pytest.param( + TestInput( + expected_pw_count=2, + expected_static_count=2, + rendering_types=cycle(['static', 'client only']), + detection_probability_recommendation=cycle([1]), + ), + id='Enforced rendering type detection', + ), + ], +) +async def test_change_context_state_after_handling(test_input: TestInput, server_url: URL) -> None: + """Test that context state is saved after handling the request.""" + predictor = _SimpleRenderingTypePredictor( + rendering_types=test_input.rendering_types, + detection_probability_recommendation=test_input.detection_probability_recommendation, + ) + + request_queue = await RequestQueue.open(name='state-test') + used_session_id = None + + async with SessionPool() as session_pool: + crawler = AdaptivePlaywrightCrawler.with_beautifulsoup_static_parser( + rendering_type_predictor=predictor, + session_pool=session_pool, + request_manager=request_queue, + ) + + @crawler.router.default_handler + async def request_handler(context: AdaptivePlaywrightCrawlingContext) -> None: + nonlocal used_session_id + + if context.session is not None: + used_session_id = context.session.id + context.session.user_data['session_state'] = True + + if isinstance(context.request.user_data['request_state'], list): + context.request.user_data['request_state'].append('handler') + + request = Request.from_url(str(server_url), user_data={'request_state': ['initial']}) + + await crawler.run([request]) + + assert used_session_id is not None + + session = await session_pool.get_session_by_id(used_session_id) + check_request = await request_queue.get_request(request.unique_key) + + assert session is not None + assert check_request is not None + assert session.user_data.get('session_state') is True + # Check that request user data was updated in the handler and only onse. + assert check_request.user_data.get('request_state') == ['initial', 'handler'] + + await request_queue.drop() + + async def test_adaptive_playwright_crawler_with_sql_storage(test_urls: list[str], tmp_path: Path) -> None: """Tests that AdaptivePlaywrightCrawler can be initialized with SqlStorageClient.""" storage_dir = tmp_path / 'test_table.db'