refactor UIA2 / WDA driver request

2026-05-10 17:43:00 +08:00 · 2023-07-24 17:38:32 +08:00
parent 0cbc5eec5a
commit c91be64240
5 changed files with 120 additions and 144 deletions
--- a/docs/CHANGELOG.md
+++ b/docs/CHANGELOG.md
@@ -1,12 +1,14 @@
 # Release History

+## v4.3.6 (2023-07-24)
+
+- feat: support to reset driver (or session only) automatically when UIA2 / WDA crashed or WebDriver request failed
+
 ## v4.3.5 (2023-07-23)

 - refactor: send events to Google Analytics 4, replace GA v1
 - fix: failure unittests caused by httpbin.org, replace with docker service
 - fix: handle unstable unittests, restore github actions pipeline
- feat: support to reset driver automatically when uia2 crashed
- feat: support to reset session when wda request failed

 **go version**

--- a/hrp/pkg/uixt/android_uia2_driver.go
+++ b/hrp/pkg/uixt/android_uia2_driver.go
@@ -6,9 +6,11 @@ import (
 	"encoding/json"
 	"fmt"
 	"net"
+	"net/http"
 	"net/url"
 	"strconv"
 	"strings"
+	"time"

 	"github.com/pkg/errors"
 	"github.com/rs/zerolog/log"
@@ -81,11 +83,65 @@ func (bs BatteryStatus) String() string {
 	}
 }

