refactor: TakeScreenShot saves all screenshots

This commit is contained in:
lilong.129
2023-04-16 22:11:35 +08:00
parent 3770ce12bd
commit 9e73328448
9 changed files with 50 additions and 72 deletions

View File

@@ -217,7 +217,7 @@ type DriverExt struct {
doneMjpegStream chan bool
scale float64
ocrService OCRService // used to get text from image
ScreenShots []string // save screenshots path
screenShots []string // cache screenshot paths
CVArgs
}
@@ -253,7 +253,9 @@ func NewDriverExt(device Device, driver WebDriver) (dExt *DriverExt, err error)
return dExt, nil
}
func (dExt *DriverExt) takeScreenShot() (raw *bytes.Buffer, err error) {
// 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(fileName ...string) (raw *bytes.Buffer, err error) {
// wait for action done
time.Sleep(500 * time.Millisecond)
@@ -263,14 +265,27 @@ func (dExt *DriverExt) takeScreenShot() (raw *bytes.Buffer, err error) {
return dExt.frame, nil
}
if raw, err = dExt.Driver.Screenshot(); err != nil {
log.Error().Err(err).Msg("takeScreenShot failed")
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 saveScreenShot(raw *bytes.Buffer, fileName string) (string, error) {
func (dExt *DriverExt) saveScreenShot(raw *bytes.Buffer, fileName string) (string, error) {
img, format, err := image.Decode(raw)
if err != nil {
return "", errors.Wrap(err, "decode screenshot image failed")
@@ -302,19 +317,11 @@ func saveScreenShot(raw *bytes.Buffer, fileName string) (string, error) {
return screenshotPath, nil
}
// ScreenShot takes screenshot and saves image file to $CWD/screenshots/ folder
func (dExt *DriverExt) ScreenShot(fileName string) (string, error) {
raw, err := dExt.takeScreenShot()
if err != nil {
return "", errors.Wrap(err, "screenshot failed")
}
fileName = filepath.Join(env.ScreenShotsPath, fileName)
path, err := saveScreenShot(raw, fileName)
if err != nil {
return "", errors.Wrap(err, "save screenshot failed")
}
return path, nil
func (dExt *DriverExt) GetScreenShots() []string {
defer func() {
dExt.screenShots = nil
}()
return dExt.screenShots
}
// isPathExists returns true if path exists, whether path is file or dir
@@ -689,15 +696,9 @@ func (dExt *DriverExt) DoAction(action MobileAction) error {
}
}
case CtlScreenShot:
// take snapshot
log.Info().Msg("take snapshot for current screen")
screenshotPath, err := dExt.ScreenShot(fmt.Sprintf("screenshot_%d",
time.Now().Unix()))
if err != nil {
return errors.Wrap(err, "take screenshot failed")
}
log.Info().Str("path", screenshotPath).Msg("take screenshot")
dExt.ScreenShots = append(dExt.ScreenShots, screenshotPath)
// take screenshot
log.Info().Msg("take screenshot for current screen")
_, err := dExt.TakeScreenShot(builtin.GenNameWithTimestamp("screenshot_"))
return err
case CtlStartCamera:
return dExt.Driver.StartCamera()

View File

@@ -2,7 +2,6 @@ package uixt
import (
"bytes"
"fmt"
"math"
"strings"
"time"
@@ -437,7 +436,6 @@ type DataOptions struct {
IgnoreNotFoundError bool // ignore error if target element not found
MaxRetryTimes int // max retry times if target element not found
Interval float64 // interval between retries in seconds
ScreenShotFilename string // turn on screenshot and specify file name
}
type DataOption func(data *DataOptions)
@@ -514,16 +512,6 @@ func WithDataWaitTime(sec float64) DataOption {
}
}
func WithScreenShot(fileName ...string) DataOption {
return func(data *DataOptions) {
if len(fileName) > 0 {
data.ScreenShotFilename = fileName[0]
} else {
data.ScreenShotFilename = fmt.Sprintf("screenshot_%d", time.Now().Unix())
}
}
}
func NewDataOptions(options ...DataOption) *DataOptions {
dataOptions := &DataOptions{
Data: make(map[string]interface{}),

View File

@@ -175,14 +175,6 @@ func (s *veDEMOCRService) GetTexts(imageBuf *bytes.Buffer, options ...DataOption
dataOptions := NewDataOptions(options...)
if dataOptions.ScreenShotFilename != "" {
path, err := saveScreenShot(imageBuf, dataOptions.ScreenShotFilename)
if err != nil {
return nil, errors.Wrap(err, "save screenshot failed")
}
log.Debug().Str("path", path).Msg("save screenshot")
}
for _, ocrResult := range ocrResults {
rect := image.Rectangle{
// ocrResult.Points 顺序:左上 -> 右上 -> 右下 -> 左下
@@ -313,8 +305,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 {
err = fmt.Errorf("takeScreenShot error: %v", err)
if bufSource, err = dExt.TakeScreenShot(builtin.GenNameWithTimestamp("ocr_")); err != nil {
return
}
@@ -329,8 +320,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 {
err = fmt.Errorf("takeScreenShot error: %v", err)
if bufSource, err = dExt.TakeScreenShot(builtin.GenNameWithTimestamp("ocr_")); err != nil {
return
}
@@ -348,8 +338,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 {
err = fmt.Errorf("takeScreenShot error: %v", err)
if bufSource, err = dExt.TakeScreenShot(builtin.GenNameWithTimestamp("ocr_")); err != nil {
return
}

View File

@@ -14,6 +14,8 @@ import (
"github.com/pkg/errors"
"gocv.io/x/gocv"
"github.com/httprunner/httprunner/v4/hrp/internal/builtin"
)
const (
@@ -101,7 +103,7 @@ func (dExt *DriverExt) FindAllImageRect(search string) (rects []image.Rectangle,
if bufSearch, err = getBufFromDisk(search); err != nil {
return nil, err
}
if bufSource, err = dExt.takeScreenShot(); err != nil {
if bufSource, err = dExt.TakeScreenShot(builtin.GenNameWithTimestamp("cv_")); err != nil {
return nil, err
}
@@ -116,7 +118,7 @@ func (dExt *DriverExt) FindImageRectInUIKit(imagePath string, options ...DataOpt
if bufSearch, err = getBufFromDisk(imagePath); err != nil {
return 0, 0, 0, 0, err
}
if bufSource, err = dExt.takeScreenShot(); err != nil {
if bufSource, err = dExt.TakeScreenShot(builtin.GenNameWithTimestamp("cv_")); err != nil {
return 0, 0, 0, 0, err
}