fix: setup device and driver

internal/version/VERSION

@@ -1 +1 @@
-v5.0.0+2502181539
+v5.0.0+2502181804

pkg/uixt/android_device.go

@@ -84,6 +84,11 @@ func NewAndroidDevice(opts ...option.AndroidDeviceOption) (device *AndroidDevice
 		Logcat:  NewAdbLogcat(androidOptions.SerialNumber),
 	}
 	log.Info().Str("serial", device.Options.SerialNumber).Msg("init android device")
+
+	// setup device
+	if err := device.Setup(); err != nil {
+		return nil, errors.Wrap(err, "setup android device failed")
+	}
 	return device, nil
 }
 
@@ -107,32 +112,6 @@ func (dev *AndroidDevice) Setup() error {
 		return errors.Wrap(code.DeviceShellExecError, err.Error())
 	}
 
-	if dev.Options.UIA2 {
-		// uiautomator2 server must be started before
-
-		// check uiautomator server package installed
-		if !dev.Device.IsPackageInstalled(dev.Options.UIA2ServerPackageName) {
-			return errors.Wrapf(code.MobileUIDriverAppNotInstalled,
-				"%s not installed", dev.Options.UIA2ServerPackageName)
-		}
-		if !dev.Device.IsPackageInstalled(dev.Options.UIA2ServerTestPackageName) {
-			return errors.Wrapf(code.MobileUIDriverAppNotInstalled,
-				"%s not installed", dev.Options.UIA2ServerTestPackageName)
-		}
-
-		// TODO: check uiautomator server package running
-		// if dev.IsPackageRunning(UIA2ServerPackageName) {
-		// 	return nil
-		// }
-
-		// start uiautomator2 server
-		go func() {
-			if err := dev.startUIA2Server(); err != nil {
-				log.Error().Err(err).Msg("start UIA2 failed")
-			}
-		}()
-		time.Sleep(5 * time.Second) // wait for uiautomator2 server start
-	}
 	return nil
 }
 
@@ -164,10 +143,6 @@ func (dev *AndroidDevice) NewDriver() (driver IDriver, err error) {
 			return nil, err
 		}
 	}
-	// setup driver
-	if err := driver.Setup(); err != nil {
-		return nil, err
-	}
 	return driver, nil
 }
 
@@ -404,33 +379,6 @@ func (dev *AndroidDevice) getPackageMD5(packagePath string) (string, error) {
 	return "", errors.New("failed to get package md5")
 }
 
-func (dev *AndroidDevice) startUIA2Server() error {
-	const maxRetries = 3
-	for attempt := 1; attempt <= maxRetries; attempt++ {
-		log.Info().Str("package", dev.Options.UIA2ServerTestPackageName).
-			Int("attempt", attempt).Msg("start uiautomator server")
-		// $ adb shell am instrument -w $UIA2ServerTestPackageName
-		// -w: wait for instrumentation to finish before returning.
-		// Required for test runners.
-		out, err := dev.Device.RunShellCommand("am", "instrument", "-w",
-			dev.Options.UIA2ServerTestPackageName)
-		if err != nil {
-			return errors.Wrap(err, "start uiautomator server failed")
-		}
-		if strings.Contains(out, "Process crashed") {
-			log.Error().Msg("uiautomator server crashed, retrying...")
-		}
-	}
-
-	return errors.Wrapf(code.MobileUIDriverAppCrashed,
-		"uiautomator server crashed %d times", maxRetries)
-}
-
-func (dev *AndroidDevice) stopUIA2Server() error {
-	_, err := dev.Device.RunShellCommand("am", "force-stop", dev.Options.UIA2ServerPackageName)
-	return err
-}
-
 type LineCallback func(string)
 
 type AdbLogcat struct {

pkg/uixt/android_driver_adb.go

@@ -33,6 +33,10 @@ func NewADBDriver(device *AndroidDevice) (*ADBDriver, error) {
 		Device:  device,
 		Session: NewDriverSession(),
 	}
+	// setup driver
+	if err := driver.Setup(); err != nil {
+		return nil, err
+	}
 	return driver, nil
 }
 

pkg/uixt/android_driver_uia2.go

