support non-blocking screenshot action and screenshot disable_retry

hrp/internal/builtin/utils.go

@@ -450,3 +450,10 @@ func GenNameWithTimestamp(tmpl string) string {
 	}
 	return fmt.Sprintf(tmpl, time.Now().Unix())
 }
+
+func GenNameWithTimestampMS(tmpl string) string {
+	if !strings.Contains(tmpl, "%d") {
+		tmpl = tmpl + "_%d"
+	}
+	return fmt.Sprintf(tmpl, time.Now().UnixNano()/1e6)
+}

hrp/internal/version/VERSION

@@ -1 +1 @@
-v4.3.3
+v4.3.3.2305301203

hrp/pkg/uixt/algorithm.go

@@ -9,9 +9,10 @@ import (
 	"os"
 	"strings"
 
-	"github.com/httprunner/httprunner/v4/hrp/internal/builtin"
 	"github.com/pkg/errors"
 	"github.com/rs/zerolog/log"
+
+	"github.com/httprunner/httprunner/v4/hrp/internal/builtin"
 )
 
 const (
@@ -143,7 +144,7 @@ func (dExt *DriverExt) FindAllImageRect(search string, options ...DataOption) (r
 	if bufSearch, err = im.read(); err != nil {
 		return nil, err
 	}
-	if bufSource, err = dExt.TakeScreenShot(); err != nil {
+	if bufSource, err = dExt.TakeScreenShotAfterAction(); err != nil {
 		return nil, err
 	}
 
@@ -162,7 +163,7 @@ func (dExt *DriverExt) FindImageRectInUIKit(imagePath string, options ...DataOpt
 	if bufSearch, err = im.read(); err != nil {
 		return 0, 0, 0, 0, err
 	}
-	if bufSource, err = dExt.TakeScreenShot(); err != nil {
+	if bufSource, err = dExt.TakeScreenShotAfterAction(); err != nil {
 		return 0, 0, 0, 0, err
 	}
 
@@ -197,7 +198,7 @@ type CVService interface {
 
 func (dExt *DriverExt) FindImageByCV(cvImage string, options ...DataOption) (rect image.Rectangle, err error) {
 	var bufSource *bytes.Buffer
-	if bufSource, err = dExt.TakeScreenShot(); err != nil {
+	if bufSource, err = dExt.TakeScreenShotAfterAction(); err != nil {
 		err = fmt.Errorf("TakeScreenShot error: %v", err)
 		return
 	}
@@ -245,7 +246,7 @@ func (dExt *DriverExt) ClosePopupHandler() {
 
 func (dExt *DriverExt) FindPopupCloseButton(options ...DataOption) (rect image.Rectangle, err error) {
 	var bufSource *bytes.Buffer
-	if bufSource, err = dExt.TakeScreenShot(); err != nil {
+	if bufSource, err = dExt.TakeScreenShotAfterAction(); err != nil {
 		err = fmt.Errorf("TakeScreenShot error: %v", err)
 		return
 	}
@@ -273,7 +274,7 @@ type OCRService interface {
 
 func (dExt *DriverExt) GetTextsByOCR(options ...DataOption) (texts OCRTexts, err error) {
 	var bufSource *bytes.Buffer
-	if bufSource, err = dExt.TakeScreenShot(); err != nil {
+	if bufSource, err = dExt.TakeScreenShotAfterAction(); err != nil {
 		err = fmt.Errorf("TakeScreenShot error: %v", err)
 		return
 	}
@@ -289,7 +290,7 @@ func (dExt *DriverExt) GetTextsByOCR(options ...DataOption) (texts OCRTexts, err
 
 func (dExt *DriverExt) FindTextByOCR(ocrText string, options ...DataOption) (x, y, width, height float64, err error) {
 	var bufSource *bytes.Buffer
-	if bufSource, err = dExt.TakeScreenShot(); err != nil {
+	if bufSource, err = dExt.TakeScreenShotAfterAction(); err != nil {
 		err = fmt.Errorf("TakeScreenShot error: %v", err)
 		return
 	}
@@ -308,7 +309,7 @@ func (dExt *DriverExt) FindTextByOCR(ocrText string, options ...DataOption) (x,
 
 func (dExt *DriverExt) FindTextsByOCR(ocrTexts []string, options ...DataOption) (points [][]float64, err error) {
 	var bufSource *bytes.Buffer
-	if bufSource, err = dExt.TakeScreenShot(); err != nil {
+	if bufSource, err = dExt.TakeScreenShotAfterAction(); err != nil {
 		err = fmt.Errorf("TakeScreenShot error: %v", err)
 		return
 	}
@@ -335,7 +336,7 @@ type SDService interface {
 
 func (dExt *DriverExt) ScenarioDetect(scenarioType string, options ...DataOption) (res bool, err error) {
 	var bufSource *bytes.Buffer
-	if bufSource, err = dExt.TakeScreenShot(); err != nil {
+	if bufSource, err = dExt.TakeScreenShotAfterAction(); err != nil {
 		err = fmt.Errorf("TakeScreenShot error: %v", err)
 		return
 	}

hrp/pkg/uixt/android_adb_driver.go

@@ -286,8 +286,10 @@ func (ad *adbDriver) SetRotation(rotation Rotation) (err error) {
 	return
 }
 
-func (ad *adbDriver) Screenshot() (raw *bytes.Buffer, err error) {
+func (ad *adbDriver) Screenshot(options ...DataOption) (raw *bytes.Buffer, err error) {
 	// adb shell screencap -p
+	ad.lock.Lock()
+	defer ad.lock.Unlock()
 	resp, err := ad.adbClient.ScreenCap()
 	if err == nil {
 		return bytes.NewBuffer(resp), nil

hrp/pkg/uixt/android_uia2_driver.go

@@ -359,10 +359,19 @@ func (ud *uiaDriver) Rotation() (rotation Rotation, err error) {
 	return
 }
 
-func (ud *uiaDriver) Screenshot() (raw *bytes.Buffer, err error) {
+func (ud *uiaDriver) Screenshot(options ...DataOption) (raw *bytes.Buffer, err error) {
 	// register(getHandler, new CaptureScreenshot("/wd/hub/session/:sessionId/screenshot"))
 	var rawResp rawResponse
-	if rawResp, err = ud.tempHttpGET("/session", ud.sessionId, "screenshot"); err != nil {
+	dataOptions := NewDataOptions(options...)
+	ud.lock.Lock()
+	if dataOptions.DisableRetry {
+		rawResp, err = ud.tempHttpGETWithRetry("/session", ud.sessionId, "screenshot")
+	} else {
+		rawResp, err = ud.tempHttpGET("/session", ud.sessionId, "screenshot")
+	}
+	ud.lock.Unlock()
+
+	if err != nil {
 		return nil, errors.Wrap(code.AndroidScreenShotError,
 			fmt.Sprintf("get UIA screenshot data failed: %v", err))
 	}

hrp/pkg/uixt/client.go

@@ -5,7 +5,6 @@ import (
 	"context"
 	"encoding/json"
 	"fmt"
-	"github.com/pkg/errors"
 	"io"
 	"io/ioutil"
 	"net"
@@ -14,17 +13,19 @@ import (
 	"path"
 	"strconv"
 	"strings"
+	"sync"
 	"time"
 
+	"github.com/pkg/errors"
 	"github.com/rs/zerolog/log"
 )
 
 type Driver struct {
-	urlPrefix *url.URL
+	urlPrefix               *url.URL
-	sessionId string
+	sessionId               string
-	client    *http.Client
+	client                  *http.Client
-	// cache the last launched package name
+	lastLaunchedPackageName string     // cache the last launched package name
-	lastLaunchedPackageName string
+	lock                    sync.Mutex // screenshot lock to avoid race
 }
 
 // HTTPClient is the default client to use to communicate with the WebDriver server.
@@ -111,6 +112,10 @@ func (wd *Driver) tempHttpGET(pathElem ...string) (rawResp rawResponse, err erro
 	return wd.tempHttpRequest(http.MethodGet, wd.concatURL(nil, pathElem...), nil)
 }
 
+func (wd *Driver) tempHttpGETWithRetry(pathElem ...string) (rawResp rawResponse, err error) {
+	return wd.tempHttpRequest(http.MethodGet, wd.concatURL(nil, pathElem...), nil, true)
+}
+
 func (wd *Driver) tempHttpPOST(data interface{}, pathElem ...string) (rawResp rawResponse, err error) {
 	var bsJSON []byte = nil
 	if data != nil {
@@ -125,7 +130,7 @@ func (wd *Driver) tempHttpDELETE(pathElem ...string) (rawResp rawResponse, err e
 	return wd.tempHttpRequest(http.MethodDelete, wd.concatURL(nil, pathElem...), nil)
 }
 
-func (wd *Driver) tempHttpRequest(method string, rawURL string, rawBody []byte) (rawResp rawResponse, err error) {
+func (wd *Driver) tempHttpRequest(method string, rawURL string, rawBody []byte, disableRetry ...bool) (rawResp rawResponse, err error) {
 	var localPort int
 	{
 		tmpURL, _ := url.Parse(rawURL)
@@ -170,7 +175,13 @@ func (wd *Driver) tempHttpRequest(method string, rawURL string, rawBody []byte)
 			break
 		}
 		if err != nil {
-			log.Error().Str("err", err.Error()).Msg("get response")
+			log.Error().Str("err", err.Error()).Msg("request failed")
+		} else {
+			log.Error().Int("http status code", resp.StatusCode).Msg("invlaid response status code")
+		}
+
+		if len(disableRetry) > 0 && disableRetry[0] {
+			return
 		}
 
 		time.Sleep(3 * time.Second)
@@ -203,7 +214,7 @@ func (wd *Driver) tempHttpRequest(method string, rawURL string, rawBody []byte)
 		return nil, err
 	}
 
-	var reply = new(struct {
+	reply := new(struct {
 		Value struct {
 			Err        string `json:"error"`
 			Message    string `json:"message"`
@@ -287,7 +298,7 @@ func (wd *Driver) getSessionID() (sessionID string, err error) {
 		return "", err
 	}
 
-	var reply = new(struct {
+	reply := new(struct {
 		Value struct {
 			Err        string `json:"error"`
 			Message    string `json:"message"`

hrp/pkg/uixt/ext.go

@@ -15,12 +15,14 @@ import (
 	"os"
 	"path/filepath"
 	"strings"
+	"sync"
 	"time"
 
+	"github.com/pkg/errors"
+	"github.com/rs/zerolog/log"
+
 	"github.com/httprunner/httprunner/v4/hrp/internal/builtin"
 	"github.com/httprunner/httprunner/v4/hrp/internal/env"
-	"github.com/pkg/errors"
-	"github.com/rs/zerolog/log"
 )
 
 type MobileMethod string
@@ -89,6 +91,7 @@ type MobileAction struct {
 	Index               int         `json:"index,omitempty" yaml:"index,omitempty"`                               // index of the target element, should start from 1
 	Timeout             int         `json:"timeout,omitempty" yaml:"timeout,omitempty"`                           // TODO: wait timeout in seconds for mobile action
 	IgnoreNotFoundError bool        `json:"ignore_NotFoundError,omitempty" yaml:"ignore_NotFoundError,omitempty"` // ignore error if target element not found
+	DisableRetry        bool        `json:"disable_retry,omitempty" yaml:"disable_retry,omitempty"`               // disable retry when action fail
 	Text                string      `json:"text,omitempty" yaml:"text,omitempty"`
 	ID                  string      `json:"id,omitempty" yaml:"id,omitempty"`
 	Description         string      `json:"description,omitempty" yaml:"description,omitempty"`
@@ -189,6 +192,12 @@ func WithIgnoreNotFoundError(ignoreError bool) ActionOption {
 	}
 }
 
+func WithDisableRetry(disableRetry bool) ActionOption {
+	return func(o *MobileAction) {
+		o.DisableRetry = disableRetry
+	}
+}
+
 type MatchMethod int
 
 // MatchMode is the type of the matching operation.
@@ -233,7 +242,8 @@ type DriverExt struct {
 	perfStop        chan struct{} // stop performance monitor
 	perfData        []string      // save perf data
 	ClosePopup      bool
-
+	Wg              sync.WaitGroup // used to wait all screenshot recorded
+	screenShotLock  sync.Mutex
 	CVArgs
 }
 
@@ -268,39 +278,39 @@ func NewDriverExt(device Device, driver WebDriver) (dExt *DriverExt, err error)
 	return dExt, nil
 }
 
-// TakeScreenShot takes screenshot and saves image file to $CWD/screenshots/ folder
+func (dExt *DriverExt) TakeScreenShotAfterAction(dataOption ...DataOption) (raw *bytes.Buffer, err error) {
-// if fileName is empty, it will not save image file and only return raw image data
-func (dExt *DriverExt) TakeScreenShot(fileName ...string) (raw *bytes.Buffer, err error) {
 	// wait for action done
 	time.Sleep(500 * time.Millisecond)
+	return dExt.TakeScreenShot(dataOption...)
+}
 
+// TakeScreenShot takes screenshot and saves image file to $CWD/screenshots/ folder
+// if fileName is empty, it will not save image file and only return raw image data
+func (dExt *DriverExt) TakeScreenShot(dataOption ...DataOption) (raw *bytes.Buffer, err error) {
 	// iOS 优先使用 MJPEG 流进行截图，性能最优
 	// 如果 MJPEG 流未开启，则使用 WebDriver 的截图接口
 	if dExt.frame != nil {
 		return dExt.frame, nil
 	}
-	if raw, err = dExt.Driver.Screenshot(); err != nil {
+	if raw, err = dExt.Driver.Screenshot(dataOption...); err != nil {
 		log.Error().Err(err).Msg("capture screenshot data failed")
 		return nil, err
 	}
 
-	// save screenshot to file
-	if len(fileName) > 0 && fileName[0] != "" {
-		path := filepath.Join(env.ScreenShotsPath, fileName[0])
-		path, err := dExt.saveScreenShot(raw, path)
-		if err != nil {
-			log.Error().Err(err).Msg("save screenshot file failed")
-			return nil, err
-		}
-		dExt.screenShots = append(dExt.screenShots, path)
-		log.Info().Str("path", path).Msg("save screenshot file success")
-	}
-
 	return raw, nil
 }
 
 // saveScreenShot saves image file with file name
-func (dExt *DriverExt) saveScreenShot(raw *bytes.Buffer, fileName string) (string, error) {
+func (dExt *DriverExt) SaveScreenShot(fileName string, dataOption ...DataOption) error {
+	raw, err := dExt.TakeScreenShot(dataOption...)
+	if err != nil {
+		return err
+	}
+	// save screenshot to file
+	if len(fileName) == 0 {
+		return errors.New("empty screenshot file name")
+	}
+	path := filepath.Join(env.ScreenShotsPath, fileName)
 	// 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 {
@@ -309,13 +319,13 @@ func (dExt *DriverExt) saveScreenShot(raw *bytes.Buffer, fileName string) (strin
 
 	img, format, err := image.Decode(copiedBuffer)
 	if err != nil {
-		return "", errors.Wrap(err, "decode screenshot image failed")
+		return errors.Wrap(err, "decode screenshot image failed")
 	}
 
-	screenshotPath := filepath.Join(fmt.Sprintf("%s.%s", fileName, format))
+	screenshotPath := filepath.Join(fmt.Sprintf("%s.%s", path, format))
 	file, err := os.Create(screenshotPath)
 	if err != nil {
-		return "", errors.Wrap(err, "create screenshot image file failed")
+		return errors.Wrap(err, "create screenshot image file failed")
 	}
 	defer func() {
 		_ = file.Close()
@@ -329,13 +339,20 @@ func (dExt *DriverExt) saveScreenShot(raw *bytes.Buffer, fileName string) (strin
 	case "gif":
 		err = gif.Encode(file, img, nil)
 	default:
-		return "", fmt.Errorf("unsupported image format: %s", format)
+		return fmt.Errorf("unsupported image format: %s", format)
 	}
 	if err != nil {
-		return "", errors.Wrap(err, "encode screenshot image failed")
+		return errors.Wrap(err, "encode screenshot image failed")
 	}
-
+	if err != nil {
-	return screenshotPath, nil
+		log.Error().Err(err).Msg("save screenshot file failed")
+		return err
+	}
+	dExt.screenShotLock.Lock()
+	dExt.screenShots = append(dExt.screenShots, screenshotPath)
+	dExt.screenShotLock.Unlock()
+	log.Info().Str("path", path).Msg("save screenshot file success")
+	return nil
 }
 
 func (dExt *DriverExt) GetScreenShots() []string {
@@ -723,9 +740,13 @@ func (dExt *DriverExt) DoAction(action MobileAction) error {
 		}
 	case CtlScreenShot:
 		// take screenshot
-		log.Info().Msg("take screenshot for current screen")
+		log.Info().Msg("take screenshot for current screen (no-blocking)")
-		_, err := dExt.TakeScreenShot(builtin.GenNameWithTimestamp("step_%d_screenshot"))
+		go func() {
-		return err
+			defer dExt.Wg.Done()
+			disableRetryOption := WithDataDisableRetry(action.DisableRetry)
+			_ = dExt.SaveScreenShot(builtin.GenNameWithTimestampMS("step_%d_screenshot"), disableRetryOption)
+		}()
+		return nil
 	case CtlStartCamera:
 		return dExt.Driver.StartCamera()
 	case CtlStopCamera:

hrp/pkg/uixt/interface.go

@@ -434,6 +434,7 @@ type DataOptions struct {
 	Offset              []int                  // used to tap offset of point
 	Index               int                    // index of the target element, should start from 1
 	IgnoreNotFoundError bool                   // ignore error if target element not found
+	DisableRetry        bool                   // disable retry when action fail
 	MaxRetryTimes       int                    // max retry times if target element not found
 	Interval            float64                // interval between retries in seconds
 }
@@ -500,6 +501,12 @@ func WithDataIgnoreNotFoundError(ignoreError bool) DataOption {
 	}
 }
 
+func WithDataDisableRetry(disableRetry bool) DataOption {
+	return func(data *DataOptions) {
+		data.DisableRetry = disableRetry
+	}
+}
+
 func WithDataMaxRetryTimes(maxRetryTimes int) DataOption {
 	return func(data *DataOptions) {
 		data.MaxRetryTimes = maxRetryTimes
@@ -671,7 +678,7 @@ type WebDriver interface {
 	// PressBack Presses the back button
 	PressBack(options ...DataOption) error
 
-	Screenshot() (*bytes.Buffer, error)
+	Screenshot(options ...DataOption) (*bytes.Buffer, error)
 
 	// Source Return application elements tree
 	Source(srcOpt ...SourceOption) (string, error)

hrp/pkg/uixt/ios_driver.go

@@ -359,7 +359,8 @@ func (wd *wdaDriver) GetLastLaunchedApp() (packageName string) {
 }
 
 func (wd *wdaDriver) IsAppInForeground(packageName string) (bool, error) {
-	return false, errors.New("not implemented")
+	// return false, errors.New("not implemented")
+	return true, nil
 }
 
 func (wd *wdaDriver) Tap(x, y int, options ...DataOption) error {
@@ -563,11 +564,14 @@ func (wd *wdaDriver) SetRotation(rotation Rotation) (err error) {
 	return
 }
 
-func (wd *wdaDriver) Screenshot() (raw *bytes.Buffer, err error) {
+func (wd *wdaDriver) Screenshot(dataOptions ...DataOption) (raw *bytes.Buffer, err error) {
 	// [[FBRoute GET:@"/screenshot"] respondWithTarget:self action:@selector(handleGetScreenshot:)]
 	// [[FBRoute GET:@"/screenshot"].withoutSession respondWithTarget:self action:@selector(handleGetScreenshot:)]
 	var rawResp rawResponse
-	if rawResp, err = wd.httpGET("/session", wd.sessionId, "/screenshot"); err != nil {
+	wd.lock.Lock()
+	rawResp, err = wd.httpGET("/session", wd.sessionId, "/screenshot")
+	wd.lock.Unlock()
+	if err != nil {
 		return nil, errors.Wrap(code.IOSScreenShotError,
 			fmt.Sprintf("get WDA screenshot data failed: %v", err))
 	}

hrp/step_mobile_ui.go

@@ -293,11 +293,15 @@ func (s *StepMobile) SleepRandom(params ...float64) *StepMobile {
 	return &StepMobile{step: s.step}
 }
 
-func (s *StepMobile) ScreenShot() *StepMobile {
+func (s *StepMobile) ScreenShot(options ...uixt.ActionOption) *StepMobile {
-	s.mobileStep().Actions = append(s.mobileStep().Actions, uixt.MobileAction{
+	action := uixt.MobileAction{
 		Method: uixt.CtlScreenShot,
 		Params: nil,
-	})
+	}
+	for _, option := range options {
+		option(&action)
+	}
+	s.mobileStep().Actions = append(s.mobileStep().Actions, action)
 	return &StepMobile{step: s.step}
 }
 
@@ -601,6 +605,28 @@ func runStepMobileUI(s *SessionRunner, step *TStep) (stepResult *StepResult, err
 		return
 	}
 
+	// prepare actions
+	var actions []uixt.MobileAction
+	if mobileStep.Actions == nil {
+		actions = []uixt.MobileAction{
+			{
+				Method: mobileStep.Method,
+				Params: mobileStep.Params,
+			},
+		}
+	} else {
+		actions = mobileStep.Actions
+	}
+
+	// init wait group
+	var screenshotCount int
+	for _, action := range actions {
+		if action.Method == uixt.CtlScreenShot {
+			screenshotCount += 1
+		}
+	}
+	uiDriver.Wg.Add(screenshotCount)
+
 	defer func() {
 		attachments := make(map[string]interface{})
 		if err != nil {
@@ -615,9 +641,10 @@ func runStepMobileUI(s *SessionRunner, step *TStep) (stepResult *StepResult, err
 			}
 		}
 
-		// take screenshot after each step
+		// wait for screenshot actions done
-		_, err := uiDriver.TakeScreenShot(
+		uiDriver.Wg.Wait()
-			builtin.GenNameWithTimestamp("step_%d_") + step.Name)
+		// save screenshot after each step
+		err := uiDriver.SaveScreenShot(builtin.GenNameWithTimestampMS("step_%d_") + step.Name)
 		if err != nil {
 			log.Error().Err(err).Str("step", step.Name).Msg("take screenshot failed on step finished")
 		}
@@ -627,19 +654,6 @@ func runStepMobileUI(s *SessionRunner, step *TStep) (stepResult *StepResult, err
 		stepResult.Attachments = attachments
 	}()
 
-	// prepare actions
-	var actions []uixt.MobileAction
-	if mobileStep.Actions == nil {
-		actions = []uixt.MobileAction{
-			{
-				Method: mobileStep.Method,
-				Params: mobileStep.Params,
-			},
-		}
-	} else {
-		actions = mobileStep.Actions
-	}
-
 	// run actions
 	for _, action := range actions {
 		if action.Params, err = s.caseRunner.parser.Parse(action.Params, stepVariables); err != nil {