From 267d53500724286185277d047fb4c7577f935ffd Mon Sep 17 00:00:00 2001
From: "lilong.129" <lilong.129@bytedance.com>
Date: Wed, 23 Aug 2023 17:27:33 +0800
Subject: [PATCH] refactor: pass ui type params to WithScreenShotUITypes

---
 hrp/pkg/uixt/action.go        | 21 +++++++++++----------
 hrp/pkg/uixt/service_vedem.go | 11 +++++++----
 2 files changed, 18 insertions(+), 14 deletions(-)

diff --git a/hrp/pkg/uixt/action.go b/hrp/pkg/uixt/action.go
index 280949a4..6ccb0b5e 100644
--- a/hrp/pkg/uixt/action.go
+++ b/hrp/pkg/uixt/action.go
@@ -113,10 +113,10 @@ type ActionOptions struct {
 	Custom map[string]interface{} `json:"custom,omitempty" yaml:"custom,omitempty"`
 
 	// screenshot related
-	ScreenShotWithOCR      bool `json:"screenshot_with_ocr,omitempty" yaml:"screenshot_with_ocr,omitempty"`
-	ScreenShotWithUpload   bool `json:"screenshot_with_upload,omitempty" yaml:"screenshot_with_upload,omitempty"`
-	ScreenShotWithLiveType bool `json:"screenshot_with_live_type,omitempty" yaml:"screenshot_with_live_type,omitempty"`
-	ScreenShotWithUIType   bool `json:"screenshot_with_ui_type,omitempty" yaml:"screenshot_with_ui_type,omitempty"`
+	ScreenShotWithOCR      bool     `json:"screenshot_with_ocr,omitempty" yaml:"screenshot_with_ocr,omitempty"`
+	ScreenShotWithUpload   bool     `json:"screenshot_with_upload,omitempty" yaml:"screenshot_with_upload,omitempty"`
+	ScreenShotWithLiveType bool     `json:"screenshot_with_live_type,omitempty" yaml:"screenshot_with_live_type,omitempty"`
+	ScreenShotWithUITypes  []string `json:"screenshot_with_ui_types,omitempty" yaml:"screenshot_with_ui_types,omitempty"`
 }
 
 func (o *ActionOptions) Options() []ActionOption {
@@ -201,14 +201,14 @@ func (o *ActionOptions) Options() []ActionOption {
 	if o.ScreenShotWithLiveType {
 		options = append(options, WithScreenShotLiveType(true))
 	}
-	if o.ScreenShotWithUIType {
-		options = append(options, WithScreenShotUIType(true))
+	if len(o.ScreenShotWithUITypes) > 0 {
+		options = append(options, WithScreenShotUITypes(o.ScreenShotWithUITypes...))
 	}
 
 	return options
 }
 
-func (o *ActionOptions) screenshotOptions() screenshotActionOptions {
+func (o *ActionOptions) screenshotActionOptions() screenshotActionOptions {
 	options := screenshotActionOptions{}
 	if o.ScreenShotWithOCR {
 		options = append(options, "ocr")
@@ -219,7 +219,8 @@ func (o *ActionOptions) screenshotOptions() screenshotActionOptions {
 	if o.ScreenShotWithLiveType {
 		options = append(options, "liveType")
 	}
-	if o.ScreenShotWithUIType {
+	// UI detection
+	if len(o.ScreenShotWithUITypes) > 0 {
 		options = append(options, "ui")
 	}
 	return options
@@ -418,9 +419,9 @@ func WithScreenShotLiveType(liveTypeOn bool) ActionOption {
 	}
 }
 
-func WithScreenShotUIType(uiTypeOn bool) ActionOption {
+func WithScreenShotUITypes(uiTypes ...string) ActionOption {
 	return func(o *ActionOptions) {
-		o.ScreenShotWithUIType = uiTypeOn
+		o.ScreenShotWithUITypes = uiTypes
 	}
 }
 
diff --git a/hrp/pkg/uixt/service_vedem.go b/hrp/pkg/uixt/service_vedem.go
index 535f87a5..eaeefe53 100644
--- a/hrp/pkg/uixt/service_vedem.go
+++ b/hrp/pkg/uixt/service_vedem.go
@@ -342,10 +342,12 @@ type IImageService interface {
 
 // GetScreenResult takes a screenshot, returns the image recognization result
 func (dExt *DriverExt) GetScreenResult(options ...ActionOption) (screenResult *ScreenResult, err error) {
-	screenshotOptions := NewActionOptions(options...).screenshotOptions()
+	actionOptions := NewActionOptions(options...)
+	screenActionOptions := actionOptions.screenshotActionOptions()
+	screenUITypeOptions := screenshotUITypeOptions(actionOptions.ScreenShotWithUITypes)
 
 	var fileName string
-	if len(screenshotOptions) == 0 {
+	if len(screenActionOptions) == 0 {
 		fileName = builtin.GenNameWithTimestamp("%d_screenshot")
 	} else {
 		fileName = builtin.GenNameWithTimestamp("%d_cv")
@@ -366,8 +368,9 @@ func (dExt *DriverExt) GetScreenResult(options ...ActionOption) (screenResult *S
 	}
 
 	var imageUrl string
-	if len(screenshotOptions) > 0 {
-		imageResult, err := dExt.ImageService.GetImage(bufSource, screenshotOptions)
+	if len(screenActionOptions) > 0 {
+		imageResult, err := dExt.ImageService.GetImage(bufSource,
+			screenActionOptions, screenUITypeOptions)
 		if err != nil {
 			log.Error().Err(err).Msg("GetImage from ImageService failed")
 			return nil, err