change: use context to contral ScreenRecord timeout or cancel

internal/version/VERSION

@@ -1 +1 @@
-v5.0.0-beta-2503061758
+v5.0.0-beta-2503062216

pkg/gadb/device.go

@@ -2,6 +2,7 @@ package gadb
 
 import (
 	"bytes"
+	"context"
 	"encoding/binary"
 	"fmt"
 	"io"
@@ -761,16 +762,38 @@ func (d *Device) ScreenCap() ([]byte, error) {
 	return buffer.Bytes(), nil
 }
 
-func (d *Device) ScreenRecord(seconds float64) ([]byte, error) {
+func (d *Device) ScreenRecord(ctx context.Context) ([]byte, error) {
 	videoPath := fmt.Sprintf("/sdcard/screenrecord_%d.mp4", time.Now().Unix())
-	_, err := d.RunShellCommandWithBytes("screenrecord",
-		"--time-limit", fmt.Sprintf("%.1f", seconds), videoPath)
-	if err != nil {
-		return nil, errors.Wrap(err, "screenrecord failed")
+
+	done := make(chan error, 1)
+	go func() {
+		_, err := d.RunShellCommandWithBytes("screenrecord", videoPath)
+		done <- err
+	}()
+
+	select {
+	case <-ctx.Done():
+		// timeout or cancelled
+		pid, err := d.RunShellCommand("pidof", "screenrecord")
+		if err == nil && pid != "" {
+			// 发送 SIGINT 信号终止录屏
+			_, _ = d.RunShellCommand("kill", "-2", strings.TrimSpace(pid))
+		}
+		<-done // 等待进程完全退出
+	case err := <-done:
+		// adb screenrecord will exit on reached 180s
+		if err != nil {
+			return nil, errors.Wrap(err, "screenrecord failed")
+		}
 	}
 
+	// remove temp file
+	defer func() {
+		go d.RunShellCommand("rm", videoPath)
+	}()
+
 	buffer := bytes.NewBuffer(nil)
-	err = d.Pull(videoPath, buffer)
+	err := d.Pull(videoPath, buffer)
 	if err != nil {
 		return nil, errors.Wrap(err, "pull video failed")
 	}

pkg/gadb/device_test.go

@@ -4,6 +4,7 @@ package gadb
 
 import (
 	"bytes"
+	"context"
 	"os"
 	"reflect"
 	"strings"
@@ -295,6 +296,36 @@ func TestDevice_Pull(t *testing.T) {
 	}
 }
 
+func TestDevice_ScreenRecord(t *testing.T) {
+	setupDevices(t)
+
+	for _, dev := range devices {
+		// screen record with time limit 5 seconds
+		ctx, cancel := context.WithTimeout(context.Background(), time.Second*5)
+		if _, err := dev.ScreenRecord(ctx); err != nil {
+			assert.Nil(t, err)
+		}
+		cancel()
+	}
+
+	for _, dev := range devices {
+		// screen record with cancel signal
+		ctx, cancel := context.WithCancel(context.Background())
+		done := make(chan error)
+		go func() {
+			_, err := dev.ScreenRecord(ctx)
+			done <- err
+		}()
+
+		// record for 3 seconds
+		time.Sleep(time.Second * 3)
+		cancel()
+
+		err := <-done
+		assert.Nil(t, err)
+	}
+}
+
 func TestDevice_RunShellCommandBackgroundWithBytes(t *testing.T) {
 	type fields struct {
 		adbClient Client

uixt/android_driver_adb.go

@@ -791,11 +791,25 @@ func (ad *ADBDriver) ScreenRecord(opts ...option.ActionOption) (videoPath string
 		filePath = filepath.Join(config.GetConfig().ScreenShotsPath, fmt.Sprintf("%s.mp4", timestamp))
 	}
 
+	var ctx context.Context
+	if options.Context != nil {
+		ctx = options.Context
+	} else {
+		ctx = context.Background()
+	}
+
+	var cancel context.CancelFunc
 	duration := options.ScreenRecordDuration
 	if duration == 0 {
 		duration = options.Duration
 	}
-	audioOn := options.ScreenRecordWithAudio
+	if duration != 0 {
+		ctx, cancel = context.WithTimeout(ctx,
+			time.Duration(duration*float64(time.Second)))
+	} else {
+		ctx, cancel = context.WithCancel(ctx)
+	}
+	defer cancel()
 
 	// get android system version
 	var sysVersion int
@@ -809,6 +823,7 @@ func (ad *ADBDriver) ScreenRecord(opts ...option.ActionOption) (videoPath string
 	}
 
 	var useAdbScreenRecord bool
+	audioOn := options.ScreenRecordWithAudio
 	if options.ScreenRecordWithScrcpy {
 		useAdbScreenRecord = false
 	} else if !audioOn {
@@ -834,7 +849,9 @@ func (ad *ADBDriver) ScreenRecord(opts ...option.ActionOption) (videoPath string
 	}()
 
 	if useAdbScreenRecord {
-		res, err := ad.Device.ScreenRecord(duration)
+		// screen record with adb screenrecord
+		// adb screenrecord duration is limited in range [1,180] seconds
+		res, err := ad.Device.ScreenRecord(ctx)
 		if err != nil {
 			return "", errors.Wrap(err, "screen record failed")
 		}
@@ -844,16 +861,8 @@ func (ad *ADBDriver) ScreenRecord(opts ...option.ActionOption) (videoPath string
 		return filePath, nil
 	}
 
-	// screen record with audio
-	log.Info().Float64("duration(s)", duration).Msg("screen record with audio, use scrcpy")
-
-	file, err := os.Create(filePath)
-	if err != nil {
-		return "", errors.Wrap(err, "create screen record file failed")
-	}
-	defer func() {
-		_ = file.Close()
-	}()
+	// screen record with scrcpy
+	log.Info().Float64("duration(s)", duration).Msg("screen record with scrcpy")
 
 	// start scrcpy
 	cmd := exec.Command(
@@ -866,15 +875,10 @@ func (ad *ADBDriver) ScreenRecord(opts ...option.ActionOption) (videoPath string
 	)
 	cmd.Stdout = io.Discard
 	cmd.Stderr = io.Discard
-
 	if err := cmd.Start(); err != nil {
 		return "", errors.Wrap(err, "start screen record failed")
 	}
 
-	ctx, cancel := context.WithTimeout(context.Background(),
-		time.Duration(duration*float64(time.Second)))
-	defer cancel()
-
 	done := make(chan error, 1)
 	go func() {
 		done <- cmd.Wait()
@@ -882,15 +886,17 @@ func (ad *ADBDriver) ScreenRecord(opts ...option.ActionOption) (videoPath string
 
 	select {
 	case <-ctx.Done():
-		// 超时，优雅停止 scrcpy 进程
+		// timeout or cancelled
+		log.Info().Msg("screen recording stopped")
 		if err := cmd.Process.Signal(syscall.SIGINT); err != nil {
 			log.Error().Err(err).Msg("failed to stop scrcpy process")
 			_ = cmd.Process.Kill() // 强制结束进程
 		}
 		<-done // 等待进程完全退出
 	case err := <-done:
+		log.Info().Msg("scrcpy exited")
 		if err != nil {
-			return "", errors.Wrap(err, "screen record failed")
+			return "", errors.Wrap(err, "screen record with scrcpy failed")
 		}
 	}
 

uixt/android_test.go

@@ -3,6 +3,7 @@
 package uixt
 
 import (
+	"context"
 	"os"
 	"testing"
 	"time"
@@ -195,12 +196,15 @@ func TestDriver_ADB_ForegroundInfo(t *testing.T) {
 
 func TestDriver_ADB_ScreenRecord(t *testing.T) {
 	driver := setupADBDriverExt(t)
+
+	// adb screenrecord --time-limit 5
 	path1, err := driver.ScreenRecord(
 		option.WithScreenRecordDuation(5))
 	assert.Nil(t, err)
 	defer os.Remove(path1)
 	t.Log(path1)
 
+	// scrcpy with time limit
 	path2, err := driver.ScreenRecord(
 		option.WithScreenRecordDuation(5),
 		option.WithScreenRecordAudio(true),
@@ -209,6 +213,7 @@ func TestDriver_ADB_ScreenRecord(t *testing.T) {
 	defer os.Remove(path2)
 	t.Log(path2)
 
+	// scrcpy with time limit
 	path3, err := driver.ScreenRecord(
 		option.WithScreenRecordDuation(5),
 		option.WithScreenRecordScrcpy(true),
@@ -216,6 +221,27 @@ func TestDriver_ADB_ScreenRecord(t *testing.T) {
 	assert.Nil(t, err)
 	defer os.Remove(path3)
 	t.Log(path3)
+
+	// scrcpy with cancel signal
+	ctx, cancel := context.WithCancel(context.Background())
+	done := make(chan error)
+	go func() {
+		path4, err := driver.ScreenRecord(
+			option.WithContext(ctx),
+			option.WithScreenRecordScrcpy(true),
+		)
+		assert.Nil(t, err)
+		defer os.Remove(path4)
+		t.Log(path4)
+		done <- err
+	}()
+
+	// record for 3 seconds
+	time.Sleep(time.Second * 3)
+	cancel()
+
+	err = <-done
+	assert.Nil(t, err)
 }
 
 func TestDriver_ADB_Backspace(t *testing.T) {

uixt/option/action.go

@@ -1,12 +1,14 @@
 package option
 
 import (
+	"context"
 	"math/rand/v2"
 
 	"github.com/httprunner/httprunner/v5/internal/builtin"
 )
 
 type ActionOptions struct {
+	Context context.Context
 	// log
 	Identifier string `json:"identifier,omitempty" yaml:"identifier,omitempty"` // used to identify the action in log
 
@@ -33,6 +35,9 @@ func (o *ActionOptions) Options() []ActionOption {
 		return options
 	}
 
+	if o.Context != nil {
+		options = append(options, WithContext(o.Context))
+	}
 	if o.Identifier != "" {
 		options = append(options, WithIdentifier(o.Identifier))
 	}
@@ -206,6 +211,12 @@ func NewActionOptions(opts ...ActionOption) *ActionOptions {
 
 type ActionOption func(o *ActionOptions)
 
+func WithContext(ctx context.Context) ActionOption {
+	return func(o *ActionOptions) {
+		o.Context = ctx
+	}
+}
+
 func WithCustomOption(key string, value interface{}) ActionOption {
 	return func(o *ActionOptions) {
 		if o.Custom == nil {

uixt/option/screen.go

@@ -1,6 +1,8 @@
 package option
 
-import "github.com/httprunner/httprunner/v5/uixt/types"
+import (
+	"github.com/httprunner/httprunner/v5/uixt/types"
+)
 
 type ScreenOptions struct {
 	ScreenShotOptions