from io import BytesIO

import allure
from PIL import Image, ImageDraw
from playwright.sync_api import sync_playwright

def highlight_screenshot(page, element, tips='高亮截圖'):
    # 滾動(dòng)頁(yè)面將元素移動(dòng)到視野中
    element.scroll_into_view_if_needed()

    # 獲取元素的位置信息
    box = element.bounding_box()
    x, y, width, height = box['x'], box['y'], box['width'], box['height']

    # 獲取頁(yè)面截圖，并讀取為PIL.Image對(duì)象
    png = page.screenshot(type='png')
    img = Image.open(BytesIO(png))

    # 利用元素的位置信息在頁(yè)面截圖中標(biāo)記出元素
    draw = ImageDraw.Draw(img)
    draw.rectangle((x, y, x + width, y + height), outline='red', width=3)

    # 將標(biāo)記后的截圖保存至allure報(bào)告中
    with BytesIO() as output:
        img.save(output, format='png')
        allure.attach(output.getvalue(), tips, allure.attachment_type.PNG)

def test_screenshot():
    with sync_playwright() as pw:
        # 打開(kāi)百度進(jìn)行搜索
        browser = pw.chromium.launch()
        page = browser.new_page()
        page.goto('https://www.baidu.com')  # 跳轉(zhuǎn)至百度頁(yè)面
        page.fill('#kw', 'playwright 選擇器')  # 搜索框中輸入內(nèi)容
        page.click('#su')  # 點(diǎn)擊搜索按鈕

        # 等待搜索結(jié)果加載
        page.wait_for_load_state('networkidle')

        # 判斷當(dāng)前頁(yè)面顯示10條搜索結(jié)果
        results = page.query_selector_all('h3')
        assert len(results) == 10

        # 將當(dāng)前頁(yè)面的搜索結(jié)果分別進(jìn)行截圖
        for i, result in enumerate(results):
            highlight_screenshot(page, result, f'搜索結(jié)果{i}')