feat: save screenshot after action

uixt/android_driver_adb.go

@@ -312,7 +312,7 @@ func (ad *ADBDriver) TapAbsXY(x, y float64, opts ...option.ActionOption) error {
 	if err != nil {
 		return err
 	}
-	defer postHandler(ad, actionOptions)
+	defer postHandler(ad, ACTION_TapAbsXY, actionOptions)
 
 	// adb shell input tap x y
 	xStr := fmt.Sprintf("%.1f", x)
@@ -331,7 +331,7 @@ func (ad *ADBDriver) DoubleTap(x, y float64, opts ...option.ActionOption) error
 	if err != nil {
 		return err
 	}
-	defer postHandler(ad, actionOptions)
+	defer postHandler(ad, ACTION_DoubleTapXY, actionOptions)
 
 	// adb shell input tap x y
 	xStr := fmt.Sprintf("%.1f", x)
@@ -380,7 +380,7 @@ func (ad *ADBDriver) Drag(fromX, fromY, toX, toY float64, opts ...option.ActionO
 	if err != nil {
 		return err
 	}
-	defer postHandler(ad, actionOptions)
+	defer postHandler(ad, ACTION_Drag, actionOptions)
 
 	duration := 200.0
 	if actionOptions.Duration > 0 {
@@ -412,7 +412,7 @@ func (ad *ADBDriver) Swipe(fromX, fromY, toX, toY float64, opts ...option.Action
 	if err != nil {
 		return err
 	}
-	defer postHandler(ad, actionOptions)
+	defer postHandler(ad, ACTION_Swipe, actionOptions)
 
 	// adb shell input swipe fromX fromY toX toY
 	_, err = ad.runShellCommand(

uixt/android_driver_uia2.go

@@ -262,7 +262,7 @@ func (ud *UIA2Driver) DoubleTap(x, y float64, opts ...option.ActionOption) error
 	if err != nil {
 		return err
 	}
-	defer postHandler(ud, actionOptions)
+	defer postHandler(ud, ACTION_DoubleTapXY, actionOptions)
 
 	data := map[string]interface{}{
 		"actions": []interface{}{
@@ -304,7 +304,7 @@ func (ud *UIA2Driver) TapAbsXY(x, y float64, opts ...option.ActionOption) error
 	if err != nil {
 		return err
 	}
-	defer postHandler(ud, actionOptions)
+	defer postHandler(ud, ACTION_TapAbsXY, actionOptions)
 
 	duration := 100.0
 	if actionOptions.PressDuration > 0 {
@@ -367,7 +367,7 @@ func (ud *UIA2Driver) Drag(fromX, fromY, toX, toY float64, opts ...option.Action
 	if err != nil {
 		return err
 	}
-	defer postHandler(ud, actionOptions)
+	defer postHandler(ud, ACTION_Drag, actionOptions)
 
 	data := map[string]interface{}{
 		"startX": fromX,
@@ -398,7 +398,7 @@ func (ud *UIA2Driver) Swipe(fromX, fromY, toX, toY float64, opts ...option.Actio
 	if err != nil {
 		return err
 	}
-	defer postHandler(ud, actionOptions)
+	defer postHandler(ud, ACTION_Swipe, actionOptions)
 
 	duration := 200.0
 	if actionOptions.PressDuration > 0 {

uixt/browser_driver.go

@@ -119,7 +119,7 @@ func (wd *BrowserDriver) Drag(fromX, fromY, toX, toY float64, options ...option.
 	if err != nil {
 		return err
 	}
-	defer postHandler(wd, actionOptions)
+	defer postHandler(wd, ACTION_Drag, actionOptions)
 
 	data := map[string]interface{}{
 		"from_x": fromX,
@@ -518,7 +518,7 @@ func (wd *BrowserDriver) TapFloat(x, y float64, opts ...option.ActionOption) err
 	if err != nil {
 		return err
 	}
-	defer postHandler(wd, actionOptions)
+	defer postHandler(wd, ACTION_TapAbsXY, actionOptions)
 
 	duration := 0.1
 	if actionOptions.Duration > 0 {
@@ -542,7 +542,7 @@ func (wd *BrowserDriver) DoubleTap(x, y float64, options ...option.ActionOption)
 	if err != nil {
 		return err
 	}
-	defer postHandler(wd, actionOptions)
+	defer postHandler(wd, ACTION_DoubleTapXY, actionOptions)
 
 	data := map[string]interface{}{
 		"x": x,

uixt/driver_ext_ai.go

@@ -60,7 +60,7 @@ func (dExt *XTDriver) PlanNextAction(text string, opts ...option.ActionOption) (
 		return nil, errors.New("LLM service is not initialized")
 	}
 
-	compressedBufSource, err := dExt.GetScreenShotBuffer()
+	compressedBufSource, err := getScreenShotBuffer(dExt.IDriver)
 	if err != nil {
 		return nil, err
 	}
@@ -118,7 +118,7 @@ func (dExt *XTDriver) AIAssert(assertion string, opts ...option.ActionOption) er
 		return errors.New("LLM service is not initialized")
 	}
 
-	compressedBufSource, err := dExt.GetScreenShotBuffer()
+	compressedBufSource, err := getScreenShotBuffer(dExt.IDriver)
 	if err != nil {
 		return err
 	}

uixt/driver_ext_screenshot.go

@@ -48,28 +48,10 @@ func (s *ScreenResult) FilterTextsByScope(x1, y1, x2, y2 float64) ai.OCRTexts {
 	})
 }
 
-func (dExt *XTDriver) GetScreenShotBuffer() (compressedBufSource *bytes.Buffer, err error) {
-	// take screenshot
-	bufSource, err := dExt.ScreenShot()
-	if err != nil {
-		return nil, errors.Wrapf(code.DeviceScreenShotError,
-			"take screenshot failed %v", err)
-	}
-
-	// compress screenshot
-	compressBufSource, err := compressImageBuffer(bufSource)
-	if err != nil {
-		return nil, errors.Wrapf(code.DeviceScreenShotError,
-			"compress screenshot failed %v", err)
-	}
-
-	return compressBufSource, nil
-}
-
 // GetScreenResult takes a screenshot, returns the image recognition result
 func (dExt *XTDriver) GetScreenResult(opts ...option.ActionOption) (screenResult *ScreenResult, err error) {
 	// get compressed screenshot buffer
-	compressBufSource, err := dExt.GetScreenShotBuffer()
+	compressBufSource, err := getScreenShotBuffer(dExt.IDriver)
 	if err != nil {
 		return nil, err
 	}
@@ -220,6 +202,25 @@ func (dExt *XTDriver) FindUIResult(opts ...option.ActionOption) (uiResult ai.UIR
 	return
 }
 
+// getScreenShotBuffer takes a screenshot, returns the compressed image buffer
+func getScreenShotBuffer(driver IDriver) (compressedBufSource *bytes.Buffer, err error) {
+	// take screenshot
+	bufSource, err := driver.ScreenShot()
+	if err != nil {
+		return nil, errors.Wrapf(code.DeviceScreenShotError,
+			"take screenshot failed %v", err)
+	}
+
+	// compress screenshot
+	compressBufSource, err := compressImageBuffer(bufSource)
+	if err != nil {
+		return nil, errors.Wrapf(code.DeviceScreenShotError,
+			"compress screenshot failed %v", err)
+	}
+
+	return compressBufSource, nil
+}
+
 // saveScreenShot saves compressed image file with file name
 func saveScreenShot(raw *bytes.Buffer, screenshotPath string) error {
 	// notice: screenshot data is a stream, so we need to copy it to a new buffer
@@ -314,17 +315,16 @@ func MarkUIOperation(driver IDriver, actionType ActionMethod, actionCoordinates
 	}
 
 	// create screenshot save path
-	timestamp := builtin.GenNameWithTimestamp("action_%d")
-	var imagePath string
+	timestamp := builtin.GenNameWithTimestamp("%d")
+	imagePath := filepath.Join(
+		config.GetConfig().ScreenShotsPath,
+		fmt.Sprintf("action_%s_pre_%s.png", timestamp, actionType),
+	)
 
 	if actionType == ACTION_TapAbsXY || actionType == ACTION_DoubleTapXY {
 		if len(actionCoordinates) != 2 {
 			return fmt.Errorf("invalid tap action coordinates: %v", actionCoordinates)
 		}
-		imagePath = filepath.Join(
-			config.GetConfig().ScreenShotsPath,
-			fmt.Sprintf("%s_%s.png", timestamp, actionType),
-		)
 		x, y := actionCoordinates[0], actionCoordinates[1]
 		point := image.Point{X: int(x), Y: int(y)}
 		err = SaveImageWithCircleMarker(compressedBufSource, point, imagePath)
@@ -332,10 +332,6 @@ func MarkUIOperation(driver IDriver, actionType ActionMethod, actionCoordinates
 		if len(actionCoordinates) != 4 {
 			return fmt.Errorf("invalid swipe action coordinates: %v", actionCoordinates)
 		}
-		imagePath = filepath.Join(
-			config.GetConfig().ScreenShotsPath,
-			fmt.Sprintf("%s_%s.png", timestamp, actionType),
-		)
 		fromX, fromY := actionCoordinates[0], actionCoordinates[1]
 		toX, toY := actionCoordinates[2], actionCoordinates[3]
 		from := image.Point{X: int(fromX), Y: int(fromY)}

uixt/driver_handler.go

@@ -2,8 +2,11 @@ package uixt
 
 import (
 	"fmt"
+	"path/filepath"
 	"time"
 
+	"github.com/httprunner/httprunner/v5/internal/builtin"
+	"github.com/httprunner/httprunner/v5/internal/config"
 	"github.com/httprunner/httprunner/v5/uixt/option"
 	"github.com/rs/zerolog/log"
 )
@@ -104,7 +107,7 @@ func preHandler_Swipe(driver IDriver, options *option.ActionOptions, rawFomX, ra
 	}
 	fromX, fromY, toX, toY = options.ApplySwipeOffset(fromX, fromY, toX, toY)
 
-	// mark UI operation
+	// save screenshot before action and mark UI operation
 	if options.PreMarkOperation {
 		if markErr := MarkUIOperation(driver, ACTION_Swipe, []float64{fromX, fromY, toX, toY}); markErr != nil {
 			log.Warn().Err(markErr).Msg("Failed to mark swipe operation")
@@ -114,5 +117,28 @@ func preHandler_Swipe(driver IDriver, options *option.ActionOptions, rawFomX, ra
 	return fromX, fromY, toX, toY, nil
 }
 
-func postHandler(_ IDriver, options *option.ActionOptions) {
+func postHandler(driver IDriver, actionType ActionMethod, options *option.ActionOptions) error {
+	// save screenshot after action
+	if options.PostMarkOperation {
+		// get compressed screenshot buffer
+		compressBufSource, err := getScreenShotBuffer(driver)
+		if err != nil {
+			return err
+		}
+
+		// save compressed screenshot to file
+		timestamp := builtin.GenNameWithTimestamp("%d")
+		imagePath := filepath.Join(
+			config.GetConfig().ScreenShotsPath,
+			fmt.Sprintf("action_%s_post_%s.png", timestamp, actionType),
+		)
+
+		go func() {
+			err := saveScreenShot(compressBufSource, imagePath)
+			if err != nil {
+				log.Error().Err(err).Msg("save screenshot file failed")
+			}
+		}()
+	}
+	return nil
 }

uixt/harmony_driver_hdc.go

@@ -159,7 +159,7 @@ func (hd *HDCDriver) TapAbsXY(x, y float64, opts ...option.ActionOption) error {
 	if err != nil {
 		return err
 	}
-	defer postHandler(hd, actionOptions)
+	defer postHandler(hd, ACTION_TapAbsXY, actionOptions)
 
 	if actionOptions.Identifier != "" {
 		startTime := int(time.Now().UnixMilli())
@@ -191,7 +191,7 @@ func (hd *HDCDriver) Swipe(fromX, fromY, toX, toY float64, opts ...option.Action
 	if err != nil {
 		return err
 	}
-	defer postHandler(hd, actionOptions)
+	defer postHandler(hd, ACTION_Swipe, actionOptions)
 
 	duration := 200
 	if actionOptions.PressDuration > 0 {

uixt/ios_driver_wda.go

@@ -602,7 +602,7 @@ func (wd *WDADriver) TapAbsXY(x, y float64, opts ...option.ActionOption) error {
 	if err != nil {
 		return err
 	}
-	defer postHandler(wd, actionOptions)
+	defer postHandler(wd, ACTION_TapAbsXY, actionOptions)
 
 	data := map[string]interface{}{
 		"x": x,
@@ -627,7 +627,7 @@ func (wd *WDADriver) DoubleTap(x, y float64, opts ...option.ActionOption) error
 	if err != nil {
 		return err
 	}
-	defer postHandler(wd, actionOptions)
+	defer postHandler(wd, ACTION_DoubleTapXY, actionOptions)
 
 	data := map[string]interface{}{
 		"x": x,
@@ -664,7 +664,7 @@ func (wd *WDADriver) Drag(fromX, fromY, toX, toY float64, opts ...option.ActionO
 	if err != nil {
 		return err
 	}
-	defer postHandler(wd, actionOptions)
+	defer postHandler(wd, ACTION_Drag, actionOptions)
 
 	data := map[string]interface{}{
 		"fromX": math.Round(fromX*10) / 10,