From 5b0f158d29320fda86421c76fb8175152150e5c1 Mon Sep 17 00:00:00 2001
From: "lilong.129" <lilong.129@bytedance.com>
Date: Mon, 28 Aug 2023 23:16:05 +0800
Subject: [PATCH] feat: add live type

---
 hrp/pkg/uixt/ext.go           | 16 +++++++++-------
 hrp/pkg/uixt/service_vedem.go |  6 ++----
 hrp/pkg/uixt/video_crawler.go |  8 ++------
 3 files changed, 13 insertions(+), 17 deletions(-)

diff --git a/hrp/pkg/uixt/ext.go b/hrp/pkg/uixt/ext.go
index 5b26f157..2cd23504 100644
--- a/hrp/pkg/uixt/ext.go
+++ b/hrp/pkg/uixt/ext.go
@@ -52,9 +52,11 @@ func WithThreshold(threshold float64) CVOption {
 }
 
 type ScreenResult struct {
-	bufSource *bytes.Buffer // raw image buffer bytes
-	imagePath string        // image file path
+	bufSource   *bytes.Buffer // raw image buffer bytes
+	imagePath   string        // image file path
+	imageResult *ImageResult  // image result
 
+	LiveType    string
 	UploadedURL string      `json:"uploaded_url"`         // uploaded image url
 	Texts       OCRTexts    `json:"texts"`                // dumped raw OCRTexts
 	Icons       UIResultMap `json:"icons"`                // CV 识别的图标
@@ -74,16 +76,16 @@ type ScreenResult struct {
 	TotalElapsed int64 `json:"total_elapsed"` // current_swipe_finish -> next_swipe_start 整体耗时(ms)
 }
 
-type ScreenResultMap map[string]*ScreenResult
+type ScreenResultMap map[string]*ScreenResult // key is date time
 
 // getScreenShotUrls returns screenShotsUrls using imagePath as key and uploaded URL as value
 func (screenResults ScreenResultMap) getScreenShotUrls() map[string]string {
 	screenShotsUrls := make(map[string]string)
-	for imagePath, screenResult := range screenResults {
+	for dateTime, screenResult := range screenResults {
 		if screenResult.UploadedURL == "" {
 			continue
 		}
-		screenShotsUrls[imagePath] = screenResult.UploadedURL
+		screenShotsUrls[dateTime] = screenResult.UploadedURL
 	}
 	return screenShotsUrls
 }
@@ -282,7 +284,7 @@ func (dExt *DriverExt) GetStepCacheData() map[string]interface{} {
 		screenSize = Size{}
 	}
 	screenResults := make(map[string]interface{})
-	for imagePath, screenResult := range dExt.cacheStepData.screenResults {
+	for dateTime, screenResult := range dExt.cacheStepData.screenResults {
 		o, _ := json.Marshal(screenResult.Texts)
 		data := map[string]interface{}{
 			"tags":  screenResult.Tags,
@@ -303,7 +305,7 @@ func (dExt *DriverExt) GetStepCacheData() map[string]interface{} {
 			"popup":                   screenResult.Popup,
 		}
 
-		screenResults[imagePath] = data
+		screenResults[dateTime] = data
 	}
 	cacheData["screen_results"] = screenResults
 
diff --git a/hrp/pkg/uixt/service_vedem.go b/hrp/pkg/uixt/service_vedem.go
index 8b6f6ce1..fae52514 100644
--- a/hrp/pkg/uixt/service_vedem.go
+++ b/hrp/pkg/uixt/service_vedem.go
@@ -388,9 +388,7 @@ func (dExt *DriverExt) GetScreenResult(options ...ActionOption) (screenResult *S
 		screenResult.Icons = imageResult.UIResult
 
 		if imageResult.LiveType != "" && imageResult.LiveType != "NoLive" {
-			screenResult.Live = &LiveRoom{
-				LiveType: imageResult.LiveType,
-			}
+			screenResult.LiveType = imageResult.LiveType
 		}
 		if actionOptions.ScreenShotWithClosePopups {
 			screenResult.Popup = &PopupInfo{
@@ -405,7 +403,7 @@ func (dExt *DriverExt) GetScreenResult(options ...ActionOption) (screenResult *S
 
 	}
 
-	dExt.cacheStepData.screenResults[imagePath] = screenResult
+	dExt.cacheStepData.screenResults[time.Now().String()] = screenResult
 
 	log.Debug().
 		Str("imagePath", imagePath).
diff --git a/hrp/pkg/uixt/video_crawler.go b/hrp/pkg/uixt/video_crawler.go
index 77e90a61..640928a7 100644
--- a/hrp/pkg/uixt/video_crawler.go
+++ b/hrp/pkg/uixt/video_crawler.go
@@ -178,9 +178,6 @@ func (vc *VideoCrawler) startLiveCrawler(enterPoint PointF) error {
 			}
 			swipeFinishTime := time.Now()
 
-			// wait for live video loading
-			time.Sleep(5 * time.Second)
-
 			liveRoom, err := vc.getCurrentLiveRoom()
 			if err != nil {
 				return errors.Wrap(err, "get current live event trackings failed")
@@ -188,16 +185,15 @@ func (vc *VideoCrawler) startLiveCrawler(enterPoint PointF) error {
 
 			// take screenshot and get screen texts by OCR
 			screenResult, err := vc.driverExt.GetScreenResult(
-				WithScreenShotOCR(true), WithScreenShotUpload(true))
+				WithScreenShotOCR(true), WithScreenShotUpload(true), WithScreenShotLiveType(true))
 			if err != nil {
 				log.Error().Err(err).Msg("OCR GetTexts failed")
 				time.Sleep(3 * time.Second)
 				continue
 			}
+			liveRoom.LiveType = screenResult.LiveType
 			screenResult.Live = liveRoom
 
-			// TODO: check live type
-
 			// incr live count
 			screenResult.VideoType = "live"
 			vc.LiveCount++