From a73976a98d6b83019cb9bf5d09bfb0fbaf901f36 Mon Sep 17 00:00:00 2001
From: buyuxiang <347586493@qq.com>
Date: Mon, 24 Jul 2023 13:06:22 +0800
Subject: [PATCH] fix wda auto reset session

---
 hrp/pkg/uixt/client.go     | 28 ++++++++++++++++++++++------
 hrp/pkg/uixt/ios_device.go |  8 ++++++++
 2 files changed, 30 insertions(+), 6 deletions(-)

diff --git a/hrp/pkg/uixt/client.go b/hrp/pkg/uixt/client.go
index 7e9ce9b4..27b0fac7 100644
--- a/hrp/pkg/uixt/client.go
+++ b/hrp/pkg/uixt/client.go
@@ -147,7 +147,7 @@ func (wd *Driver) uia2HttpDELETE(pathElem ...string) (rawResp rawResponse, err e
 	return wd.uia2HttpRequest(http.MethodDelete, wd.concatURL(nil, pathElem...), nil)
 }
 
-func (wd *Driver) resetWDASession() error {
+func (wd *Driver) resetWDASession() (err error) {
 	capabilities := NewCapabilities()
 	capabilities.WithDefaultAlertAction(AlertActionAccept)
 
@@ -156,7 +156,6 @@ func (wd *Driver) resetWDASession() error {
 	data["capabilities"] = map[string]interface{}{"alwaysMatch": capabilities}
 
 	var rawResp rawResponse
-	var err error
 	if rawResp, err = wd.httpPOST(data, "/session"); err != nil {
 		return err
 	}
@@ -165,6 +164,19 @@ func (wd *Driver) resetWDASession() error {
 		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
 }
 
@@ -175,13 +187,17 @@ func (wd *Driver) wdaHttpRequest(method string, rawURL string, rawBody []byte, d
 		if err == nil || disableRetryBool {
 			return
 		}
-		// wait for WDA to resume automatically
-		time.Sleep(20 * time.Second)
+		// 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 session, retry count: %v", retryCount)
+			log.Err(err).Msgf("failed to reset wda driver, retry count: %v", retryCount)
 			continue
 		}
-		log.Debug().Str("new session", wd.sessionId).Msgf("successful to reset wda session, retry count: %v", retryCount)
+		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
 }
diff --git a/hrp/pkg/uixt/ios_device.go b/hrp/pkg/uixt/ios_device.go
index a0339e8e..03a3eb6a 100644
--- a/hrp/pkg/uixt/ios_device.go
+++ b/hrp/pkg/uixt/ios_device.go
@@ -67,6 +67,8 @@ var (
 	WithIOSPcapBundleID = gidevice.WithPcapBundleID
 )
 
+var WDALocalPort, WDALocalMjpegPort int
+
 type IOSDeviceOption func(*IOSDevice)
 
 func WithUDID(udid string) IOSDeviceOption {
@@ -592,6 +594,7 @@ 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)
@@ -609,7 +612,12 @@ 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")