screenshot delay and click selector

app/models/schemas.py

@@ -22,6 +22,20 @@ class VisitRequest(BaseModel):
     proxy: Optional[str] = None
     proxyAuth: Optional[dict[str, Any]] = None
     incognito: bool = False
+    screenshotSelector: Optional[str] = None
+    screenshotDelay: int = Field(
+        default=0,
+        ge=0,
+        le=60000,
+        description="Milliseconds to wait before taking the screenshot.",
+    )
+    clickSelector: Optional[str] = None
+    selectorTimeout: int = Field(
+        default=1000,
+        ge=100,
+        le=30000,
+        description="Milliseconds to wait for selectors used in screenshotSelector or clickSelector.",
+    )
 
 
 class ArasResponse(BaseModel):

app/use_cases/camoufox/wrapper.py

@@ -174,13 +174,15 @@ class CamoufoxWrapper:
                 html = await page.content()
 
             screenshot: Optional[str] = None
-            if options.get("screenshot"):
-                screenshot_bytes = await page.screenshot(
-                    full_page=True,
-                    type="jpeg",
-                    quality=config.SCREENSHOT_QUALITY,
-                )
-                screenshot = base64.b64encode(screenshot_bytes).decode("utf-8")
+            if (
+                options.get("screenshot")
+                or options.get("screenshotSelector")
+                or options.get("clickSelector")
+            ):
+                await self._click_element(page, options)
+                screenshot_bytes = await self._take_screenshot(page, options)
+                if screenshot_bytes:
+                    screenshot = base64.b64encode(screenshot_bytes).decode("utf-8")
 
             duration = time.time() - start_time
             log.logger.info(f"HTML Size: {len(html) / 1000:.2f} KB from '{url}'")
@@ -188,3 +190,56 @@ class CamoufoxWrapper:
 
             return {"html": html, "screenshot": screenshot}
 
+    def _get_screenshot_quality(self) -> int:
+        return config.SCREENSHOT_QUALITY
+
+    async def _take_screenshot(
+        self, page, options: dict[str, Any]
+    ) -> Optional[bytes]:
+        screenshot_options = {
+            "type": "jpeg",
+            "quality": self._get_screenshot_quality(),
+        }
+        selector = options.get("screenshotSelector")
+        selector_timeout = options.get("selectorTimeout", 1000)
+
+        if selector:
+            log.logger.info(f"Taking screenshot of selector: {selector}")
+            try:
+                await page.wait_for_selector(selector, timeout=selector_timeout)
+                element = await page.query_selector(selector)
+                if element:
+                    await self._apply_screenshot_delay(options)
+                    return await element.screenshot(**screenshot_options)
+                log.logger.warning(
+                    "Selector not found after waitForSelector, falling back to full page screenshot"
+                )
+            except Exception as err:
+                log.logger.warning(
+                    f'Selector "{selector}" not found within timeout, falling back to full page screenshot: {err}'
+                )
+
+        await self._apply_screenshot_delay(options)
+        return await page.screenshot(full_page=True, **screenshot_options)
+
+    async def _apply_screenshot_delay(self, options: dict[str, Any]) -> None:
+        delay = options.get("screenshotDelay", 0) or 0
+        if delay > 0:
+            log.logger.info(f"Waiting {delay} ms before taking screenshot")
+            await asyncio.sleep(delay / 1000)
+
+    async def _click_element(self, page, options: dict[str, Any]) -> None:
+        selector = options.get("clickSelector")
+        if not selector:
+            return
+
+        log.logger.info(f"Clicking element: {selector}")
+        try:
+            selector_timeout = options.get("selectorTimeout", 1000)
+            await page.wait_for_selector(selector, timeout=selector_timeout)
+            await page.click(selector)
+        except Exception as err:
+            log.logger.warning(
+                f'Click selector "{selector}" not found or not clickable: {err}'
+            )
+

tests/test_api_controller.py

@@ -25,6 +25,9 @@ def mocked_client(mock_response=None):
         mock_page.goto = AsyncMock()
         mock_page.wait_for_load_state = AsyncMock()
         mock_page.set_default_navigation_timeout = AsyncMock()
+        mock_page.wait_for_selector = AsyncMock()
+        mock_page.query_selector = AsyncMock(return_value=None)
+        mock_page.click = AsyncMock()
         mock_browser.new_page = AsyncMock(return_value=mock_page)
 
         MockClass.return_value.__aenter__ = AsyncMock(return_value=mock_browser)
@@ -78,6 +81,38 @@ class TestVisitEndpoint:
             assert body["status"] == "OK"
             assert "screenshot" in body["data"]
 
