refactor: move websocket related logic to step_websocket

hrp/runner.go

@@ -418,11 +418,7 @@ func (r *CaseRunner) NewSession() *SessionRunner {
 		sessionVariables: make(map[string]interface{}),
 		summary:          newSummary(),
 
-		transactions:      make(map[string]map[transactionType]time.Time),
-		wsConnMap:         make(map[string]*websocket.Conn),
-		inheritWsConnMap:  make(map[string]*websocket.Conn),
-		pongResponseChan:  make(chan string, 1),
-		closeResponseChan: make(chan *wsCloseRespObject, 1),
+		transactions: make(map[string]map[transactionType]time.Time),
 	}
 	return sessionRunner
 }
@@ -438,22 +434,6 @@ type SessionRunner struct {
 	// transactions stores transaction timing info.
 	// key is transaction name, value is map of transaction type and time, e.g. start time and end time.
 	transactions map[string]map[transactionType]time.Time
-
-	wsConnMap         map[string]*websocket.Conn // save all websocket connections
-	inheritWsConnMap  map[string]*websocket.Conn // inherit all websocket connections
-	pongResponseChan  chan string                // channel used to receive pong response message
-	closeResponseChan chan *wsCloseRespObject    // channel used to receive close response message
-}
-
-func (r *SessionRunner) inheritConnection(src *SessionRunner) {
-	log.Info().Msg("inherit session runner")
-	r.inheritWsConnMap = make(map[string]*websocket.Conn, len(src.wsConnMap)+len(src.inheritWsConnMap))
-	for k, v := range src.wsConnMap {
-		r.inheritWsConnMap[k] = v
-	}
-	for k, v := range src.inheritWsConnMap {
-		r.inheritWsConnMap[k] = v
-	}
 }
 
 // Start runs the test steps in sequential order.
