diff --git a/main.py b/main.py index 847fc06..5071ae7 100644 --- a/main.py +++ b/main.py @@ -31,25 +31,23 @@ async def read_item(request: LinkRequest): logger.info(f"Request: {request}") start_time = int(time.time() * 1000) browser = await new_browser() - try: - page = await browser.get(request.url) - await page.bring_to_front() + page = await browser.get(request.url) + await page.bring_to_front() + timeout = request.maxTimeout + if timeout == 0: + timeout = None - challenged = await asyncio.wait_for( - bypass_cloudflare(page), timeout=request.maxTimeout - ) + challenged = await asyncio.wait_for(bypass_cloudflare(page), timeout=timeout) - logger.info(f"Got webpage: {request.url}") + logger.info(f"Got webpage: {request.url}") - response = await LinkResponse.create( - page=page, - start_timestamp=start_time, - challenged=challenged, - ) - except asyncio.TimeoutError: - logger.fatal("Couldn't complete the request") - finally: - browser.stop() + response = await LinkResponse.create( + page=page, + start_timestamp=start_time, + challenged=challenged, + ) + + browser.stop() return response diff --git a/tests/main_test.py b/tests/main_test.py index f6050f5..20e6160 100644 --- a/tests/main_test.py +++ b/tests/main_test.py @@ -20,6 +20,8 @@ test_websites = [ def test_bypass(website: str): response = client.post( "/v1", - json=LinkRequest(url=website, maxTimeout=60, cmd="request.get").model_dump(), + json=LinkRequest( + url=website, maxTimeout=60 * len(test_websites), cmd="request.get" + ).model_dump(), ) assert response.status_code == HTTPStatus.OK