feat: support websocket protocol

hrp/session.go

@@ -4,6 +4,7 @@ import (
 	_ "embed"
 	"time"
 
+	"github.com/gorilla/websocket"
 	"github.com/pkg/errors"
 	"github.com/rs/zerolog/log"
 )
@@ -15,9 +16,12 @@ type SessionRunner struct {
 	sessionVariables map[string]interface{}
 	// 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
-	startTime    time.Time        // record start time of the testcase
-	summary      *TestCaseSummary // record test case summary
+	transactions      map[string]map[transactionType]time.Time
+	startTime         time.Time               // record start time of the testcase
+	summary           *TestCaseSummary        // record test case summary
+	wsConn            *websocket.Conn         // one websocket connection each session
+	pongResponseChan  chan string             // channel used to receive pong response message
+	closeResponseChan chan *wsCloseRespObject // channel used to receive close response message
 }
 
 func (r *SessionRunner) resetSession() {
@@ -26,6 +30,8 @@ func (r *SessionRunner) resetSession() {
 	r.transactions = make(map[string]map[transactionType]time.Time)
 	r.startTime = time.Now()
 	r.summary = newSummary()
+	r.pongResponseChan = make(chan string, 1)
+	r.closeResponseChan = make(chan *wsCloseRespObject, 1)
 }
 
 func (r *SessionRunner) GetParser() *Parser {
@@ -83,6 +89,17 @@ func (r *SessionRunner) Start(givenVars map[string]interface{}) error {
 			Msg("run step end")
 	}
 
+	// close websocket connection after all steps done
+	defer func() {
+		if r.wsConn != nil {
+			log.Info().Str("testcase", config.Name).Msg("websocket disconnected")
+			err := r.wsConn.Close()
+			if err != nil {
+				log.Error().Err(err).Msg("websocket disconnection failed")
+			}
+		}
+	}()
+
 	log.Info().Str("testcase", config.Name).Msg("run testcase end")
 	return nil
 }