BrowserUse14-源码-ScreenShot模块-整理

BrowserUse14-源码-ScreenShot模块-整理

---

ScreenShot模块

1-源码部分

【下载】基于图片的Base64进行图片的，【查看】指定本地的文件路径进行查看

"""
浏览器使用代理的截图存储服务。
"""

import base64
from pathlib import Path

import anyio

from browser_use_manual.observability import observe_debug


class ScreenshotService:
	"""简单的截图存储服务，将截图保存到磁盘"""

	def __init__(self, agent_directory: str | Path):
		"""初始化时传入代理目录路径"""
		self.agent_directory = Path(agent_directory) if isinstance(agent_directory, str) else agent_directory

		# 创建 screenshots 子目录
		self.screenshots_dir = self.agent_directory / 'screenshots'
		self.screenshots_dir.mkdir(parents=True, exist_ok=True)

	@observe_debug(ignore_input=True, ignore_output=True, name='store_screenshot')
	async def store_screenshot(self, screenshot_b64: str, step_number: int) -> str:
		"""将截图保存到磁盘，并返回完整路径字符串"""
		screenshot_filename = f'step_{step_number}.png'
		screenshot_path = self.screenshots_dir / screenshot_filename

		# 解码 base64 并保存到磁盘
		screenshot_data = base64.b64decode(screenshot_b64)

		async with await anyio.open_file(screenshot_path, 'wb') as f:
			await f.write(screenshot_data)

		return str(screenshot_path)

	@observe_debug(ignore_input=True, ignore_output=True, name='get_screenshot_from_disk')
	async def get_screenshot(self, screenshot_path: str) -> str | None:
		"""从磁盘路径加载截图，并以 base64 格式返回"""
		if not screenshot_path:
			return None

		path = Path(screenshot_path)
		if not path.exists():
			return None

		# 从磁盘读取并编码为 base64
		async with await anyio.open_file(path, 'rb') as f:
			screenshot_data = await f.read()

		return base64.b64encode(screenshot_data).decode('utf-8')

---

2-测试用例

测试【截图】模块的【本地图片的下载和读取】和【远程文件的下载和读取】

#!/usr/bin/env python3
"""
Screenshot 截图服务测试用例
测试读取本地图片文件并转换为base64编码，然后使用截图服务进行测试
"""

# 添加项目根目录到 Python 路径以支持导入
import sys
from pathlib import Path

sys.path.insert(0, str(Path(__file__).parent.parent.parent))

import base64
import asyncio
import logging
import httpx
from browser_use_manual.screenshots.service import ScreenshotService

# 配置日志
logging.basicConfig(level=logging.INFO, format='%(levelname)s - %(name)s - %(message)s')
logger = logging.getLogger(__name__)


async def test_get_screenshot_from_file(image_path: str):
    """测试从文件读取截图并转换为base64"""
    # 检查文件是否存在
    path = Path(image_path)
    if not path.exists():
        logger.error(f"测试图片文件不存在: {image_path}")
        return False

    # 读取图片文件并转换为base64
    with open(path, "rb") as image_file:
        image_data = image_file.read()
        image_b64 = base64.b64encode(image_data).decode('utf-8')

    logger.info(f"图片文件 '{image_path}' 已转换为base64，长度: {len(image_b64)} 字符")

    # 创建截图服务实例
    agent_dir = Path(__file__).parent.parent / "test_data"
    screenshot_service = ScreenshotService(agent_dir)
    
    # 使用截图服务存储截图
    stored_path = await screenshot_service.store_screenshot(image_b64, 1)
    logger.info(f"截图已存储到: {stored_path}")

    # 使用截图服务读取截图
    retrieved_b64 = await screenshot_service.get_screenshot(stored_path)
    
    if retrieved_b64 == image_b64:
        logger.info("✓ 成功: 从磁盘读取的截图与原图一致")
        return True
    else:
        logger.error("✗ 失败: 从磁盘读取的截图与原图不一致")
        return False


async def test_get_screenshot_from_url(image_url: str):
    """测试从URL下载图片并转换为base64"""
    try:
        # 下载远程图片
        async with httpx.AsyncClient() as client:
            response = await client.get(image_url)
            response.raise_for_status()
            image_data = response.content
            
        image_b64 = base64.b64encode(image_data).decode('utf-8')
        logger.info(f"远程图片 '{image_url}' 已下载并转换为base64，长度: {len(image_b64)} 字符")

        # 创建截图服务实例
        agent_dir = Path(__file__).parent.parent / "test_data"
        screenshot_service = ScreenshotService(agent_dir)
        
        # 使用截图服务存储截图
        stored_path = await screenshot_service.store_screenshot(image_b64, 2)
        logger.info(f"截图已存储到: {stored_path}")

        # 使用截图服务读取截图
        retrieved_b64 = await screenshot_service.get_screenshot(stored_path)
        
        if retrieved_b64 == image_b64:
            logger.info("✓ 成功: 从磁盘读取的截图与下载的图片一致")
            return True
        else:
            logger.error("✗ 失败: 从磁盘读取的截图与下载的图片不一致")
            return False
            
    except Exception as e:
        logger.error(f"下载或处理远程图片时出错: {e}")
        return False


async def run_all_tests():
    """运行所有测试"""
    # 1-测试本地文件（如果指定了有效路径）
    test_image_path = "/Users/rong/Documents/EnzoApplication/WorkSpace/Python/20251209_1_Python_playwright_manual/browser-use-manual-file/test_screen_image.jpg"
    path = Path(test_image_path)
    if path.exists():
        logger.info(f"开始测试本地图片文件: {test_image_path}")
        local_success = await test_get_screenshot_from_file(test_image_path)
    else:
        logger.warning(f"本地测试图片文件不存在: {test_image_path}")
        local_success = True  # 不影响远程图片测试
    
    # 2-测试远程图片
    remote_image_url = "https://i-blog.csdnimg.cn/direct/bee22ad9c3b74f07b8568843e48fe08e.png"
    logger.info(f"开始测试远程图片URL: {remote_image_url}")
    remote_success = await test_get_screenshot_from_url(remote_image_url)
    
    return local_success and remote_success


def main():
    # 指定要测试的图片文件路径
    logger.info("测试截图服务，包括本地文件和远程图片")
    
    # 运行所有测试
    success = asyncio.run(run_all_tests())
    
    if success:
        print("\n🎉 截图服务测试全部通过!")
    else:
        print("\n❌ 截图服务某些测试失败!")


if __name__ == "__main__":
    main()