feature: 增加ws连接可以继承到引用的testcases中

hrp/runner.go

@@ -461,6 +461,7 @@ type SessionRunner struct {
 	startTime         time.Time                  // record start time of the testcase
 	summary           *TestCaseSummary           // record test case summary
 	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
 }
@@ -472,19 +473,28 @@ func (r *SessionRunner) resetSession() {
 	r.startTime = time.Now()
 	r.summary = newSummary()
 	r.wsConnMap = make(map[string]*websocket.Conn)
+	r.inheritWsConnMap = make(map[string]*websocket.Conn)
 	r.pongResponseChan = make(chan string, 1)
 	r.closeResponseChan = make(chan *wsCloseRespObject, 1)
 }
 
+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.
 // givenVars is used for data driven
 func (r *SessionRunner) Start(givenVars map[string]interface{}) error {
 	config := r.caseRunner.testCase.Config
 	log.Info().Str("testcase", config.Name).Msg("run testcase start")
 
-	// reset session runner
-	r.resetSession()
-
 	// update config variables with given variables
 	r.InitWithParameters(givenVars)
 
@@ -652,3 +662,15 @@ func (r *SessionRunner) releaseResources() {
 		}
 	}
 }
+
+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,6 +89,8 @@ 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

@@ -314,7 +314,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.wsConnMap[parsedURL] != nil {
+		if r.getWsClient(parsedURL) != nil {
 			break
 		}
 		resp, err = openWithTimeout(parsedURL, parsedHeader, r, step)
@@ -476,10 +476,15 @@ func openWithTimeout(urlStr string, requestHeader http.Header, r *SessionRunner,
 		conn.SetCloseHandler(func(code int, text string) error {
 			message := websocket.FormatCloseMessage(code, "")
 			conn.WriteControl(websocket.CloseMessage, message, time.Now().Add(defaultWriteWait))
-			r.closeResponseChan <- &wsCloseRespObject{
+			select {
+			case r.closeResponseChan <- &wsCloseRespObject{
 				StatusCode: code,
 				Text:       text,
+			}:
+			default:
+				log.Warn().Msg("close response channel is block, drop the response")
 			}
+
 			return nil
 		})
 		r.wsConnMap[urlStr] = conn
@@ -499,7 +504,7 @@ func openWithTimeout(urlStr string, requestHeader http.Header, r *SessionRunner,
 }
 
 func readMessageWithTimeout(urlString string, r *SessionRunner, step *TStep) (*wsReadRespObject, error) {
-	wsConn := r.wsConnMap[urlString]
+	wsConn := r.getWsClient(urlString)
 	if wsConn == nil {
 		return nil, errors.New("try to use existing connection, but there is no connection")
 	}
@@ -529,7 +534,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.wsConnMap[urlString]
+	wsConn := r.getWsClient(urlString)
 	if wsConn == nil {
 		return errors.New("try to use existing connection, but there is no connection")
 	}
@@ -595,7 +600,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.wsConnMap[urlString]
+	wsConn := r.getWsClient(urlString)
 	if wsConn == nil {
 		return nil, errors.New("no connection needs to be closed")
 	}