feat: add RequestWithRetry for DriverSession

pkg/uixt/android_driver_uia2.go

@@ -30,6 +30,10 @@ func NewUIA2Driver(device *AndroidDevice) (*UIA2Driver, error) {
 	if err := driver.Setup(); err != nil {
 		return nil, err
 	}
+
+	// register driver session reset handler
+	driver.Session.RegisterResetHandler(driver.Setup)
+
 	return driver, nil
 }
 

pkg/uixt/driver_ext/android_stub_driver.go

@@ -45,6 +45,14 @@ func NewStubAndroidDriver(dev *uixt.AndroidDevice) (*StubAndroidDriver, error) {
 		ADBDriver: adbDriver,
 	}
 
+	// setup driver
+	if err := driver.Setup(); err != nil {
+		return nil, err
+	}
+
+	// register driver session reset handler
+	driver.Session.RegisterResetHandler(driver.Setup)
+
 	return driver, nil
 }
 
@@ -272,11 +280,6 @@ func (sad *StubAndroidDriver) LogoutNoneUI(packageName string) error {
 		log.Err(err).Msgf("%v", res)
 		return err
 	}
-	fmt.Printf("%v", resp)
-	if err != nil {
-		return err
-	}
-	time.Sleep(3 * time.Second)
 	return nil
 }
 

pkg/uixt/driver_ext/ios_stub_driver.go

@@ -44,6 +44,9 @@ func NewStubIOSDriver(dev *uixt.IOSDevice) (*StubIOSDriver, error) {
 		return nil, err
 	}
 
+	// register driver session reset handler
+	driver.Session.RegisterResetHandler(driver.Setup)
+
 	return driver, nil
 }
 

pkg/uixt/driver_session.go

@@ -20,7 +20,6 @@ import (
 	"github.com/rs/zerolog/log"
 
 	"github.com/httprunner/httprunner/v5/internal/json"
-	"github.com/httprunner/httprunner/v5/pkg/uixt/option"
 )
 
 type Attachments map[string]interface{}
@@ -61,38 +60,13 @@ type DriverSession struct {
 	timeout  time.Duration
 	maxRetry int
 
+	// used to reset driver session when request failed
+	resetFn func() error
+
 	// cache driver request and response
 	requests []*DriverRequests
 }
 
-func (s *DriverSession) Init(capabilities option.Capabilities) (err error) {
-	data := make(map[string]interface{})
-	if len(capabilities) == 0 {
-		data["capabilities"] = make(map[string]interface{})
-	} else {
-		data["capabilities"] = map[string]interface{}{"alwaysMatch": capabilities}
-	}
-	var rawResp DriverRawResponse
-	if rawResp, err = s.POST(data, "/session"); err != nil {
-		return err
-	}
-	reply := new(struct{ Value struct{ SessionId string } })
-	if err = json.Unmarshal(rawResp, reply); err != nil {
-		return err
-	}
-	s.ID = reply.Value.SessionId
-
-	// WDA
-	// sessionInfo, err := rawResp.ValueConvertToSessionInfo()
-	// if err != nil {
-	// 	return err
-	// }
-	// // update session ID
-	// wd.Session.ID = sessionInfo.ID
-
-	return nil
-}
-
 func (s *DriverSession) Reset() {
 	s.requests = make([]*DriverRequests, 0)
 }
@@ -101,6 +75,10 @@ func (s *DriverSession) SetBaseURL(baseUrl string) {
 	s.baseUrl = baseUrl
 }
 
+func (s *DriverSession) RegisterResetHandler(fn func() error) {
+	s.resetFn = fn
+}
+
 func (s *DriverSession) addRequestResult(driverResult *DriverRequests) {
 	s.requests = append(s.requests, driverResult)
 }
@@ -179,7 +157,7 @@ func (s *DriverSession) GET(pathElem ...string) (rawResp DriverRawResponse, err
 	if err != nil {
 		return nil, err
 	}
-	return s.Request(http.MethodGet, urlStr, nil)
+	return s.RequestWithRetry(http.MethodGet, urlStr, nil)
 }
 
 func (s *DriverSession) POST(data interface{}, pathElem ...string) (rawResp DriverRawResponse, err error) {
@@ -193,7 +171,7 @@ func (s *DriverSession) POST(data interface{}, pathElem ...string) (rawResp Driv
 			return nil, err
 		}
 	}
-	return s.Request(http.MethodPost, urlStr, bsJSON)
+	return s.RequestWithRetry(http.MethodPost, urlStr, bsJSON)
 }
 
 func (s *DriverSession) DELETE(pathElem ...string) (rawResp DriverRawResponse, err error) {
@@ -201,7 +179,30 @@ func (s *DriverSession) DELETE(pathElem ...string) (rawResp DriverRawResponse, e
 	if err != nil {
 		return nil, err
 	}
-	return s.Request(http.MethodDelete, urlStr, nil)
+	return s.RequestWithRetry(http.MethodDelete, urlStr, nil)
+}
+
+func (s *DriverSession) RequestWithRetry(method string, rawURL string, rawBody []byte) (
+	rawResp DriverRawResponse, err error) {
+	for count := 1; count <= s.maxRetry; count++ {
+		rawResp, err = s.Request(method, rawURL, rawBody)
+		if err == nil {
+			return
+		}
+		time.Sleep(3 * time.Second)
+
+		if s.resetFn != nil {
+			log.Warn().Msg("reset driver session")
+			if err2 := s.resetFn(); err2 != nil {
+				log.Error().Err(err2).Msgf(
+					"failed to reset session, try count %v", count)
+			} else {
+				log.Info().Msgf(
+					"reset session success, try count %v", count)
+			}
+		}
+	}
+	return
 }
 
 func (s *DriverSession) Request(method string, rawURL string, rawBody []byte) (
@@ -285,30 +286,6 @@ func (s *DriverSession) Request(method string, rawURL string, rawBody []byte) (
 	return
 }
 
-// TODO: FIXME
-func (s *DriverSession) RequestWithRetry(method string, rawURL string, rawBody []byte) (
-	rawResp DriverRawResponse, err error) {
-	for count := 1; count <= s.maxRetry; count++ {
-		rawResp, err = s.Request(method, rawURL, rawBody)
-		if err == nil {
-			return
-		}
-		time.Sleep(3 * time.Second)
-		oldSessionID := s.ID
-		if err2 := s.Init(nil); err2 != nil {
-			log.Error().Err(err2).Msgf(
-				"failed to reset session, try count %v", count)
-			continue
-		}
-		log.Debug().Str("new session", s.ID).Str("old session", oldSessionID).Msgf(
-			"reset session successfully, try count %v", count)
-		if oldSessionID != "" {
-			rawURL = strings.Replace(rawURL, oldSessionID, s.ID, 1)
-		}
-	}
-	return
-}
-
 func (s *DriverSession) SetupPortForward(localPort int) error {
 	conn, err := net.Dial("tcp", fmt.Sprintf("127.0.0.1:%d", localPort))
 	if err != nil {

pkg/uixt/ios_driver_wda.go

@@ -42,6 +42,9 @@ func NewWDADriver(device *IOSDevice) (*WDADriver, error) {
 		return nil, err
 	}
 
+	// register driver session reset handler
+	driver.Session.RegisterResetHandler(driver.Setup)
+
 	return driver, nil
 }