+    def test_visit_with_screenshot_selector(self):
+        with mocked_client() as client:
+            response = client.post(
+                "/api/v1/visit",
+                json={
+                    "url": "https://example.com",
+                    "screenshotSelector": "#product",
+                    "selectorTimeout": 2000,
+                },
+            )
+            assert response.status_code == 200
+
+    def test_visit_with_click_selector(self):
+        with mocked_client() as client:
+            response = client.post(
+                "/api/v1/visit",
+                json={
+                    "url": "https://example.com",
+                    "clickSelector": "button#load-more",
+                    "screenshotDelay": 500,
+                },
+            )
+            assert response.status_code == 200
+
+    def test_visit_with_invalid_selector_timeout(self):
+        with mocked_client() as client:
+            response = client.post(
+                "/api/v1/visit",
+                json={"url": "https://example.com", "selectorTimeout": 50},
+            )
+            assert response.status_code == 422
+
     def test_visit_with_invalid_waituntil(self):
         with mocked_client() as client:
             response = client.post(

tests/test_schemas.py

@@ -23,6 +23,10 @@ class TestVisitRequest:
         assert req.proxy is None
         assert req.proxyAuth is None
         assert req.incognito is False
+        assert req.screenshotSelector is None
+        assert req.screenshotDelay == 0
+        assert req.clickSelector is None
+        assert req.selectorTimeout == 1000
 
     def test_full_request(self):
         req = VisitRequest(
@@ -33,6 +37,10 @@ class TestVisitRequest:
             proxy="http://proxy:8080",
             proxyAuth={"username": "user", "password": "pass"},
             incognito=True,
+            screenshotSelector="#product",
+            screenshotDelay=500,
+            clickSelector="button#load-more",
+            selectorTimeout=2000,
         )
         assert req.screenshot is True
         assert req.waitUntil == "networkidle"
@@ -40,6 +48,30 @@ class TestVisitRequest:
         assert req.proxy == "http://proxy:8080"
         assert req.proxyAuth["username"] == "user"
         assert req.incognito is True
+        assert req.screenshotSelector == "#product"
+        assert req.screenshotDelay == 500
+        assert req.clickSelector == "button#load-more"
+        assert req.selectorTimeout == 2000
+
+    def test_screenshot_delay_boundary_values(self):
+        VisitRequest(url="https://example.com", screenshotDelay=0)
+        VisitRequest(url="https://example.com", screenshotDelay=60000)
+
+    def test_screenshot_delay_invalid_values(self):
+        with pytest.raises(ValidationError):
+            VisitRequest(url="https://example.com", screenshotDelay=-1)
+        with pytest.raises(ValidationError):
+            VisitRequest(url="https://example.com", screenshotDelay=60001)
+
+    def test_selector_timeout_boundary_values(self):
+        VisitRequest(url="https://example.com", selectorTimeout=100)
+        VisitRequest(url="https://example.com", selectorTimeout=30000)
+
+    def test_selector_timeout_invalid_values(self):
+        with pytest.raises(ValidationError):
+            VisitRequest(url="https://example.com", selectorTimeout=99)
+        with pytest.raises(ValidationError):
+            VisitRequest(url="https://example.com", selectorTimeout=30001)
 
     def test_wait_time_negative_rejected(self):
         with pytest.raises(ValidationError):

tests/test_wrapper.py

@@ -1,4 +1,6 @@
-from unittest.mock import MagicMock, patch
+import asyncio
+import pytest
+from unittest.mock import AsyncMock, MagicMock, patch
 from app.use_cases.camoufox.wrapper import (
     CamoufoxWrapper,
     _detect_headless_mode,
@@ -6,6 +8,10 @@ from app.use_cases.camoufox.wrapper import (
 )
 
 
+def _run_async(coro):
+    return asyncio.run(coro)
+
+
 class TestDetectHeadlessMode:
     @patch("sys.platform", "darwin")
     def test_macos_returns_true(self):
@@ -94,3 +100,134 @@ class TestCamoufoxWrapper:
         assert result["server"] == "http://host:8080"
         assert result["username"] == "user"
         assert "password" not in result
+
+
+class TestCamoufoxWrapperScreenshot:
+    def test_get_screenshot_quality(self):
+        wrapper = CamoufoxWrapper()
+        assert wrapper._get_screenshot_quality() == 60
+
+    def test_apply_screenshot_delay_zero(self):
+        wrapper = CamoufoxWrapper()
+        with patch("asyncio.sleep", new_callable=AsyncMock) as mock_sleep:
+            _run_async(wrapper._apply_screenshot_delay({"screenshotDelay": 0}))
+            mock_sleep.assert_not_awaited()
+
+    def test_apply_screenshot_delay_positive(self):
+        wrapper = CamoufoxWrapper()
+        with patch("asyncio.sleep", new_callable=AsyncMock) as mock_sleep:
+            _run_async(wrapper._apply_screenshot_delay({"screenshotDelay": 500}))
+            mock_sleep.assert_awaited_once_with(0.5)
+
+    def test_click_element_no_selector(self):
+        wrapper = CamoufoxWrapper()
+        page = MagicMock()
+        _run_async(wrapper._click_element(page, {}))
+        page.wait_for_selector.assert_not_called()
+        page.click.assert_not_called()
+
+    def test_click_element_success(self):
+        wrapper = CamoufoxWrapper()
+        page = AsyncMock()
+        _run_async(wrapper._click_element(page, {"clickSelector": "button#btn", "selectorTimeout": 2000}))
+        page.wait_for_selector.assert_awaited_once_with("button#btn", timeout=2000)
+        page.click.assert_awaited_once_with("button#btn")
+
+    def test_click_element_failure_logs_warning(self):
+        wrapper = CamoufoxWrapper()
+        page = AsyncMock()
+        page.wait_for_selector.side_effect = Exception("timeout")
+        with patch("app.use_cases.camoufox.wrapper.log.logger.warning") as mock_warn:
+            _run_async(wrapper._click_element(page, {"clickSelector": "button#btn"}))
+            mock_warn.assert_called_once()
+
+    def test_take_screenshot_full_page(self):
+        wrapper = CamoufoxWrapper()
+        page = AsyncMock()
+        page.screenshot.return_value = b"fullpage"
+        result = _run_async(wrapper._take_screenshot(page, {}))
+        page.screenshot.assert_awaited_once_with(full_page=True, type="jpeg", quality=60)
+        assert result == b"fullpage"
+
+    def test_take_screenshot_by_selector(self):
+        wrapper = CamoufoxWrapper()
+        page = AsyncMock()
+        element = AsyncMock()
+        element.screenshot.return_value = b"element"
+        page.query_selector.return_value = element
+        result = _run_async(
+            wrapper._take_screenshot(page, {"screenshotSelector": "#product", "selectorTimeout": 2000})
+        )
+        page.wait_for_selector.assert_awaited_once_with("#product", timeout=2000)
+        page.query_selector.assert_awaited_once_with("#product")
+        element.screenshot.assert_awaited_once_with(type="jpeg", quality=60)
+        assert result == b"element"
+
+    def test_take_screenshot_selector_fallback_to_full_page(self):
+        wrapper = CamoufoxWrapper()
+        page = AsyncMock()
+        page.query_selector.return_value = None
+        page.screenshot.return_value = b"fallback"
+        with patch("app.use_cases.camoufox.wrapper.log.logger.warning") as mock_warn:
+            result = _run_async(
+                wrapper._take_screenshot(page, {"screenshotSelector": "#missing"})
+            )
+            mock_warn.assert_called_once()
+            page.screenshot.assert_awaited_once_with(full_page=True, type="jpeg", quality=60)
+            assert result == b"fallback"
+
+    def test_take_screenshot_with_delay(self):
+        wrapper = CamoufoxWrapper()
+        page = AsyncMock()
+        page.screenshot.return_value = b"delayed"
+        with patch("asyncio.sleep", new_callable=AsyncMock) as mock_sleep:
+            result = _run_async(wrapper._take_screenshot(page, {"screenshotDelay": 300}))
+            mock_sleep.assert_awaited_once_with(0.3)
+            assert result == b"delayed"
+
+
+class TestCamoufoxWrapperVisitScreenshotFlow:
+    @patch("app.use_cases.camoufox.wrapper._resolve_headless_mode", return_value=True)
+    @patch("app.use_cases.camoufox.wrapper.AsyncCamoufox")
+    @patch("app.use_cases.camoufox.wrapper.base64.b64encode", return_value=b"base64data")
+    def test_visit_triggers_screenshot_with_screenshot_selector(
+        self, mock_b64encode, mock_browser_cls, mock_headless
+    ):
+        wrapper = CamoufoxWrapper()
+        page = AsyncMock()
+        element = AsyncMock()
+        element.screenshot.return_value = b"element-shot"
+        page.query_selector.return_value = element
+        page.content.return_value = "<html></html>"
+        browser = AsyncMock()
+        browser.new_page.return_value = page
+        mock_browser_cls.return_value.__aenter__ = AsyncMock(return_value=browser)
+        mock_browser_cls.return_value.__aexit__ = AsyncMock(return_value=False)
+
+        result = _run_async(
+            wrapper.visit("https://example.com", {"screenshotSelector": "#product"})
+        )
+
+        page.wait_for_selector.assert_awaited_once_with("#product", timeout=1000)
+        element.screenshot.assert_awaited_once()
+        assert result["screenshot"] == "base64data"
+
+    @patch("app.use_cases.camoufox.wrapper._resolve_headless_mode", return_value=True)
+    @patch("app.use_cases.camoufox.wrapper.AsyncCamoufox")
+    def test_visit_triggers_click_then_screenshot(
+        self, mock_browser_cls, mock_headless
+    ):
+        wrapper = CamoufoxWrapper()
+        page = AsyncMock()
+        page.content.return_value = "<html></html>"
+        page.screenshot.return_value = b"shot"
+        browser = AsyncMock()
+        browser.new_page.return_value = page
+        mock_browser_cls.return_value.__aenter__ = AsyncMock(return_value=browser)
+        mock_browser_cls.return_value.__aexit__ = AsyncMock(return_value=False)
+
+        _run_async(wrapper.visit("https://example.com", {"clickSelector": "button"}))
+
+        page.wait_for_selector.assert_any_await("button", timeout=1000)
+        page.click.assert_awaited_once_with("button")
+        page.screenshot.assert_awaited_once()