refactor: save screenshot in async

This commit is contained in:
lilong.129 2025-03-05 16:29:36 +08:00
parent 98d1ecb478
commit e187ba824a
7 changed files with 49 additions and 90 deletions

View File

@ -1 +1 @@
v5.0.0+2503051452
v5.0.0+2503051629

View File

@ -544,18 +544,6 @@ func (ad *ADBDriver) ScreenShot(opts ...option.ActionOption) (raw *bytes.Buffer,
"adb screencap failed %v", err)
}
raw = bytes.NewBuffer(resp)
actionOptions := option.NewActionOptions(opts...)
if actionOptions.ScreenShotFileName != "" {
// save screenshot to file
path, err := saveScreenShot(raw, actionOptions.ScreenShotFileName)
if err != nil {
return nil, errors.Wrapf(code.DeviceScreenShotError,
"save screenshot file failed %v", err)
}
log.Info().Str("path", path).Msg("screenshot saved")
}
return raw, nil
}

View File

@ -307,9 +307,10 @@ func (wd *BrowserDriver) ScreenShot(options ...option.ActionOption) (*bytes.Buff
data := resp.Data.(map[string]interface{})
screenshotBase64 := data["screenshot"].(string)
screenRaw, err := base64.StdEncoding.DecodeString(screenshotBase64)
res := bytes.NewBuffer(screenRaw)
return res, err
if err != nil {
return nil, err
}
return bytes.NewBuffer(screenRaw), nil
}
func (wd *BrowserDriver) HttpPOST(data interface{}, pathElem ...string) (response *WebAgentResponse, err error) {

View File

@ -49,6 +49,21 @@ func (s *ScreenResult) FilterTextsByScope(x1, y1, x2, y2 float64) ai.OCRTexts {
func (dExt *XTDriver) GetScreenResult(opts ...option.ActionOption) (screenResult *ScreenResult, err error) {
screenshotOptions := option.NewActionOptions(opts...)
// 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)
}
// save compressed screenshot to file
var fileName string
optionsList := screenshotOptions.List()
if screenshotOptions.ScreenShotFileName != "" {
@ -58,57 +73,33 @@ func (dExt *XTDriver) GetScreenResult(opts ...option.ActionOption) (screenResult
} else {
fileName = builtin.GenNameWithTimestamp("%d_screenshot")
}
var bufSource *bytes.Buffer
var compressBufSource *bytes.Buffer
var imageResult *ai.CVResult
var imagePath string
var windowSize types.Size
var lastErr error
// get screenshot info with retry
for i := 0; i < 3; i++ {
imagePath = filepath.Join(config.GetConfig().ScreenShotsPath, fileName)
bufSource, err = dExt.ScreenShot(option.WithScreenShotFileName(imagePath))
imagePath := filepath.Join(config.GetConfig().ScreenShotsPath, fileName)
go func() {
path, err := saveScreenShot(compressBufSource, imagePath)
if err != nil {
lastErr = err
continue
}
compressBufSource, err = compressImageBuffer(bufSource)
if err != nil {
lastErr = err
continue
log.Error().Err(err).Msg("save screenshot file failed")
} else {
log.Info().Str("path", path).Msg("screenshot saved")
}
}()
windowSize, err = dExt.WindowSize()
if err != nil {
lastErr = errors.Wrap(code.DeviceGetInfoError, err.Error())
continue
}
screenResult = &ScreenResult{
bufSource: compressBufSource,
ImagePath: imagePath,
Tags: nil,
Resolution: windowSize,
}
imageResult, err = dExt.CVService.ReadFromBuffer(compressBufSource, opts...)
if err != nil {
log.Error().Err(err).Msg("ReadFromBuffer from ImageService failed")
lastErr = err
continue
}
// success, break the loop
lastErr = nil
break
}
if lastErr != nil {
return nil, lastErr
windowSize, err := dExt.WindowSize()
if err != nil {
return nil, errors.Wrap(code.DeviceGetInfoError, err.Error())
}
// cache screen result
dExt.screenResults = append(dExt.screenResults, screenResult)
// read image from buffer with CV
screenResult = &ScreenResult{
bufSource: compressBufSource,
ImagePath: imagePath,
Tags: nil,
Resolution: windowSize,
}
imageResult, err := dExt.CVService.ReadFromBuffer(compressBufSource, opts...)
if err != nil {
log.Error().Err(err).Msg("ReadFromBuffer from ImageService failed")
return nil, err
}
if imageResult != nil {
screenResult.Texts = imageResult.OCRResult.ToOCRTexts()
screenResult.UploadedURL = imageResult.URL
@ -128,6 +119,9 @@ func (dExt *XTDriver) GetScreenResult(opts ...option.ActionOption) (screenResult
}
}
// cache screen result
dExt.screenResults = append(dExt.screenResults, screenResult)
log.Debug().
Str("imagePath", imagePath).
Str("imageUrl", screenResult.UploadedURL).

View File

@ -42,7 +42,7 @@ func NewDriverSession() *DriverSession {
timeout := 30 * time.Second
session := &DriverSession{
ctx: context.Background(),
ID: "",
ID: "<SessionNotInit>",
timeout: timeout,
client: &http.Client{
Timeout: timeout,

View File

@ -237,18 +237,6 @@ func (hd *HDCDriver) ScreenShot(opts ...option.ActionOption) (*bytes.Buffer, err
return nil, err
}
rawBuffer := bytes.NewBuffer(raw)
actionOptions := option.NewActionOptions(opts...)
if actionOptions.ScreenShotFileName != "" {
// save screenshot to file
path, err := saveScreenShot(rawBuffer, actionOptions.ScreenShotFileName)
if err != nil {
return nil, errors.Wrapf(code.DeviceScreenShotError,
"save screenshot file failed %v", err)
}
log.Info().Str("path", path).Msg("screenshot saved")
}
return rawBuffer, nil
}

View File

@ -334,19 +334,7 @@ func (wd *WDADriver) ScreenShot(opts ...option.ActionOption) (raw *bytes.Buffer,
return nil, errors.Wrap(code.DeviceScreenShotError,
fmt.Sprintf("decode WDA screenshot data failed: %v", err))
}
actionOptions := option.NewActionOptions(opts...)
if actionOptions.ScreenShotFileName != "" {
// save screenshot to file
path, err := saveScreenShot(raw, actionOptions.ScreenShotFileName)
if err != nil {
return nil, errors.Wrapf(code.DeviceScreenShotError,
"save screenshot file failed %v", err)
}
log.Info().Str("path", path).Msg("screenshot saved")
}
return
return raw, nil
}
func (wd *WDADriver) toScale(x float64) float64 {
@ -614,13 +602,13 @@ func (wd *WDADriver) TapAbsXY(x, y float64, opts ...option.ActionOption) error {
func (wd *WDADriver) DoubleTap(x, y float64, opts ...option.ActionOption) error {
// [[FBRoute POST:@"/wda/doubleTap"] respondWithTarget:self action:@selector(handleDoubleTapCoordinate:)]
var err error
actionOptions := option.NewActionOptions(opts...)
x, y, err = convertToAbsolutePoint(wd, x, y)
if err != nil {
return err
}
x, y = actionOptions.ApplyOffset(x, y)
actionOptions := option.NewActionOptions(opts...)
x, y = actionOptions.ApplyOffset(x, y)
x = wd.toScale(x)
y = wd.toScale(y)
data := map[string]interface{}{