@@ -5,6 +5,7 @@ import (
 	"encoding/base64"
 	"encoding/json"
 	"fmt"
+	"strings"
 
 	"github.com/pkg/errors"
 	"github.com/rs/zerolog/log"
@@ -24,20 +25,9 @@ func NewUIA2Driver(device *AndroidDevice) (*UIA2Driver, error) {
 	driver := &UIA2Driver{
 		ADBDriver: adbDriver,
 	}
-	err = driver.InitSession(nil)
-	if err != nil {
-		return nil, err
-	}
 
-	// forward port
+	// setup driver
-	localPort, err := device.Device.Forward(device.Options.UIA2Port)
+	if err := driver.Setup(); err != nil {
-	if err != nil {
-		return nil, errors.Wrap(code.DeviceConnectionError,
-			fmt.Sprintf("forward port %d->%d failed: %v",
-				localPort, device.Options.UIA2Port, err))
-	}
-	err = driver.Session.InitConnection(localPort)
-	if err != nil {
 		return nil, err
 	}
 	return driver, nil
@@ -50,6 +40,59 @@ type UIA2Driver struct {
 	windowSize types.Size
 }
 
+func (ud *UIA2Driver) Setup() error {
+	localPort, err := ud.Device.Forward(ud.Device.Options.UIA2Port)
+	if err != nil {
+		return errors.Wrap(code.DeviceConnectionError,
+			fmt.Sprintf("forward port %d->%d failed: %v",
+				localPort, ud.Device.Options.UIA2Port, err))
+	}
+	err = ud.Session.SetupPortForward(localPort)
+	if err != nil {
+		return err
+	}
+	ud.Session.SetBaseURL(
+		fmt.Sprintf("http://forward-to-%d:%d/wd/hub",
+			localPort, ud.Device.Options.UIA2Port))
+
+	// uiautomator2 server must be started before
+
+	// check uiautomator server package installed
+	if !ud.Device.IsPackageInstalled(ud.Device.Options.UIA2ServerPackageName) {
+		return errors.Wrapf(code.MobileUIDriverAppNotInstalled,
+			"%s not installed", ud.Device.Options.UIA2ServerPackageName)
+	}
+	if !ud.Device.IsPackageInstalled(ud.Device.Options.UIA2ServerTestPackageName) {
+		return errors.Wrapf(code.MobileUIDriverAppNotInstalled,
+			"%s not installed", ud.Device.Options.UIA2ServerTestPackageName)
+	}
+
+	// TODO: check uiautomator server package running
+	// if dev.IsPackageRunning(UIA2ServerPackageName) {
+	// 	return nil
+	// }
+
+	// start uiautomator2 server
+	// go func() {
+	// 	if err := ud.startUIA2Server(); err != nil {
+	// 		log.Fatal().Err(err).Msg("start UIA2 failed")
+	// 	}
+	// }()
+	// time.Sleep(5 * time.Second) // wait for uiautomator2 server start
+
+	// create new session
+	err = ud.InitSession(nil)
+	if err != nil {
+		return err
+	}
+	return nil
+}
+
+func (ud *UIA2Driver) TearDown() error {
+	log.Warn().Msg("TearDown not implemented in UIA2Driver")
+	return nil
+}
+
 func (ud *UIA2Driver) InitSession(capabilities option.Capabilities) (err error) {
 	// register(postHandler, new InitSession("/wd/hub/session"))
 	var rawResp DriverRawResponse
@@ -497,3 +540,31 @@ func (ud *UIA2Driver) Source(srcOpt ...option.SourceOption) (source string, err
 	source = reply.Value
 	return
 }
+
+func (ud *UIA2Driver) startUIA2Server() error {
+	const maxRetries = 3
+	for attempt := 1; attempt <= maxRetries; attempt++ {
+		log.Info().Str("package", ud.Device.Options.UIA2ServerTestPackageName).
+			Int("attempt", attempt).Msg("start uiautomator server")
+		// $ adb shell am instrument -w $UIA2ServerTestPackageName
+		// -w: wait for instrumentation to finish before returning.
+		// Required for test runners.
+		out, err := ud.Device.RunShellCommand("am", "instrument", "-w",
+			ud.Device.Options.UIA2ServerTestPackageName)
+		if err != nil {
+			return errors.Wrap(err, "start uiautomator server failed")
+		}
+		if strings.Contains(out, "Process crashed") {
+			log.Error().Msg("uiautomator server crashed, retrying...")
+		}
+	}
+
+	return errors.Wrapf(code.MobileUIDriverAppCrashed,
+		"uiautomator server crashed %d times", maxRetries)
+}
+
+func (ud *UIA2Driver) stopUIA2Server() error {
+	_, err := ud.Device.RunShellCommand("am", "force-stop",
+		ud.Device.Options.UIA2ServerPackageName)
+	return err
+}

pkg/uixt/driver_session.go

@@ -270,24 +270,17 @@ func (s *DriverSession) RequestWithRetry(method string, rawURL string, rawBody [
 	return
 }
 
-func (s *DriverSession) InitConnection(localPort int) error {
+func (s *DriverSession) SetupPortForward(localPort int) error {
 	conn, err := net.Dial("tcp", fmt.Sprintf("127.0.0.1:%d", localPort))
 	if err != nil {
 		return fmt.Errorf("create tcp connection error %v", err)
 	}
-	s.client = NewHTTPClientWithConnection(conn, s.timeout)
+	s.client.Transport = &http.Transport{
-	return nil
+		DialContext: func(_ context.Context, _, _ string) (net.Conn, error) {
-}
+			return conn, nil
-
-func NewHTTPClientWithConnection(conn net.Conn, timeout time.Duration) *http.Client {
-	return &http.Client{
-		Transport: &http.Transport{
-			DialContext: func(_ context.Context, _, _ string) (net.Conn, error) {
-				return conn, nil
-			},
 		},
-		Timeout: timeout,
 	}
+	return nil
 }
 
 type DriverRawResponse []byte

pkg/uixt/harmony_device.go

@@ -61,6 +61,11 @@ func NewHarmonyDevice(opts ...option.HarmonyDeviceOption) (device *HarmonyDevice
 		Device:  harmonyDevice,
 	}
 	log.Info().Str("connectKey", device.Options.ConnectKey).Msg("init harmony device")
+
+	// setup device
+	if err := device.Setup(); err != nil {
+		return nil, errors.Wrap(err, "setup harmony device failed")
+	}
 	return device, nil
 }
 
@@ -99,9 +104,5 @@ func (dev *HarmonyDevice) NewDriver() (IDriver, error) {
 	if err != nil {
 		return nil, errors.Wrap(err, "init harmony driver failed")
 	}
-	// setup driver
-	if err := driver.Setup(); err != nil {
-		return nil, err
-	}
 	return driver, nil
 }

pkg/uixt/harmony_driver_hdc.go

@@ -29,6 +29,11 @@ func NewHDCDriver(device *HarmonyDevice) (*HDCDriver, error) {
 	}
 	driver.uiDriver = uiDriver
 
+	// setup driver
+	if err := driver.Setup(); err != nil {
+		return nil, err
+	}
+
 	return driver, nil
 }
 

pkg/uixt/ios_device.go

@@ -107,6 +107,11 @@ func NewIOSDevice(opts ...option.IOSDeviceOption) (device *IOSDevice, err error)
 		listeners:   make(map[int]*forward.ConnListener),
 	}
 	log.Info().Str("udid", device.Options.UDID).Msg("init ios device")
+
+	// setup device
+	if err := device.Setup(); err != nil {
+		return nil, errors.Wrap(err, "setup ios device failed")
+	}
 	return device, nil
 }
 
@@ -224,10 +229,6 @@ func (dev *IOSDevice) NewDriver() (driver IDriver, err error) {
 			return nil, err
 		}
 	}
-	// setup driver
-	if err := wdaDriver.Setup(); err != nil {
-		return nil, err
-	}
 	return wdaDriver, nil
 }
 

pkg/uixt/ios_driver_wda.go

@@ -2,6 +2,7 @@ package uixt
 
 import (
 	"bytes"
+	"context"
 	"encoding/base64"
 	builtinJSON "encoding/json"
 	"fmt"
@@ -35,24 +36,8 @@ func NewWDADriver(device *IOSDevice) (*WDADriver, error) {
 		Session: NewDriverSession(),
 	}
 
-	host := "localhost"
+	// setup driver
-	localPort, err := driver.getLocalPort()
+	if err := driver.Setup(); err != nil {
-	if err != nil {
-		return nil, err
-	}
-	driver.Session.SetBaseURL(fmt.Sprintf("http://%s:%d", host, localPort))
-
-	if err = driver.initMjpegClient(); err != nil {
-		return nil, err
-	}
-
-	// create new session
-	if err = driver.InitSession(nil); err != nil {
-		return nil, errors.Wrap(code.DeviceHTTPDriverError, err.Error())
-	}
-
-	// init WDA scale
-	if driver.scale, err = driver.Scale(); err != nil {
 		return nil, err
 	}
 
@@ -67,9 +52,8 @@ type WDADriver struct {
 	windowSize types.Size
 	scale      float64
 
-	mjpegHTTPConn net.Conn // via HTTP
+	mjpegClient *http.Client
-	mjpegClient   *http.Client
+	mjpegUrl    string
-	mjpegUrl      string
 }
 
 func (wd *WDADriver) getLocalPort() (int, error) {
@@ -116,13 +100,21 @@ func (wd *WDADriver) initMjpegClient() error {
 	if err != nil {
 		return err
 	}
-	if wd.mjpegHTTPConn, err = net.Dial(
+	mjpegHTTPConn, err := net.Dial(
 		"tcp",
 		fmt.Sprintf("%s:%d", host, localMjpegPort),
-	); err != nil {
+	)
+	if err != nil {
 		return errors.Wrap(code.DeviceHTTPDriverError, err.Error())
 	}
-	wd.mjpegClient = NewHTTPClientWithConnection(wd.mjpegHTTPConn, 30*time.Second)
+	wd.mjpegClient = &http.Client{
+		Transport: &http.Transport{
+			DialContext: func(_ context.Context, _, _ string) (net.Conn, error) {
+				return mjpegHTTPConn, nil
+			},
+		},
+		Timeout: 30 * time.Second,
+	}
 	wd.mjpegUrl = fmt.Sprintf("http://%s:%d", host, localMjpegPort)
 	return nil
 }
@@ -131,6 +123,37 @@ func (wd *WDADriver) GetMjpegClient() *http.Client {
 	return wd.mjpegClient
 }
 
+func (wd *WDADriver) Setup() error {
+	localPort, err := wd.getLocalPort()
+	if err != nil {
+		return err
+	}
+	err = wd.Session.SetupPortForward(localPort)
+	if err != nil {
+		return err
+	}
+	wd.Session.SetBaseURL(fmt.Sprintf("http://127.0.0.1:%d", localPort))
+
+	if err = wd.initMjpegClient(); err != nil {
+		return err
+	}
+
+	// create new session
+	if err := wd.InitSession(nil); err != nil {
+		return errors.Wrap(code.DeviceHTTPDriverError, err.Error())
+	}
+
+	// init WDA scale
+	if wd.scale, err = wd.Scale(); err != nil {
+		return err
+	}
+	return nil
+}
+
+func (wd *WDADriver) TearDown() error {
+	return wd.DeleteSession()
+}
+
 func (wd *WDADriver) InitSession(capabilities option.Capabilities) error {
 	// [[FBRoute POST:@"/session"].withoutSession respondWithTarget:self action:@selector(handleCreateSession:)]
 	data := make(map[string]interface{})
@@ -156,9 +179,6 @@ func (wd *WDADriver) DeleteSession() (err error) {
 	if wd.mjpegClient != nil {
 		wd.mjpegClient.CloseIdleConnections()
 	}
-	if wd.mjpegHTTPConn != nil {
-		wd.mjpegHTTPConn.Close()
-	}
 
 	// [[FBRoute DELETE:@""] respondWithTarget:self action:@selector(handleDeleteSession:)]
 	_, err = wd.Session.DELETE("/session", wd.Session.ID)
@@ -934,11 +954,3 @@ func (wd *WDADriver) StopCaptureLog() (result interface{}, err error) {
 func (wd *WDADriver) GetSession() *DriverSession {
 	return wd.Session
 }
-
-func (wd *WDADriver) Setup() error {
-	return nil
-}
-
-func (wd *WDADriver) TearDown() error {
-	return wd.DeleteSession()
-}

pkg/uixt/ios_test.go

@@ -65,6 +65,7 @@ func TestDevice_IOS_GetPackageInfo(t *testing.T) {
 	require.Nil(t, err)
 	appInfo, err := device.GetPackageInfo("com.ss.iphone.ugc.Aweme")
 	assert.Nil(t, err)
+	assert.Equal(t, "com.ss.iphone.ugc.Aweme", appInfo.PackageName)
 	t.Logf("%+v", appInfo)
 }
 

pkg/uixt/option/android.go

@@ -58,7 +58,7 @@ func NewAndroidDeviceOptions(opts ...AndroidDeviceOption) *AndroidDeviceOptions
 		opt(config)
 	}
 
-	// set default
+	// adb default
 	if config.AdbServerHost == "" {
 		config.AdbServerHost = defaultAdbServerHost
 	}
@@ -66,18 +66,18 @@ func NewAndroidDeviceOptions(opts ...AndroidDeviceOption) *AndroidDeviceOptions
 		config.AdbServerPort = defaultAdbServerPort
 	}
 
-	if config.UIA2 {
+	// uiautomator2 default
-		if config.UIA2IP == "" && config.UIA2Port == 0 {
+	if config.UIA2IP == "" && config.UIA2Port == 0 {
-			config.UIA2IP = defaultUIA2ServerHost
+		config.UIA2IP = defaultUIA2ServerHost
-			config.UIA2Port = defaultUIA2ServerPort
+		config.UIA2Port = defaultUIA2ServerPort
-		}
-		if config.UIA2ServerPackageName == "" {
-			config.UIA2ServerPackageName = defaultUIA2ServerPackageName
-		}
-		if config.UIA2ServerTestPackageName == "" {
-			config.UIA2ServerTestPackageName = defaultUIA2ServerTestPackageName
-		}
 	}
+	if config.UIA2ServerPackageName == "" {
+		config.UIA2ServerPackageName = defaultUIA2ServerPackageName
+	}
+	if config.UIA2ServerTestPackageName == "" {
+		config.UIA2ServerTestPackageName = defaultUIA2ServerTestPackageName
+	}
+
 	return config
 }
 

runner.go

@@ -424,13 +424,11 @@ func (r *CaseRunner) parseConfig() (parsedConfig *TConfig, err error) {
 			return nil, errors.Wrap(code.InvalidCaseError,
 				fmt.Sprintf("parse android config failed: %v", err))
 		}
+
 		device, err := uixt.NewAndroidDevice(androidDevice.Options.Options()...)
 		if err != nil {
 			return nil, errors.Wrap(err, "init android device failed")
 		}
-		if err := device.Setup(); err != nil {
-			return nil, errors.Wrap(err, "setup android device failed")
-		}
 		driver, err := device.NewDriver()
 		if err != nil {
 			return nil, errors.Wrap(err, "init android driver failed")
@@ -451,9 +449,6 @@ func (r *CaseRunner) parseConfig() (parsedConfig *TConfig, err error) {
 		if err != nil {
 			return nil, errors.Wrap(err, "init ios device failed")
 		}
-		if err := device.Setup(); err != nil {
-			return nil, errors.Wrap(err, "setup ios device failed")
-		}
 		driver, err := device.NewDriver()
 		if err != nil {
 			return nil, errors.Wrap(err, "init ios driver failed")
@@ -474,9 +469,6 @@ func (r *CaseRunner) parseConfig() (parsedConfig *TConfig, err error) {
 		if err != nil {
 			return nil, errors.Wrap(err, "init harmony device failed")
 		}
-		if err := device.Setup(); err != nil {
-			return nil, errors.Wrap(err, "setup harmony device failed")
-		}
 		driver, err := device.NewDriver()
 		if err != nil {
 			return nil, errors.Wrap(err, "init harmony driver failed")

server/context.go

@@ -52,7 +52,6 @@ func (r *Router) HandleDeviceContext() gin.HandlerFunc {
 				c.Abort()
 				return
 			}
-			device.Setup()
 
 			driver, err := device.NewDriver()
 			if err != nil {

server/ext/shoots.go

@@ -194,7 +194,6 @@ func handleDeviceContext() gin.HandlerFunc {
 				c.Abort()
 				return
 			}
-			device.Setup()
 
 			driver, err := driver_ext.NewShootsAndroidDriver(device)
 			if err != nil {