Repository: apify/crawlee-python
File: tests/unit/playwright_crawler/test_playwright_crawler.py (Python)

Code before:
    assert 'User-Agent' in headers
    assert 'Sec-Ch-Ua' not in headers
    assert 'Sec-Ch-Ua-Mobile' not in headers
    assert 'Sec-Ch-Ua-Platform' not in headers

    assert 'headless' not in headers['User-Agent'].lower()

    assert headers['User-Agent'] == PW_FIREFOX_HEADLESS_DEFAULT_USER_AGENT


async def test_custom_headers() -> None:
    crawler = PlaywrightCrawler()
    headers = dict[str, str]()
    custom_headers = {'Power-Header': 'ring', 'Library': 'storm', 'My-Test-Header': 'fuzz'}

    @crawler.router.default_handler
    async def request_handler(context: PlaywrightCrawlingContext) -> None:
        response = await context.response.text()
        response_headers = dict(json.loads(response)).get('headers', {})

        for key, val in response_headers.items():
            headers[key] = val

    await crawler.run([Request.from_url('https://httpbin.org/get', headers=custom_headers)])
    for key, value in custom_headers.items():
        assert headers[key] == value


async def test_pre_navigation_hook() -> None:
    crawler = PlaywrightCrawler()
    mock_hook = mock.AsyncMock(return_value=None)

    crawler.pre_navigation_hook(mock_hook)

    @crawler.router.default_handler
    async def request_handler(_context: PlaywrightCrawlingContext) -> None:
        pass

    await crawler.run(['https://example.com', 'https://httpbin.org'])

    assert mock_hook.call_count == 2


Review comment (suggestion):
I believe we could compare the dicts as a whole like this:

```suggestion
    assert headers == custom_headers
```

also could we use a better names, e.g. `request_headers` and `response_headers`? :slightly_smiling_face: 

Code after:
    assert 'User-Agent' in headers
    assert 'Sec-Ch-Ua' not in headers
    assert 'Sec-Ch-Ua-Mobile' not in headers
    assert 'Sec-Ch-Ua-Platform' not in headers

    assert 'headless' not in headers['User-Agent'].lower()

    assert headers['User-Agent'] == PW_FIREFOX_HEADLESS_DEFAULT_USER_AGENT


async def test_custom_headers() -> None:
    crawler = PlaywrightCrawler()
    response_headers = dict[str, str]()
    request_headers = {'Power-Header': 'ring', 'Library': 'storm', 'My-Test-Header': 'fuzz'}

    @crawler.router.default_handler
    async def request_handler(context: PlaywrightCrawlingContext) -> None:
        response = await context.response.text()
        context_response_headers = dict(json.loads(response)).get('headers', {})

        for key, val in context_response_headers.items():
            response_headers[key] = val

    await crawler.run([Request.from_url('https://httpbin.org/get', headers=request_headers)])

    assert response_headers.get('Power-Header') == request_headers['Power-Header']
    assert response_headers.get('Library') == request_headers['Library']
    assert response_headers.get('My-Test-Header') == request_headers['My-Test-Header']


async def test_pre_navigation_hook() -> None:
    crawler = PlaywrightCrawler()
    mock_hook = mock.AsyncMock(return_value=None)

    crawler.pre_navigation_hook(mock_hook)

    @crawler.router.default_handler
    async def request_handler(_context: PlaywrightCrawlingContext) -> None:
        pass

    await crawler.run(['https://example.com', 'https://httpbin.org'])

    assert mock_hook.call_count == 2