From 6bb67e8a68dfac441ae45540967fa9a3cbd0a435 Mon Sep 17 00:00:00 2001
From: "lilong.129" <lilong.129@bytedance.com>
Date: Fri, 13 Sep 2024 23:08:19 +0800
Subject: [PATCH] feat: GetScreenResult with max retry times option

---
 hrp/pkg/uixt/screenshot.go | 51 +++++++++++++++++++++++++-------------
 1 file changed, 34 insertions(+), 17 deletions(-)

diff --git a/hrp/pkg/uixt/screenshot.go b/hrp/pkg/uixt/screenshot.go
index 5fd1de14..f6a13890 100644
--- a/hrp/pkg/uixt/screenshot.go
+++ b/hrp/pkg/uixt/screenshot.go
@@ -42,30 +42,47 @@ func (dExt *DriverExt) GetScreenResult(options ...ActionOption) (screenResult *S
 	if len(screenshotActions) != 0 {
 		fileName = builtin.GenNameWithTimestamp("%d_" + strings.Join(screenshotActions, "_"))
 	}
-	bufSource, imagePath, err := dExt.GetScreenShot(fileName)
-	if err != nil {
-		return
+	if actionOptions.MaxRetryTimes == 0 {
+		actionOptions.MaxRetryTimes = 1
 	}
 
-	windowSize, err := dExt.Driver.WindowSize()
+	var bufSource *bytes.Buffer
+	var imageResult *ImageResult
+	var imagePath string
+	var windowSize Size
+	for i := 0; i < actionOptions.MaxRetryTimes; i++ {
+		bufSource, imagePath, err = dExt.GetScreenShot(fileName)
+		if err != nil {
+			continue
+		}
+
+		windowSize, err = dExt.Driver.WindowSize()
+		if err != nil {
+			err = errors.Wrap(code.MobileUIDriverError, err.Error())
+			continue
+		}
+
+		screenResult = &ScreenResult{
+			bufSource:  bufSource,
+			ImagePath:  imagePath,
+			Tags:       nil,
+			Resolution: windowSize,
+		}
+		imageResult, err = dExt.ImageService.GetImage(bufSource, options...)
+		if err != nil {
+			log.Error().Err(err).Msg("GetImage from ImageService failed")
+			continue
+		}
+		// success
+		break
+	}
 	if err != nil {
-		err = errors.Wrap(code.MobileUIDriverError, err.Error())
-		return
-	}
-	screenResult = &ScreenResult{
-		bufSource:  bufSource,
-		ImagePath:  imagePath,
-		Tags:       nil,
-		Resolution: windowSize,
+		return nil, err
 	}
+
 	// cache screen result
 	dExt.Driver.GetSession().addScreenResult(screenResult)
 
-	imageResult, err := dExt.ImageService.GetImage(bufSource, options...)
-	if err != nil {
-		log.Error().Err(err).Msg("GetImage from ImageService failed")
-		return screenResult, err
-	}
 	if imageResult != nil {
 		screenResult.ImageResult = imageResult
 		screenResult.Texts = imageResult.OCRResult.ToOCRTexts()