refactor: driver ScreenShot

pkg/uixt/android_driver_adb.go

@@ -520,13 +520,26 @@ func (ad *ADBDriver) SetRotation(rotation types.Rotation) (err error) {
 	return
 }
 
-func (ad *ADBDriver) ScreenShot() (raw *bytes.Buffer, err error) {
+func (ad *ADBDriver) ScreenShot(opts ...option.ActionOption) (raw *bytes.Buffer, err error) {
 	resp, err := ad.runShellCommand("screencap", "-p")
 	if err != nil {
-		return nil, errors.Wrap(err, "adb screencap failed")
+		return nil, errors.Wrapf(code.DeviceScreenShotError,
+			"adb screencap failed %v", err)
+	}
+	raw = bytes.NewBuffer([]byte(resp))
+
+	actionOptions := option.NewActionOptions(opts...)
+	if actionOptions.ScreenShotFileName != "" {
+		// save screenshot to file
+		path, err := saveScreenShot(raw, actionOptions.ScreenShotFileName)
+		if err != nil {
+			return nil, errors.Wrapf(code.DeviceScreenShotError,
+				"save screenshot file failed %v", err)
+		}
+		log.Info().Str("path", path).Msg("screenshot saved")
 	}
 
-	return bytes.NewBuffer([]byte(resp)), nil
+	return raw, nil
 }
 
 func (ad *ADBDriver) Source(srcOpt ...option.SourceOption) (source string, err error) {

pkg/uixt/android_driver_uia2.go

@@ -533,10 +533,10 @@ func (ud *UIA2Driver) Rotation() (rotation types.Rotation, err error) {
 	return
 }
 
-func (ud *UIA2Driver) ScreenShot() (raw *bytes.Buffer, err error) {
+func (ud *UIA2Driver) ScreenShot(opts ...option.ActionOption) (raw *bytes.Buffer, err error) {
 	// https://bytedance.larkoffice.com/docx/C8qEdmSHnoRvMaxZauocMiYpnLh
 	// ui2截图受内存影响，改为adb截图
-	return ud.ADBDriver.ScreenShot()
+	return ud.ADBDriver.ScreenShot(opts...)
 }
 
 func (ud *UIA2Driver) Source(srcOpt ...option.SourceOption) (source string, err error) {

pkg/uixt/driver.go

@@ -43,7 +43,7 @@ type IDriver interface {
 	BatteryInfo() (types.BatteryInfo, error)
 	ForegroundInfo() (app types.AppInfo, err error)
 	WindowSize() (types.Size, error)
-	ScreenShot() (*bytes.Buffer, error)
+	ScreenShot(opts ...option.ActionOption) (*bytes.Buffer, error)
 	ScreenRecord(duration time.Duration) (videoPath string, err error)
 	Source(srcOpt ...option.SourceOption) (string, error)
 	Orientation() (orientation types.Orientation, err error)
@@ -97,11 +97,10 @@ var _ IDriverExt = (*XTDriver)(nil)
 type IDriverExt interface {
 	GetScreenResult(opts ...option.ActionOption) (screenResult *ScreenResult, err error)
 	GetScreenTexts(opts ...option.ActionOption) (ocrTexts ai.OCRTexts, err error)
-	GetScreenShot(fileName string) (raw *bytes.Buffer, path string, err error)
 
 	// tap with AI
 	TapByOCR(text string, opts ...option.ActionOption) error
-	TapByCV(opts ...option.ActionOption) error
+	TapByCV(opts ...option.ActionOption) error // TODO: refactor
 
 	// swipe
 	SwipeRelative(fromX, fromY, toX, toY float64, opts ...option.ActionOption) error

pkg/uixt/driver_screenshot.go

@@ -67,7 +67,8 @@ func (dExt *XTDriver) GetScreenResult(opts ...option.ActionOption) (screenResult
 
 	// get screenshot info with retry
 	for i := 0; i < 3; i++ {
-		bufSource, imagePath, err = dExt.GetScreenShot(fileName)
+		imagePath = filepath.Join(config.ScreenShotsPath, fileName)
+		bufSource, err = dExt.IDriver.ScreenShot(option.WithScreenShotFileName(imagePath))
 		if err != nil {
 			lastErr = err
 			time.Sleep(time.Second * 1)
@@ -188,24 +189,6 @@ func (dExt *XTDriver) FindUIResult(opts ...option.ActionOption) (point ai.PointF
 	return
 }
 
-// GetScreenShot takes screenshot and saves image file to $CWD/screenshots/ folder
-func (dExt *XTDriver) GetScreenShot(fileName string) (raw *bytes.Buffer, path string, err error) {
-	if raw, err = dExt.ScreenShot(); err != nil {
-		log.Error().Err(err).Msg("capture screenshot data failed")
-		return nil, "", errors.Wrap(code.DeviceScreenShotError, err.Error())
-	}
-
-	// save screenshot to file
-	path = filepath.Join(config.ScreenShotsPath, fileName)
-	path, err = saveScreenShot(raw, path)
-	if err != nil {
-		log.Error().Err(err).Msg("save screenshot file failed")
-		return nil, "", errors.Wrap(code.DeviceScreenShotError,
-			fmt.Sprintf("save screenshot file failed: %s", err.Error()))
-	}
-	return raw, path, nil
-}
-
 // saveScreenShot saves compressed 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

pkg/uixt/driver_screenshot_test.go

@@ -3,21 +3,21 @@
 package uixt
 
 import (
+	"path/filepath"
 	"testing"
+
+	"github.com/httprunner/httprunner/v5/internal/config"
+	"github.com/httprunner/httprunner/v5/pkg/uixt/option"
 )
 
 func TestGetScreenShot(t *testing.T) {
 	setupAndroidAdbDriver(t)
 
-	fileName := "test_screenshot"
-	_, path, err := driverExt.GetScreenShot(fileName)
+	imagePath := filepath.Join(config.ScreenShotsPath, "test_screenshot")
+	_, err := driverExt.IDriver.ScreenShot(option.WithScreenShotFileName(imagePath))
 	if err != nil {
 		t.Fatalf("GetScreenShot failed: %v", err)
 	}
 
-	if path == "" {
-		t.Fatal("screenshot path is empty")
-	}
-
-	t.Logf("screenshot saved at: %s", path)
+	t.Logf("screenshot saved at: %s", imagePath)
 }

pkg/uixt/harmony_driver_hdc.go

@@ -8,8 +8,10 @@ import (
 	"time"
 
 	"code.byted.org/iesqa/ghdc"
+	"github.com/pkg/errors"
 	"github.com/rs/zerolog/log"
 
+	"github.com/httprunner/httprunner/v5/code"
 	"github.com/httprunner/httprunner/v5/pkg/uixt/option"
 	"github.com/httprunner/httprunner/v5/pkg/uixt/types"
 )
@@ -217,13 +219,13 @@ func (hd *HDCDriver) PressHarmonyKeyCode(keyCode ghdc.KeyCode) (err error) {
 	return hd.uiDriver.PressKey(keyCode)
 }
 
-func (hd *HDCDriver) ScreenShot() (*bytes.Buffer, error) {
+func (hd *HDCDriver) ScreenShot(opts ...option.ActionOption) (*bytes.Buffer, error) {
 	tempDir := os.TempDir()
 	screenshotPath := fmt.Sprintf("%s/screenshot_%d.png", tempDir, time.Now().Unix())
 	err := hd.uiDriver.Screenshot(screenshotPath)
 	if err != nil {
-		log.Error().Err(err).Msg("failed to screenshot")
-		return nil, err
+		return nil, errors.Wrapf(code.DeviceScreenShotError,
+			"hdc screencap failed %v", err)
 	}
 	defer func() {
 		_ = os.Remove(screenshotPath)
@@ -234,7 +236,20 @@ func (hd *HDCDriver) ScreenShot() (*bytes.Buffer, error) {
 		log.Error().Err(err).Msg("failed to screenshot")
 		return nil, err
 	}
-	return bytes.NewBuffer(raw), nil
+	rawBuffer := bytes.NewBuffer(raw)
+
+	actionOptions := option.NewActionOptions(opts...)
+	if actionOptions.ScreenShotFileName != "" {
+		// save screenshot to file
+		path, err := saveScreenShot(rawBuffer, actionOptions.ScreenShotFileName)
+		if err != nil {
+			return nil, errors.Wrapf(code.DeviceScreenShotError,
+				"save screenshot file failed %v", err)
+		}
+		log.Info().Str("path", path).Msg("screenshot saved")
+	}
+
+	return rawBuffer, nil
 }
 
 func (hd *HDCDriver) Source(srcOpt ...option.SourceOption) (string, error) {

pkg/uixt/ios_driver_wda.go

@@ -280,19 +280,31 @@ func (wd *WDADriver) Screen() (screen ai.Screen, err error) {
 	return
 }
 
-func (wd *WDADriver) ScreenShot() (raw *bytes.Buffer, err error) {
+func (wd *WDADriver) ScreenShot(opts ...option.ActionOption) (raw *bytes.Buffer, err error) {
 	// [[FBRoute GET:@"/screenshot"] respondWithTarget:self action:@selector(handleGetScreenshot:)]
 	// [[FBRoute GET:@"/screenshot"].withoutSession respondWithTarget:self action:@selector(handleGetScreenshot:)]
-	var rawResp DriverRawResponse
-	if rawResp, err = wd.httpGET("/session", wd.Session.ID, "/screenshot"); err != nil {
+	rawResp, err := wd.httpGET("/session", wd.Session.ID, "/screenshot")
+	if err != nil {
 		return nil, errors.Wrap(code.DeviceScreenShotError,
-			fmt.Sprintf("get WDA screenshot data failed: %v", err))
+			fmt.Sprintf("WDA screenshot failed %v", err))
 	}
-
-	if raw, err = rawResp.ValueDecodeAsBase64(); err != nil {
+	raw, err = rawResp.ValueDecodeAsBase64()
+	if err != nil {
 		return nil, errors.Wrap(code.DeviceScreenShotError,
 			fmt.Sprintf("decode WDA screenshot data failed: %v", err))
 	}
+
+	actionOptions := option.NewActionOptions(opts...)
+	if actionOptions.ScreenShotFileName != "" {
+		// save screenshot to file
+		path, err := saveScreenShot(raw, actionOptions.ScreenShotFileName)
+		if err != nil {
+			return nil, errors.Wrapf(code.DeviceScreenShotError,
+				"save screenshot file failed %v", err)
+		}
+		log.Info().Str("path", path).Msg("screenshot saved")
+	}
+
 	return
 }