+func (ud *uiaDriver) resetUIA2Driver() error {
+	newUIADriver, err := NewUIADriver(NewCapabilities(), ud.urlPrefix.String())
+	if err != nil {
+		return err
+	}
+	ud.client = newUIADriver.client
+	ud.sessionId = newUIADriver.sessionId
+	return nil
+}
+
+func (ud *uiaDriver) uia2HttpRequest(method string, rawURL string, rawBody []byte, disableRetry ...bool) (rawResp rawResponse, err error) {
+	disableRetryBool := len(disableRetry) > 0 && disableRetry[0]
+	for retryCount := 1; retryCount <= 5; retryCount++ {
+		rawResp, err = ud.httpRequest(method, rawURL, rawBody)
+		if err == nil || disableRetryBool {
+			return
+		}
+		// wait for UIA2 server to resume automatically
+		time.Sleep(3 * time.Second)
+		oldSessionID := ud.sessionId
+		if err = ud.resetUIA2Driver(); err != nil {
+			log.Err(err).Msgf("failed to reset uia2 driver, retry count: %v", retryCount)
+			continue
+		}
+		log.Debug().Str("new session", ud.sessionId).Str("old session", oldSessionID).Msgf("successful to reset uia2 driver, retry count: %v", retryCount)
+		if oldSessionID != "" {
+			rawURL = strings.Replace(rawURL, oldSessionID, ud.sessionId, 1)
+		}
+	}
+	return
+}
+
+func (ud *uiaDriver) uia2HttpGET(pathElem ...string) (rawResp rawResponse, err error) {
+	return ud.uia2HttpRequest(http.MethodGet, ud.concatURL(nil, pathElem...), nil)
+}
+
+func (ud *uiaDriver) uia2HttpGETWithRetry(pathElem ...string) (rawResp rawResponse, err error) {
+	return ud.uia2HttpRequest(http.MethodGet, ud.concatURL(nil, pathElem...), nil, true)
+}
+
+func (ud *uiaDriver) uia2HttpPOST(data interface{}, pathElem ...string) (rawResp rawResponse, err error) {
+	var bsJSON []byte = nil
+	if data != nil {
+		if bsJSON, err = json.Marshal(data); err != nil {
+			return nil, err
+		}
+	}
+	return ud.uia2HttpRequest(http.MethodPost, ud.concatURL(nil, pathElem...), bsJSON)
+}
+
+func (ud *uiaDriver) uia2HttpDELETE(pathElem ...string) (rawResp rawResponse, err error) {
+	return ud.uia2HttpRequest(http.MethodDelete, ud.concatURL(nil, pathElem...), nil)
+}
+
 func (ud *uiaDriver) NewSession(capabilities Capabilities) (sessionInfo SessionInfo, err error) {
 	// register(postHandler, new NewSession("/wd/hub/session"))
 	var rawResp rawResponse
 	data := map[string]interface{}{"capabilities": capabilities}
-	if rawResp, err = ud.uia2HttpPOST(data, "/session"); err != nil {
+	if rawResp, err = ud.httpPOST(data, "/session"); err != nil {
 		return SessionInfo{SessionId: ""}, err
 	}
 	reply := new(struct{ Value struct{ SessionId string } })
--- a/hrp/pkg/uixt/client.go
+++ b/hrp/pkg/uixt/client.go
@@ -93,137 +93,6 @@ func (wd *Driver) httpRequest(method string, rawURL string, rawBody []byte) (raw
 	return
 }

-func (wd *Driver) resetUIA2Driver() error {
-	ud, err := NewUIADriver(NewCapabilities(), wd.urlPrefix.String())
-	if err != nil {
-		return err
-	}
-	wd.client = ud.client
-	wd.sessionId = ud.sessionId
-	return nil
-}
-
-func (wd *Driver) uia2HttpRequest(method string, rawURL string, rawBody []byte, disableRetry ...bool) (rawResp rawResponse, err error) {
-	disableRetryBool := len(disableRetry) > 0 && disableRetry[0]
-	for retryCount := 1; retryCount <= 5; retryCount++ {
-		rawResp, err = wd.httpRequest(method, rawURL, rawBody)
-		if err == nil || disableRetryBool {
-			return
-		}
-		// wait for UIA2 server to resume automatically
-		time.Sleep(3 * time.Second)
-		oldSessionID := wd.sessionId
-		if err = wd.resetUIA2Driver(); err != nil {
-			log.Err(err).Msgf("failed to reset uia2 driver, retry count: %v", retryCount)
-			continue
-		}
-		log.Debug().Str("new session", wd.sessionId).Str("old session", oldSessionID).Msgf("successful to reset uia2 driver, retry count: %v", retryCount)
-		if oldSessionID != "" {
-			rawURL = strings.Replace(rawURL, oldSessionID, wd.sessionId, 1)
-		}
-	}
-	return
-}
-
-func (wd *Driver) uia2HttpGET(pathElem ...string) (rawResp rawResponse, err error) {
-	return wd.uia2HttpRequest(http.MethodGet, wd.concatURL(nil, pathElem...), nil)
-}
-
-func (wd *Driver) uia2HttpGETWithRetry(pathElem ...string) (rawResp rawResponse, err error) {
-	return wd.uia2HttpRequest(http.MethodGet, wd.concatURL(nil, pathElem...), nil, true)
-}
-
-func (wd *Driver) uia2HttpPOST(data interface{}, pathElem ...string) (rawResp rawResponse, err error) {
-	var bsJSON []byte = nil
-	if data != nil {
-		if bsJSON, err = json.Marshal(data); err != nil {
-			return nil, err
-		}
-	}
-	return wd.uia2HttpRequest(http.MethodPost, wd.concatURL(nil, pathElem...), bsJSON)
-}
-
-func (wd *Driver) uia2HttpDELETE(pathElem ...string) (rawResp rawResponse, err error) {
-	return wd.uia2HttpRequest(http.MethodDelete, wd.concatURL(nil, pathElem...), nil)
-}
-
-func (wd *Driver) resetWDASession() (err error) {
-	capabilities := NewCapabilities()
-	capabilities.WithDefaultAlertAction(AlertActionAccept)
-
-	// [[FBRoute POST:@"/session"].withoutSession respondWithTarget:self action:@selector(handleCreateSession:)]
-	data := make(map[string]interface{})
-	data["capabilities"] = map[string]interface{}{"alwaysMatch": capabilities}
-
-	var rawResp rawResponse
-	if rawResp, err = wd.httpPOST(data, "/session"); err != nil {
-		return err
-	}
-	var sessionInfo SessionInfo
-	if sessionInfo, err = rawResp.valueConvertToSessionInfo(); err != nil {
-		return err
-	}
-	wd.sessionId = sessionInfo.SessionId
-	return
-}
-
-func (wd *Driver) resetWDADriver() error {
-	capabilities := NewCapabilities()
-	capabilities.WithDefaultAlertAction(AlertActionAccept)
-
-	wdaDriver, err := NewWDADriver(capabilities, WDALocalPort, WDALocalMjpegPort)
-	if err != nil {
-		return err
-	}
-	wd.client = wdaDriver.client
-	wd.sessionId = wdaDriver.sessionId
-	return nil
-}
-
-func (wd *Driver) wdaHttpRequest(method string, rawURL string, rawBody []byte, disableRetry ...bool) (rawResp rawResponse, err error) {
-	disableRetryBool := len(disableRetry) > 0 && disableRetry[0]
-	for retryCount := 1; retryCount <= 5; retryCount++ {
-		rawResp, err = wd.httpRequest(method, rawURL, rawBody)
-		if err == nil || disableRetryBool {
-			return
-		}
-		// TODO: polling WDA to check if resumed automatically
-		time.Sleep(5 * time.Second)
-		oldSessionID := wd.sessionId
-		if err = wd.resetWDASession(); err != nil {
-			log.Err(err).Msgf("failed to reset wda driver, retry count: %v", retryCount)
-			continue
-		}
-		log.Debug().Str("new session", wd.sessionId).Str("old session", oldSessionID).Msgf("successful to reset wda driver, retry count: %v", retryCount)
-		if oldSessionID != "" {
-			rawURL = strings.Replace(rawURL, oldSessionID, wd.sessionId, 1)
-		}
-	}
-	return
-}
-
-func (wd *Driver) wdaHttpGET(pathElem ...string) (rawResp rawResponse, err error) {
-	return wd.wdaHttpRequest(http.MethodGet, wd.concatURL(nil, pathElem...), nil)
-}
-
-func (wd *Driver) wdaHttpGETWithRetry(pathElem ...string) (rawResp rawResponse, err error) {
-	return wd.wdaHttpRequest(http.MethodGet, wd.concatURL(nil, pathElem...), nil, true)
-}
-
-func (wd *Driver) wdaHttpPOST(data interface{}, pathElem ...string) (rawResp rawResponse, err error) {
-	var bsJSON []byte = nil
-	if data != nil {
-		if bsJSON, err = json.Marshal(data); err != nil {
-			return nil, err
-		}
-	}
-	return wd.wdaHttpRequest(http.MethodPost, wd.concatURL(nil, pathElem...), bsJSON)
-}
-
-func (wd *Driver) wdaHttpDELETE(pathElem ...string) (rawResp rawResponse, err error) {
-	return wd.wdaHttpRequest(http.MethodDelete, wd.concatURL(nil, pathElem...), nil)
-}
-
 func convertToHTTPClient(conn net.Conn) *http.Client {
 	return &http.Client{
 		Transport: &http.Transport{
--- a/hrp/pkg/uixt/ios_device.go
+++ b/hrp/pkg/uixt/ios_device.go
@@ -67,8 +67,6 @@ var (
 	WithIOSPcapBundleID = gidevice.WithPcapBundleID
 )

-var WDALocalPort, WDALocalMjpegPort int
-
 type IOSDeviceOption func(*IOSDevice)

 func WithUDID(udid string) IOSDeviceOption {
@@ -476,7 +474,7 @@ func (dev *IOSDevice) forward(localPort, remotePort int) error {
 			rInnerConn, err := device.NewConnect(remotePort)
 			if err != nil {
 				log.Error().Err(err).Msg("connect to ios device failed")
-				continue
+				os.Exit(code.GetErrorCode(code.IOSDeviceConnectionError))
 			}

 			rConn := rInnerConn.RawConn()
@@ -594,7 +592,6 @@ func (dev *IOSDevice) NewHTTPDriver(capabilities Capabilities) (driver WebDriver
 	} else {
 		log.Info().Int("WDA_LOCAL_PORT", localPort).Msg("reuse WDA local port")
 	}
-	WDALocalPort = localPort

 	var localMjpegPort int
 	localMjpegPort, err = strconv.Atoi(env.WDA_LOCAL_MJPEG_PORT)
@@ -612,12 +609,7 @@ func (dev *IOSDevice) NewHTTPDriver(capabilities Capabilities) (driver WebDriver
 		log.Info().Int("WDA_LOCAL_MJPEG_PORT", localMjpegPort).
 			Msg("reuse WDA local mjpeg port")
 	}
-	WDALocalMjpegPort = localMjpegPort

-	return NewWDADriver(capabilities, localPort, localMjpegPort)
-}
-
-func NewWDADriver(capabilities Capabilities, localPort, localMjpegPort int) (driver *wdaDriver, err error) {
 	log.Info().Interface("capabilities", capabilities).
 		Int("localPort", localPort).Int("localMjpegPort", localMjpegPort).
 		Msg("init WDA HTTP driver")
--- a/hrp/pkg/uixt/ios_driver.go
+++ b/hrp/pkg/uixt/ios_driver.go
@@ -10,6 +10,7 @@ import (
 	"net/url"
 	"regexp"
 	"strings"
+	"time"

 	"github.com/pkg/errors"
 	"github.com/rs/zerolog/log"
@@ -31,6 +32,62 @@ type wdaDriver struct {
 	mjpegClient   *http.Client
 }

+func (wd *wdaDriver) resetWDASession() error {
+	capabilities := NewCapabilities()
+	capabilities.WithDefaultAlertAction(AlertActionAccept)
+
+	sessionInfo, err := wd.NewSession(capabilities)
+	if err != nil {
+		return err
+	}
+	wd.sessionId = sessionInfo.SessionId
+	return nil
+}
+
+func (wd *wdaDriver) wdaHttpRequest(method string, rawURL string, rawBody []byte, disableRetry ...bool) (rawResp rawResponse, err error) {
+	disableRetryBool := len(disableRetry) > 0 && disableRetry[0]
+	for retryCount := 1; retryCount <= 5; retryCount++ {
+		rawResp, err = wd.httpRequest(method, rawURL, rawBody)
+		if err == nil || disableRetryBool {
+			return
+		}
+		// TODO: polling WDA to check if resumed automatically
+		time.Sleep(5 * time.Second)
+		oldSessionID := wd.sessionId
+		if err = wd.resetWDASession(); err != nil {
+			log.Err(err).Msgf("failed to reset wda driver, retry count: %v", retryCount)
+			continue
+		}
+		log.Debug().Str("new session", wd.sessionId).Str("old session", oldSessionID).Msgf("successful to reset wda driver, retry count: %v", retryCount)
+		if oldSessionID != "" {
+			rawURL = strings.Replace(rawURL, oldSessionID, wd.sessionId, 1)
+		}
+	}
+	return
+}
+
+func (wd *wdaDriver) wdaHttpGET(pathElem ...string) (rawResp rawResponse, err error) {
+	return wd.wdaHttpRequest(http.MethodGet, wd.concatURL(nil, pathElem...), nil)
+}
+
+func (wd *wdaDriver) wdaHttpGETWithRetry(pathElem ...string) (rawResp rawResponse, err error) {
+	return wd.wdaHttpRequest(http.MethodGet, wd.concatURL(nil, pathElem...), nil, true)
+}
+
+func (wd *wdaDriver) wdaHttpPOST(data interface{}, pathElem ...string) (rawResp rawResponse, err error) {
+	var bsJSON []byte = nil
+	if data != nil {
+		if bsJSON, err = json.Marshal(data); err != nil {
+			return nil, err
+		}
+	}
+	return wd.wdaHttpRequest(http.MethodPost, wd.concatURL(nil, pathElem...), bsJSON)
+}
+
+func (wd *wdaDriver) wdaHttpDELETE(pathElem ...string) (rawResp rawResponse, err error) {
+	return wd.wdaHttpRequest(http.MethodDelete, wd.concatURL(nil, pathElem...), nil)
+}
+
 func (wd *wdaDriver) GetMjpegClient() *http.Client {
 	return wd.mjpegClient
 }
@@ -45,7 +102,7 @@ func (wd *wdaDriver) NewSession(capabilities Capabilities) (sessionInfo SessionI
 	}

 	var rawResp rawResponse
-	if rawResp, err = wd.wdaHttpPOST(data, "/session"); err != nil {
+	if rawResp, err = wd.httpPOST(data, "/session"); err != nil {
 		return SessionInfo{}, err
 	}
 	if sessionInfo, err = rawResp.valueConvertToSessionInfo(); err != nil {