@@ -468,11 +448,6 @@ func (r *SessionRunner) Start(givenVars map[string]interface{}) error {
 	// update config variables with given variables
 	r.InitWithParameters(givenVars)
 
-	defer func() {
-		// close session resource after all steps done or fast fail
-		r.releaseResources()
-	}()
-
 	// run step in sequential order
 	for _, step := range r.caseRunner.TestSteps {
 		select {
@@ -675,29 +650,3 @@ func (r *SessionRunner) addSingleStepResult(stepResult *StepResult) {
 		r.summary.Success = false
 	}
 }
-
-// releaseResources releases resources used by session runner
-func (r *SessionRunner) releaseResources() {
-	// close websocket connections
-	for _, wsConn := range r.wsConnMap {
-		if wsConn != nil {
-			log.Info().Str("testcase", r.caseRunner.Config.Name).Msg("websocket disconnected")
-			err := wsConn.Close()
-			if err != nil {
-				log.Error().Err(err).Msg("websocket disconnection failed")
-			}
-		}
-	}
-}
-
-func (r *SessionRunner) getWsClient(url string) *websocket.Conn {
-	if client, ok := r.wsConnMap[url]; ok {
-		return client
-	}
-
-	if client, ok := r.inheritWsConnMap[url]; ok {
-		return client
-	}
-
-	return nil
-}

hrp/step_testcase.go

@@ -89,8 +89,6 @@ func (s *StepTestCaseWithOptionalArgs) Run(r *SessionRunner) (stepResult *StepRe
 		return stepResult, err
 	}
 	sessionRunner := caseRunner.NewSession()
-	// need to inherit some information from current session
-	sessionRunner.inheritConnection(r)
 
 	start := time.Now()
 	// run referenced testcase with step variables

hrp/step_websocket.go

@@ -16,6 +16,12 @@ import (
 	"github.com/httprunner/httprunner/v4/hrp/internal/json"
 )
 
+var (
+	wsConnMap         map[string]*websocket.Conn // save all websocket connections
+	pongResponseChan  chan string                // channel used to receive pong response message
+	closeResponseChan chan *wsCloseRespObject    // channel used to receive close response message
+)
+
 const (
 	wsOpen         ActionType = "open"
 	wsPing         ActionType = "ping"
@@ -314,7 +320,7 @@ func runStepWebSocket(r *SessionRunner, step *TStep) (stepResult *StepResult, er
 	case wsOpen:
 		log.Info().Int64("timeout(ms)", step.WebSocket.GetTimeout()).Str("url", parsedURL).Msg("open websocket connection")
 		// use the current websocket connection if existed
-		if r.getWsClient(parsedURL) != nil {
+		if getWsClient(parsedURL) != nil {
 			break
 		}
 		resp, err = openWithTimeout(parsedURL, parsedHeader, r, step)
@@ -334,7 +340,7 @@ func runStepWebSocket(r *SessionRunner, step *TStep) (stepResult *StepResult, er
 			case <-timer.C:
 				timer.Stop()
 				log.Warn().Msg("pong timeout")
-			case pongResponse := <-r.pongResponseChan:
+			case pongResponse := <-pongResponseChan:
 				resp = pongResponse
 				log.Info().Msg("pong message arrives")
 			}
@@ -426,6 +432,20 @@ func runStepWebSocket(r *SessionRunner, step *TStep) (stepResult *StepResult, er
 	return stepResult, nil
 }
 
+func getWsClient(url string) *websocket.Conn {
+	if wsConnMap == nil {
+		wsConnMap = make(map[string]*websocket.Conn)
+		pongResponseChan = make(chan string, 1)
+		closeResponseChan = make(chan *wsCloseRespObject, 1)
+	}
+
+	if client, ok := wsConnMap[url]; ok {
+		return client
+	}
+
+	return nil
+}
+
 func printWebSocketResponse(resp interface{}) error {
 	if resp == nil {
 		fmt.Println("(response body is empty in this step)")
@@ -470,14 +490,14 @@ func openWithTimeout(urlStr string, requestHeader http.Header, r *SessionRunner,
 
 		// the following handlers handle and transport control message from server
 		conn.SetPongHandler(func(appData string) error {
-			r.pongResponseChan <- appData
+			pongResponseChan <- appData
 			return nil
 		})
 		conn.SetCloseHandler(func(code int, text string) error {
 			message := websocket.FormatCloseMessage(code, "")
 			conn.WriteControl(websocket.CloseMessage, message, time.Now().Add(defaultWriteWait))
 			select {
-			case r.closeResponseChan <- &wsCloseRespObject{
+			case closeResponseChan <- &wsCloseRespObject{
 				StatusCode: code,
 				Text:       text,
 			}:
@@ -487,7 +507,7 @@ func openWithTimeout(urlStr string, requestHeader http.Header, r *SessionRunner,
 
 			return nil
 		})
-		r.wsConnMap[urlStr] = conn
+		wsConnMap[urlStr] = conn
 		openResponseChan <- resp
 	}()
 
@@ -504,7 +524,7 @@ func openWithTimeout(urlStr string, requestHeader http.Header, r *SessionRunner,
 }
 
 func readMessageWithTimeout(urlString string, r *SessionRunner, step *TStep) (*wsReadRespObject, error) {
-	wsConn := r.getWsClient(urlString)
+	wsConn := getWsClient(urlString)
 	if wsConn == nil {
 		return nil, errors.New("try to use existing connection, but there is no connection")
 	}
@@ -534,7 +554,7 @@ func readMessageWithTimeout(urlString string, r *SessionRunner, step *TStep) (*w
 }
 
 func writeWebSocket(urlString string, r *SessionRunner, step *TStep, stepVariables map[string]interface{}) error {
-	wsConn := r.getWsClient(urlString)
+	wsConn := getWsClient(urlString)
 	if wsConn == nil {
 		return errors.New("try to use existing connection, but there is no connection")
 	}
@@ -600,7 +620,7 @@ func writeWithAction(c *websocket.Conn, step *TStep, messageType int, message []
 }
 
 func closeWithTimeout(urlString string, r *SessionRunner, step *TStep, stepVariables map[string]interface{}) (*wsCloseRespObject, error) {
-	wsConn := r.getWsClient(urlString)
+	wsConn := getWsClient(urlString)
 	if wsConn == nil {
 		return nil, errors.New("no connection needs to be closed")
 	}
@@ -639,7 +659,7 @@ func closeWithTimeout(urlString string, r *SessionRunner, step *TStep, stepVaria
 		return nil, errors.New("close timeout")
 	case err := <-errorChan:
 		return nil, err
-	case closeResult := <-r.closeResponseChan:
+	case closeResult := <-closeResponseChan:
 		return closeResult, nil
 	}
 }
@@ -673,3 +693,17 @@ func getContentSize(resp interface{}) int64 {
 		return -1
 	}
 }
+
+// releaseResources releases resources used by session runner
+func (r *SessionRunner) releaseResources() {
+	// close websocket connections
+	for _, wsConn := range wsConnMap {
+		if wsConn != nil {
+			log.Info().Str("testcase", r.caseRunner.Config.Name).Msg("websocket disconnected")
+			err := wsConn.Close()
+			if err != nil {
+				log.Error().Err(err).Msg("websocket disconnection failed")
+			}
+		}
+	}
+}