feat: add driver.GetScreenShot

2026-05-30 12:59:39 +08:00 · 2024-08-30 23:24:23 +08:00
parent abe2536942
commit 8567b2edec
3 changed files with 110 additions and 97 deletions
--- a/hrp/pkg/uixt/ext.go
+++ b/hrp/pkg/uixt/ext.go
@@ -3,9 +3,7 @@ package uixt
 import (
 	"bytes"
 	"fmt"
-	"image"
 	_ "image/gif"
-	"image/jpeg"
 	_ "image/png"
 	"math/rand"
 	"mime"
@@ -217,100 +215,6 @@ func (dExt *DriverExt) Install(filePath string, opts *InstallOptions) error {
 	return dExt.Device.Install(filePath, opts)
 }

-// takeScreenShot takes screenshot and saves image file to $CWD/screenshots/ folder
-func (dExt *DriverExt) takeScreenShot(fileName string) (raw *bytes.Buffer, path string, err error) {
-	// iOS 优先使用 MJPEG 流进行截图，性能最优
-	// 如果 MJPEG 流未开启，则使用 WebDriver 的截图接口
-	if dExt.frame != nil {
-		return dExt.frame, "", nil
-	}
-	if raw, err = dExt.Driver.Screenshot(); err != nil {
-		log.Error().Err(err).Msg("capture screenshot data failed")
-		return nil, "", err
-	}
-
-	// compress image data
-	compressed, err := compressImageBuffer(raw)
-	if err != nil {
-		log.Error().Err(err).Msg("compress screenshot data failed")
-		return nil, "", err
-	}
-
-	// save screenshot to file
-	path = filepath.Join(env.ScreenShotsPath, fileName)
-	path, err = dExt.saveScreenShot(compressed, path)
-	if err != nil {
-		log.Error().Err(err).Msg("save screenshot file failed")
-		return nil, "", err
-	}
-	return compressed, path, nil
-}
-
-func compressImageBuffer(raw *bytes.Buffer) (compressed *bytes.Buffer, err error) {
-	// 解码原始图像数据
-	img, format, err := image.Decode(raw)
-	if err != nil {
-		return nil, err
-	}
-
-	// 创建一个用来保存压缩后数据的buffer
-	var buf bytes.Buffer
-
-	switch format {
-	// Convert to jpeg uniformly and compress with a compression rate of 95
-	case "jpeg", "png":
-		jpegOptions := &jpeg.Options{Quality: 95}
-		err = jpeg.Encode(&buf, img, jpegOptions)
-		if err != nil {
-			return nil, err
-		}
-	default:
-		return nil, fmt.Errorf("unsupported image format: %s", format)
-	}
-
-	// 返回压缩后的图像数据
-	return &buf, nil
-}
-
-// saveScreenShot saves image file with file name
-func (dExt *DriverExt) saveScreenShot(raw *bytes.Buffer, fileName string) (string, error) {
-	// notice: screenshot data is a stream, so we need to copy it to a new buffer
-	copiedBuffer := &bytes.Buffer{}
-	if _, err := copiedBuffer.Write(raw.Bytes()); err != nil {
-		log.Error().Err(err).Msg("copy screenshot buffer failed")
-	}
-
-	img, format, err := image.Decode(copiedBuffer)
-	if err != nil {
-		return "", errors.Wrap(err, "decode screenshot image failed")
-	}
-
-	// The default format uses jpeg for compression
-	screenshotPath := filepath.Join(fmt.Sprintf("%s.%s", fileName, "jpeg"))
-	file, err := os.Create(screenshotPath)
-	if err != nil {
-		return "", errors.Wrap(err, "create screenshot image file failed")
-	}
-	defer func() {
-		_ = file.Close()
-	}()
-
-	switch format {
-	case "jpeg", "png":
-		jpegOptions := &jpeg.Options{}
-		err = jpeg.Encode(file, img, jpegOptions)
-	default:
-		return "", fmt.Errorf("unsupported image format: %s", format)
-	}
-	if err != nil {
-		return "", errors.Wrap(err, "encode screenshot image failed")
-	}
-
-	dExt.cacheStepData.screenShots = append(dExt.cacheStepData.screenShots, screenshotPath)
-	log.Info().Str("path", screenshotPath).Msg("save screenshot file success")
-	return screenshotPath, nil
-}
-
 func (dExt *DriverExt) GetStepCacheData() map[string]interface{} {
 	cacheData := make(map[string]interface{})
 	cacheData["screenshots"] = dExt.cacheStepData.screenShots
--- a/hrp/pkg/uixt/screenshot.go
+++ b/hrp/pkg/uixt/screenshot.go
@@ -0,0 +1,109 @@
+package uixt
+
+import (
+	"bytes"
+	"fmt"
+	"image"
+	"image/jpeg"
+	"os"
+	"path/filepath"
+
+	"github.com/pkg/errors"
+	"github.com/rs/zerolog/log"
+
+	"github.com/httprunner/httprunner/v4/hrp/internal/env"
+)
+
+// GetScreenShot takes screenshot and saves image file to $CWD/screenshots/ folder
+func (dExt *DriverExt) GetScreenShot(fileName string) (raw *bytes.Buffer, path string, err error) {
+	// iOS 优先使用 MJPEG 流进行截图，性能最优
+	// 如果 MJPEG 流未开启，则使用 WebDriver 的截图接口
+	if dExt.frame != nil {
+		return dExt.frame, "", nil
+	}
+	if raw, err = dExt.Driver.Screenshot(); err != nil {
+		log.Error().Err(err).Msg("capture screenshot data failed")
+		return nil, "", err
+	}
+
+	// compress image data
+	compressed, err := compressImageBuffer(raw)
+	if err != nil {
+		log.Error().Err(err).Msg("compress screenshot data failed")
+		return nil, "", err
+	}
+
+	// save screenshot to file
+	path = filepath.Join(env.ScreenShotsPath, fileName)
+	path, err = saveScreenShot(compressed, path)
+	if err != nil {
+		log.Error().Err(err).Msg("save screenshot file failed")
+		return nil, "", err
+	}
+	dExt.cacheStepData.screenShots = append(dExt.cacheStepData.screenShots, path)
+	return compressed, path, nil
+}
+
+func compressImageBuffer(raw *bytes.Buffer) (compressed *bytes.Buffer, err error) {
+	// 解码原始图像数据
+	img, format, err := image.Decode(raw)
+	if err != nil {
+		return nil, err
+	}
+
+	// 创建一个用来保存压缩后数据的buffer
+	var buf bytes.Buffer
+
+	switch format {
+	// Convert to jpeg uniformly and compress with a compression rate of 95
+	case "jpeg", "png":
+		jpegOptions := &jpeg.Options{Quality: 95}
+		err = jpeg.Encode(&buf, img, jpegOptions)
+		if err != nil {
+			return nil, err
+		}
+	default:
+		return nil, fmt.Errorf("unsupported image format: %s", format)
+	}
+
+	// 返回压缩后的图像数据
+	return &buf, nil
+}
+
+// saveScreenShot saves image file with file name
+func saveScreenShot(raw *bytes.Buffer, fileName string) (string, error) {
+	// notice: screenshot data is a stream, so we need to copy it to a new buffer
+	copiedBuffer := &bytes.Buffer{}
+	if _, err := copiedBuffer.Write(raw.Bytes()); err != nil {
+		log.Error().Err(err).Msg("copy screenshot buffer failed")
+	}
+
+	img, format, err := image.Decode(copiedBuffer)
+	if err != nil {
+		return "", errors.Wrap(err, "decode screenshot image failed")
+	}
+
+	// The default format uses jpeg for compression
+	screenshotPath := filepath.Join(fmt.Sprintf("%s.%s", fileName, "jpeg"))
+	file, err := os.Create(screenshotPath)
+	if err != nil {
+		return "", errors.Wrap(err, "create screenshot image file failed")
+	}
+	defer func() {
+		_ = file.Close()
+	}()
+
+	switch format {
+	case "jpeg", "png":
+		jpegOptions := &jpeg.Options{}
+		err = jpeg.Encode(file, img, jpegOptions)
+	default:
+		return "", fmt.Errorf("unsupported image format: %s", format)
+	}
+	if err != nil {
+		return "", errors.Wrap(err, "encode screenshot image failed")
+	}
+
+	log.Info().Str("path", screenshotPath).Msg("save screenshot file success")
+	return screenshotPath, nil
+}
--- a/hrp/pkg/uixt/service_vedem.go
+++ b/hrp/pkg/uixt/service_vedem.go
@@ -392,7 +392,7 @@ func (dExt *DriverExt) GetScreenResult(options ...ActionOption) (screenResult *S
 	if len(screenshotActions) != 0 {
 		fileName = builtin.GenNameWithTimestamp("%d_" + strings.Join(screenshotActions, "_"))
 	}
-	bufSource, imagePath, err := dExt.takeScreenShot(fileName)
+	bufSource, imagePath, err := dExt.GetScreenShot(fileName)
 	if err != nil {
 		return
 	}