feat: support websocket protocol

docs/CHANGELOG.md

@@ -12,6 +12,7 @@
 - feat: add `--profile` flag for har2case to support overwrite headers/cookies with specified yaml/json profile file
 - feat: support run testcases in specified folder path, including testcases in sub folders
 - feat: support HTTP/2 protocol
+- feat: support WebSocket protocol
 - change: integrate [sentry sdk][sentry sdk] for panic reporting and analysis
 - change: lock funplugin version when creating scaffold project
 - fix: call referenced api/testcase with relative path

examples/hrp/http2_test.json

@@ -0,0 +1,146 @@
+{
+    "config": {
+        "name": "run request with HTTP/1.1 and HTTP/2",
+        "base_url": "https://postman-echo.com"
+    },
+    "teststeps": [
+        {
+            "name": "HTTP/1.1 get",
+            "request": {
+                "method": "GET",
+                "url": "/get",
+                "params": {
+                    "foo1": "foo1",
+                    "foo2": "foo2"
+                },
+                "headers": {
+                    "User-Agent": "HttpRunnerPlus"
+                }
+            },
+            "validate": [
+                {
+                    "check": "status_code",
+                    "assert": "equals",
+                    "expect": 200,
+                    "msg": "check status code"
+                },
+                {
+                    "check": "proto",
+                    "assert": "equals",
+                    "expect": "HTTP/1.1",
+                    "msg": "check protocol type"
+                },
+                {
+                    "check": "body.args.foo1",
+                    "assert": "length_equals",
+                    "expect": 4,
+                    "msg": "check param foo1"
+                }
+            ]
+        },
+        {
+            "name": "HTTP/1.1 post",
+            "request": {
+                "method": "POST",
+                "url": "/post",
+                "headers": {
+                    "User-Agent": "HttpRunnerPlus"
+                },
+                "body": {
+                    "foo1": "foo1",
+                    "foo2": "foo2"
+                }
+            },
+            "validate": [
+                {
+                    "check": "status_code",
+                    "assert": "equals",
+                    "expect": 200,
+                    "msg": "check status code"
+                },
+                {
+                    "check": "proto",
+                    "assert": "equals",
+                    "expect": "HTTP/1.1",
+                    "msg": "check protocol type"
+                },
+                {
+                    "check": "body.json.foo1",
+                    "assert": "length_equals",
+                    "expect": 4,
+                    "msg": "check body foo1"
+                }
+            ]
+        },
+        {
+            "name": "HTTP/2 get",
+            "request": {
+                "method": "GET",
+                "url": "/get",
+                "http2": true,
+                "params": {
+                    "foo1": "foo1",
+                    "foo2": "foo2"
+                },
+                "headers": {
+                    "User-Agent": "HttpRunnerPlus"
+                }
+            },
+            "validate": [
+                {
+                    "check": "status_code",
+                    "assert": "equals",
+                    "expect": 200,
+                    "msg": "check status code"
+                },
+                {
+                    "check": "proto",
+                    "assert": "equals",
+                    "expect": "HTTP/2.0",
+                    "msg": "check protocol type"
+                },
+                {
+                    "check": "body.args.foo1",
+                    "assert": "length_equals",
+                    "expect": 4,
+                    "msg": "check param foo1"
+                }
+            ]
+        },
+        {
+            "name": "HTTP/2 post",
+            "request": {
+                "method": "POST",
+                "url": "/post",
+                "http2": true,
+                "headers": {
+                    "User-Agent": "HttpRunnerPlus"
+                },
+                "body": {
+                    "foo1": "foo1",
+                    "foo2": "foo2"
+                }
+            },
+            "validate": [
+                {
+                    "check": "status_code",
+                    "assert": "equals",
+                    "expect": 200,
+                    "msg": "check status code"
+                },
+                {
+                    "check": "proto",
+                    "assert": "equals",
+                    "expect": "HTTP/2.0",
+                    "msg": "check protocol type"
+                },
+                {
+                    "check": "body.json.foo1",
+                    "assert": "length_equals",
+                    "expect": 4,
+                    "msg": "check body foo1"
+                }
+            ]
+        }
+    ]
+}

examples/hrp/websocket_test.json

@@ -0,0 +1,145 @@
+{
+    "config": {
+        "name": "run request with WebSocket protocol",
+        "base_url": "ws://echo.websocket.events",
+        "variables": {
+            "a": 12.3,
+            "b": 3.45,
+            "n": 5
+        }
+    },
+    "teststeps": [
+        {
+            "name": "open connection",
+            "websocket": {
+                "type": "open",
+                "url": "/",
+                "headers": {
+                    "User-Agent": "HttpRunnerPlus"
+                }
+            },
+            "validate": [
+                {
+                    "check": "status_code",
+                    "assert": "equals",
+                    "expect": 101,
+                    "msg": "check open status code"
+                },
+                {
+                    "check": "headers.Connection",
+                    "assert": "equals",
+                    "expect": "Upgrade",
+                    "msg": "check headers"
+                }
+            ]
+        },
+        {
+            "name": "ping pong test",
+            "websocket": {
+                "type": "ping",
+                "url": "/",
+                "timeout": 5000
+            }
+        },
+        {
+            "name": "read sponsor info",
+            "websocket": {
+                "type": "r",
+                "url": "/",
+                "timeout": 5000
+            },
+            "validate": [
+                {
+                    "check": "body",
+                    "assert": "contains",
+                    "expect": "Lob.com",
+                    "msg": "check sponsor message"
+                }
+            ]
+        },
+        {
+            "name": "write json",
+            "websocket": {
+                "type": "w",
+                "url": "/",
+                "text": {
+                    "foo1": "${gen_random_string($n)}",
+                    "foo2": "${max($a, $b)}"
+                }
+            }
+        },
+        {
+            "name": "read json",
+            "websocket": {
+                "type": "r",
+                "url": "/"
+            },
+            "extract": {
+                "varFoo1": "body.foo1"
+            },
+            "validate": [
+                {
+                    "check": "body.foo1",
+                    "assert": "length_equals",
+                    "expect": 5,
+                    "msg": "check json foo1"
+                },
+                {
+                    "check": "body.foo2",
+                    "assert": "equals",
+                    "expect": 12.3,
+                    "msg": "check json foo2"
+                }
+            ]
+        },
+        {
+            "name": "write and read text",
+            "websocket": {
+                "type": "wr",
+                "url": "/",
+                "text": "$varFoo1"
+            },
+            "validate": [
+                {
+                    "check": "body",
+                    "assert": "length_equals",
+                    "expect": 5,
+                    "msg": "check length equal"
+                }
+            ]
+        },
+        {
+            "name": "write something redundant",
+            "websocket": {
+                "type": "w",
+                "url": "/",
+                "text": "have a nice day!"
+            }
+        },
+        {
+            "name": "write something redundant",
+            "websocket": {
+                "type": "w",
+                "url": "/",
+                "text": "balabala ..."
+            }
+        },
+        {
+            "name": "close connection",
+            "websocket": {
+                "type": "close",
+                "url": "/",
+                "close_status": 1000,
+                "timeout": 30000
+            },
+            "validate": [
+                {
+                    "check": "status_code",
+                    "assert": "equals",
+                    "expect": 1000,
+                    "msg": "check close status code"
+                }
+            ]
+        }
+    ]
+}

go.mod

@@ -7,6 +7,7 @@ require (
 	github.com/denisbrodbeck/machineid v1.0.1
 	github.com/getsentry/sentry-go v0.13.0
 	github.com/google/uuid v1.3.0
+	github.com/gorilla/websocket v1.4.1
 	github.com/httprunner/funplugin v0.4.2
 	github.com/jinzhu/copier v0.3.2
 	github.com/jmespath/go-jmespath v0.4.0

go.sum

@@ -210,6 +210,7 @@ github.com/google/uuid v1.3.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+
 github.com/googleapis/gax-go/v2 v2.0.4/go.mod h1:0Wqv26UfaUD9n4G6kQubkQ+KchISgw+vpHVxEJEs9eg=
 github.com/googleapis/gax-go/v2 v2.0.5/go.mod h1:DWXyrwAJ9X0FpwwEdw+IPEYBICEFu5mhpdKc/us6bOk=
 github.com/gopherjs/gopherjs v0.0.0-20181017120253-0766667cb4d1/go.mod h1:wJfORRmW1u3UXTncJ5qlYoELFm8eSnnEO6hX4iZ3EWY=
+github.com/gorilla/websocket v1.4.1 h1:q7AeDBpnBk8AogcD4DSag/Ukw/KV+YhzLj2bP5HvKCM=
 github.com/gorilla/websocket v1.4.1/go.mod h1:YR8l580nyteQvAITg2hZ9XVh4b55+EU/adAjf1fMHhE=
 github.com/grpc-ecosystem/grpc-gateway v1.16.0/go.mod h1:BDjrQk3hbvj6Nolgz8mAMFbcEtjT1g+wF4CSlocrBnw=
 github.com/hashicorp/consul/api v1.1.0/go.mod h1:VmuI/Lkw1nC05EYQWNKwWGbkg+FbDBtguAZLlVdkD9Q=

hrp/config.go

@@ -25,6 +25,7 @@ type TConfig struct {
 	Parameters        map[string]interface{} `json:"parameters,omitempty" yaml:"parameters,omitempty"`
 	ParametersSetting *TParamsConfig         `json:"parameters_setting,omitempty" yaml:"parameters_setting,omitempty"`
 	ThinkTimeSetting  *ThinkTimeConfig       `json:"think_time,omitempty" yaml:"think_time,omitempty"`
+	WebSocketSetting  *WebSocketConfig       `json:"websocket,omitempty" yaml:"websocket,omitempty"`
 	Export            []string               `json:"export,omitempty" yaml:"export,omitempty"`
 	Weight            int                    `json:"weight,omitempty" yaml:"weight,omitempty"`
 	Path              string                 `json:"path,omitempty" yaml:"path,omitempty"` // testcase file path
@@ -78,6 +79,14 @@ func (c *TConfig) SetWeight(weight int) *TConfig {
 	return c
 }
 
+func (c *TConfig) SetWebSocket(times, interval, timeout, size int64) {
+	c.WebSocketSetting = &WebSocketConfig{
+		ReconnectionTimes:    times,
+		ReconnectionInterval: interval,
+		MaxMessageSize:       size,
+	}
+}
+
 type ThinkTimeConfig struct {
 	Strategy thinkTimeStrategy `json:"strategy,omitempty" yaml:"strategy,omitempty"` // default、random、limit、multiply、ignore
 	Setting  interface{}       `json:"setting,omitempty" yaml:"setting,omitempty"`   // random(map): {"min_percentage": 0.5, "max_percentage": 1.5}; 10、multiply(float64): 1.5

hrp/internal/builtin/function.go

@@ -16,6 +16,7 @@ var Functions = map[string]interface{}{
 	"md5":               MD5,             // call with one argument
 	"parameterize":      loadFromCSV,
 	"P":                 loadFromCSV,
+	"load_ws_message":   loadMessage,
 }
 
 func init() {

hrp/internal/builtin/utils.go

@@ -282,6 +282,16 @@ func loadFromCSV(path string) []map[string]interface{} {
 	return result
 }
 
+func loadMessage(path string) []byte {
+	log.Info().Str("path", path).Msg("load message file")
+	file, err := readFile(path)
+	if err != nil {
+		log.Error().Err(err).Msg("read message file failed")
+		os.Exit(1)
+	}
+	return file
+}
+
 func readFile(path string) ([]byte, error) {
 	var err error
 	path, err = filepath.Abs(path)

hrp/response.go

@@ -18,7 +18,17 @@ import (
 	"github.com/httprunner/httprunner/hrp/internal/json"
 )
 
-func newResponseObject(t *testing.T, parser *Parser, resp *http.Response) (*responseObject, error) {
+var fieldTags = []string{"proto", "status_code", "headers", "cookies", "body", textExtractorSubRegexp}
+
+type httpRespObjMeta struct {
+	Proto      string            `json:"proto"`
+	StatusCode int               `json:"status_code"`
+	Headers    map[string]string `json:"headers"`
+	Cookies    map[string]string `json:"cookies"`
+	Body       interface{}       `json:"body"`
+}
+
+func newHttpResponseObject(t *testing.T, parser *Parser, resp *http.Response) (*responseObject, error) {
 	// prepare response headers
 	headers := make(map[string]string)
 	for k, v := range resp.Header {
@@ -46,7 +56,7 @@ func newResponseObject(t *testing.T, parser *Parser, resp *http.Response) (*resp
 		body = string(respBodyBytes)
 	}
 
-	respObjMeta := respObjMeta{
+	respObjMeta := httpRespObjMeta{
 		Proto:      resp.Proto,
 		StatusCode: resp.StatusCode,
 		Headers:    headers,
@@ -54,19 +64,49 @@ func newResponseObject(t *testing.T, parser *Parser, resp *http.Response) (*resp
 		Body:       body,
 	}
 
-	// convert respObjMeta to interface{}
+	return convertToResponseObject(t, parser, respObjMeta)
+}
+
+type wsCloseRespObject struct {
+	StatusCode int    `json:"status_code"`
+	Text       string `json:"body"`
+}
+
+func newWsCloseResponseObject(t *testing.T, parser *Parser, resp *wsCloseRespObject) (*responseObject, error) {
+	return convertToResponseObject(t, parser, resp)
+}
+
+type wsReadRespObject struct {
+	Message     interface{} `json:"body"`
+	messageType int
+}
+
+func newWsReadResponseObject(t *testing.T, parser *Parser, resp *wsReadRespObject) (*responseObject, error) {
+	byteMessage, ok := resp.Message.([]byte)
+	if !ok {
+		return nil, errors.New("websocket message type should be []byte")
+	}
+	var msg interface{}
+	if err := json.Unmarshal(byteMessage, &msg); err != nil {
+		// response body is not json, use raw body
+		msg = string(byteMessage)
+	}
+	resp.Message = msg
+	return convertToResponseObject(t, parser, resp)
+}
+
+func convertToResponseObject(t *testing.T, parser *Parser, respObjMeta interface{}) (*responseObject, error) {
 	respObjMetaBytes, _ := json.Marshal(respObjMeta)
 	var data interface{}
 	decoder := json.NewDecoder(bytes.NewReader(respObjMetaBytes))
 	decoder.UseNumber()
 	if err := decoder.Decode(&data); err != nil {
 		log.Error().
-			Str("respObjMeta", string(respObjMetaBytes)).
+			Str("respObjectMeta", string(respObjMetaBytes)).
 			Err(err).
-			Msg("[NewResponseObject] convert respObjMeta to interface{} failed")
+			Msg("[convertToResponseObject] convert respObjectMeta to interface{} failed")
 		return nil, err
 	}
-
 	return &responseObject{
 		t:           t,
 		parser:      parser,
@@ -74,14 +114,6 @@ func newResponseObject(t *testing.T, parser *Parser, resp *http.Response) (*resp
 	}, nil
 }
 
-type respObjMeta struct {
-	Proto      string            `json:"proto"`
-	StatusCode int               `json:"status_code"`
-	Headers    map[string]string `json:"headers"`
-	Cookies    map[string]string `json:"cookies"`
-	Body       interface{}       `json:"body"`
-}
-
 type responseObject struct {
 	t                 *testing.T
 	parser            *Parser
@@ -194,12 +226,12 @@ func (v *responseObject) Validate(iValidators []interface{}, variablesMapping ma
 }
 
 func checkSearchField(expr string) bool {
-	return strings.Contains(expr, "proto") ||
+	for _, t := range fieldTags {
-		strings.Contains(expr, "status_code") ||
+		if strings.Contains(expr, t) {
-		strings.Contains(expr, "headers") ||
+			return true
-		strings.Contains(expr, "cookies") ||
+		}
-		strings.Contains(expr, "body") ||
+	}
-		strings.Contains(expr, textExtractorSubRegexp)
+	return false
 }
 
 func (v *responseObject) searchJmespath(expr string) interface{} {

hrp/response_test.go

@@ -22,7 +22,7 @@ func TestSearchJmespath(t *testing.T) {
 	}
 	resp := http.Response{}
 	resp.Body = io.NopCloser(strings.NewReader(testText))
-	respObj, err := newResponseObject(t, newParser(), &resp)
+	respObj, err := newHttpResponseObject(t, newParser(), &resp)
 	if err != nil {
 		t.Fatal()
 	}
@@ -47,7 +47,7 @@ func TestSearchRegexp(t *testing.T) {
 	// new response object
 	resp := http.Response{}
 	resp.Body = io.NopCloser(strings.NewReader(testText))
-	respObj, err := newResponseObject(t, newParser(), &resp)
+	respObj, err := newHttpResponseObject(t, newParser(), &resp)
 	if err != nil {
 		t.Fatal()
 	}

hrp/runner.go

@@ -8,6 +8,7 @@ import (
 	"testing"
 	"time"
 
+	"github.com/gorilla/websocket"
 	"github.com/jinzhu/copier"
 	"github.com/pkg/errors"
 	"github.com/rs/zerolog/log"
@@ -43,6 +44,10 @@ func NewRunner(t *testing.T) *HRPRunner {
 			},
 			Timeout: 30 * time.Second,
 		},
+		// use default handshake timeout (no timeout limit) here, enable timeout at step level
+		wsDialer: &websocket.Dialer{
+			TLSClientConfig: &tls.Config{InsecureSkipVerify: true},
+		},
 	}
 }
 
@@ -55,6 +60,7 @@ type HRPRunner struct {
 	genHTMLReport bool
 	httpClient    *http.Client
 	http2Client   *http.Client
+	wsDialer      *websocket.Dialer
 }
 
 // SetClientTransport configures transport of http client for high concurrency load testing
@@ -76,6 +82,7 @@ func (r *HRPRunner) SetClientTransport(maxConns int, disableKeepAlive bool, disa
 		TLSClientConfig:    &tls.Config{InsecureSkipVerify: true},
 		DisableCompression: disableCompression,
 	}
+	r.wsDialer.EnableCompression = !disableCompression
 	return r
 }
 
@@ -112,6 +119,7 @@ func (r *HRPRunner) SetProxyUrl(proxyUrl string) *HRPRunner {
 		Proxy:           http.ProxyURL(p),
 		TLSClientConfig: &tls.Config{InsecureSkipVerify: true},
 	}
+	r.wsDialer.Proxy = http.ProxyURL(p)
 	return r
 }
 
@@ -273,6 +281,9 @@ func (r *testCaseRunner) parseConfig() error {
 	// ensure correction of think time config
 	r.parsedConfig.ThinkTimeSetting.checkThinkTime()
 
+	// ensure correction of websocket config
+	r.parsedConfig.WebSocketSetting.checkWebSocket()
+
 	// parse testcase config parameters
 	parametersIterator, err := initParametersIterator(r.parsedConfig)
 	if err != nil {

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
+	transactions      map[string]map[transactionType]time.Time
-	startTime    time.Time        // record start time of the testcase
+	startTime         time.Time               // record start time of the testcase
-	summary      *TestCaseSummary // record test case summary
+	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
 }

hrp/step.go

@@ -9,6 +9,7 @@ const (
 	stepTypeTransaction StepType = "transaction"
 	stepTypeRendezvous  StepType = "rendezvous"
 	stepTypeThinkTime   StepType = "thinktime"
+	stepTypeWebSocket   StepType = "websocket"
 )
 
 type StepResult struct {
@@ -32,6 +33,7 @@ type TStep struct {
 	Transaction   *Transaction           `json:"transaction,omitempty" yaml:"transaction,omitempty"`
 	Rendezvous    *Rendezvous            `json:"rendezvous,omitempty" yaml:"rendezvous,omitempty"`
 	ThinkTime     *ThinkTime             `json:"think_time,omitempty" yaml:"think_time,omitempty"`
+	WebSocket     *WebSocketAction       `json:"websocket,omitempty" yaml:"websocket,omitempty"`
 	Variables     map[string]interface{} `json:"variables,omitempty" yaml:"variables,omitempty"`
 	SetupHooks    []string               `json:"setup_hooks,omitempty" yaml:"setup_hooks,omitempty"`
 	TeardownHooks []string               `json:"teardown_hooks,omitempty" yaml:"teardown_hooks,omitempty"`
@@ -43,7 +45,7 @@ type TStep struct {
 // IStep represents interface for all types for teststeps, includes:
 // StepRequest, StepRequestWithOptionalArgs, StepRequestValidation, StepRequestExtraction,
 // StepTestCaseWithOptionalArgs,
-// StepTransaction, StepRendezvous.
+// StepTransaction, StepRendezvous, StepWebSocket.
 type IStep interface {
 	Name() string
 	Type() StepType

hrp/step_request.go

@@ -340,7 +340,7 @@ func runStepRequest(r *SessionRunner, step *TStep) (stepResult *StepResult, err
 	}
 
 	// new response object
-	respObj, err := newResponseObject(r.hrpRunner.t, parser, resp)
+	respObj, err := newHttpResponseObject(r.hrpRunner.t, parser, resp)
 	if err != nil {
 		err = errors.Wrap(err, "init ResponseObject error")
 		return
@@ -401,7 +401,7 @@ func printRequest(req *http.Request) error {
 	}
 	fmt.Println("-------------------- request --------------------")
 	reqContent := string(reqDump)
-	if req.Body != nil && !printBody {
+	if reqContentType != "" && !printBody {
 		reqContent += fmt.Sprintf("(request body omitted for Content-Type: %v)", reqContentType)
 	}
 	fmt.Println(reqContent)
@@ -409,7 +409,7 @@ func printRequest(req *http.Request) error {
 }
 
 func printResponse(resp *http.Response) error {
-	fmt.Println("==================== response ===================")
+	fmt.Println("==================== response ====================")
 	respContentType := resp.Header.Get("Content-Type")
 	printBody := shouldPrintBody(respContentType)
 	respDump, err := httputil.DumpResponse(resp, printBody)
@@ -417,7 +417,7 @@ func printResponse(resp *http.Response) error {
 		return errors.Wrap(err, "dump response failed")
 	}
 	respContent := string(respDump)
-	if !printBody {
+	if respContentType != "" && !printBody {
 		respContent += fmt.Sprintf("(response body omitted for Content-Type: %v)", respContentType)
 	}
 	fmt.Println(respContent)
@@ -677,6 +677,14 @@ func (s *StepRequest) SetRendezvous(name string) *StepRendezvous {
 	}
 }
 
+// WebSocket creates a new websocket action
+func (s *StepRequest) WebSocket() *StepWebSocket {
+	s.step.WebSocket = &WebSocketAction{}
+	return &StepWebSocket{
+		step: s.step,
+	}
+}
+
 // StepRequestWithOptionalArgs implements IStep interface.
 type StepRequestWithOptionalArgs struct {
 	step *TStep
@@ -799,7 +807,13 @@ func (s *StepRequestExtraction) Name() string {
 }
 
 func (s *StepRequestExtraction) Type() StepType {
-	return StepType(fmt.Sprintf("request-%v", s.step.Request.Method))
+	if s.step.Request != nil {
+		return StepType(fmt.Sprintf("request-%v", s.step.Request.Method))
+	}
+	if s.step.WebSocket != nil {
+		return StepType(fmt.Sprintf("websocket-%v", s.step.WebSocket.Type))
+	}
+	return "extraction"
 }
 
 func (s *StepRequestExtraction) Struct() *TStep {
@@ -807,7 +821,13 @@ func (s *StepRequestExtraction) Struct() *TStep {
 }
 
 func (s *StepRequestExtraction) Run(r *SessionRunner) (*StepResult, error) {
-	return runStepRequest(r, s.step)
+	if s.step.Request != nil {
+		return runStepRequest(r, s.step)
+	}
+	if s.step.WebSocket != nil {
+		return runStepWebSocket(r, s.step)
+	}
+	return nil, errors.New("unexpected protocol type")
 }
 
 // StepRequestValidation implements IStep interface.
@@ -823,7 +843,13 @@ func (s *StepRequestValidation) Name() string {
 }
 
 func (s *StepRequestValidation) Type() StepType {
-	return StepType(fmt.Sprintf("request-%v", s.step.Request.Method))
+	if s.step.Request != nil {
+		return StepType(fmt.Sprintf("request-%v", s.step.Request.Method))
+	}
+	if s.step.WebSocket != nil {
+		return StepType(fmt.Sprintf("websocket-%v", s.step.WebSocket.Type))
+	}
+	return "validation"
 }
 
 func (s *StepRequestValidation) Struct() *TStep {
@@ -831,7 +857,13 @@ func (s *StepRequestValidation) Struct() *TStep {
 }
 
 func (s *StepRequestValidation) Run(r *SessionRunner) (*StepResult, error) {
-	return runStepRequest(r, s.step)
+	if s.step.Request != nil {
+		return runStepRequest(r, s.step)
+	}
+	if s.step.WebSocket != nil {
+		return runStepWebSocket(r, s.step)
+	}
+	return nil, errors.New("unexpected protocol type")
 }
 
 func (s *StepRequestValidation) AssertEqual(jmesPath string, expected interface{}, msg string) *StepRequestValidation {

hrp/step_websocket.go

@@ -0,0 +1,683 @@
+package hrp
+
+import (
+	"bytes"
+	"fmt"
+	"github.com/gorilla/websocket"
+	"github.com/httprunner/httprunner/hrp/internal/builtin"
+	"github.com/httprunner/httprunner/hrp/internal/json"
+	"github.com/pkg/errors"
+	"github.com/rs/zerolog/log"
+	"net/http"
+	"testing"
+	"time"
+	"unsafe"
+)
+
+const (
+	wsOpen         ActionType = "open"
+	wsPing         ActionType = "ping"
+	wsWriteAndRead ActionType = "wr"
+	wsRead         ActionType = "r"
+	wsWrite        ActionType = "w"
+	wsClose        ActionType = "close"
+)
+
+const (
+	defaultTimeout     = 30000                        // default timeout 30 seconds for open, read and close
+	defaultCloseStatus = websocket.CloseNormalClosure // default normal close status
+	defaultWriteWait   = 5 * time.Second              // default timeout 5 seconds for writing control message
+)
+
+type ActionType string
+
+func (at ActionType) toString() string {
+	switch at {
+	case wsOpen:
+		return "open new connection"
+	case wsPing:
+		return "send ping and expect pong"
+	case wsWriteAndRead:
+		return "write and read"
+	case wsRead:
+		return "read only"
+	case wsWrite:
+		return "write only"
+	case wsClose:
+		return "close current connection"
+	default:
+		return "unexpected action type"
+	}
+}
+
+type MessageType int
+
+func (mt MessageType) toString() string {
+	switch mt {
+	case websocket.TextMessage:
+		return "text"
+	case websocket.BinaryMessage:
+		return "binary"
+	case websocket.PingMessage:
+		return "ping"
+	case websocket.PongMessage:
+		return "pong"
+	case websocket.CloseMessage:
+		return "close"
+	case 0:
+		return "continuation"
+	case -1:
+		return "no frame"
+	default:
+		return "unexpected message type"
+	}
+}
+
+// WebSocketConfig TODO: support reconnection ability
+type WebSocketConfig struct {
+	ReconnectionTimes    int64 `json:"reconnection_times,omitempty" yaml:"reconnection_times,omitempty"`       // maximum reconnection times when the connection closed by remote server
+	ReconnectionInterval int64 `json:"reconnection_interval,omitempty" yaml:"reconnection_interval,omitempty"` // interval between each reconnection in milliseconds
+	MaxMessageSize       int64 `json:"max_message_size,omitempty" yaml:"max_message_size,omitempty"`           // maximum message size during writing/reading
+}
+
+const (
+	defaultReconnectionTimes    = 0
+	defaultReconnectionInterval = 5000
+	defaultMaxMessageSize       = 0
+)
+
+// checkWebSocket validates the websocket configuration parameters
+func (wsConfig *WebSocketConfig) checkWebSocket() {
+	if wsConfig == nil {
+		return
+	}
+	if wsConfig.ReconnectionTimes <= 0 {
+		wsConfig.ReconnectionTimes = defaultReconnectionTimes
+	}
+	if wsConfig.ReconnectionInterval <= 0 {
+		wsConfig.ReconnectionInterval = defaultReconnectionInterval
+	}
+	if wsConfig.MaxMessageSize <= 0 {
+		wsConfig.MaxMessageSize = defaultMaxMessageSize
+	}
+}
+
+// StepWebSocket implements IStep interface.
+type StepWebSocket struct {
+	step *TStep
+}
+
+func (s *StepWebSocket) Name() string {
+	if s.step.Name != "" {
+		return s.step.Name
+	}
+	return fmt.Sprintf("%s %s", s.step.WebSocket.Type, s.step.WebSocket.URL)
+}
+
+func (s *StepWebSocket) Type() StepType {
+	return StepType(fmt.Sprintf("websocket-%v", s.step.WebSocket.Type))
+}
+
+func (s *StepWebSocket) Struct() *TStep {
+	return s.step
+}
+
+func (s *StepWebSocket) Run(r *SessionRunner) (*StepResult, error) {
+	return runStepWebSocket(r, s.step)
+}
+
+func (s *StepWebSocket) OpenConnection(url string) *StepWebSocket {
+	s.step.WebSocket.Type = wsOpen
+	s.step.WebSocket.URL = url
+	return s
+}
+
+func (s *StepWebSocket) PingPong(url string) *StepWebSocket {
+	s.step.WebSocket.Type = wsPing
+	s.step.WebSocket.URL = url
+	return s
+}
+
+func (s *StepWebSocket) WriteAndRead(url string) *StepWebSocket {
+	s.step.WebSocket.Type = wsWriteAndRead
+	s.step.WebSocket.URL = url
+	return s
+}
+
+func (s *StepWebSocket) Read(url string) *StepWebSocket {
+	s.step.WebSocket.Type = wsRead
+	s.step.WebSocket.URL = url
+	return s
+}
+
+func (s *StepWebSocket) Write(url string) *StepWebSocket {
+	s.step.WebSocket.Type = wsWrite
+	s.step.WebSocket.URL = url
+	return s
+}
+
+func (s *StepWebSocket) CloseConnection(url string) *StepWebSocket {
+	s.step.WebSocket.Type = wsClose
+	s.step.WebSocket.URL = url
+	return s
+}
+
+func (s *StepWebSocket) WithParams(params map[string]interface{}) *StepWebSocket {
+	s.step.WebSocket.Params = params
+	return s
+}
+
+func (s *StepWebSocket) WithHeaders(headers map[string]string) *StepWebSocket {
+	s.step.WebSocket.Headers = headers
+	return s
+}
+
+func (s *StepWebSocket) NewConnection() *StepWebSocket {
+	s.step.WebSocket.NewConnection = true
+	return s
+}
+
+func (s *StepWebSocket) WithTextMessage(message interface{}) *StepWebSocket {
+	s.step.WebSocket.TextMessage = message
+	return s
+}
+
+func (s *StepWebSocket) WithBinaryMessage(message interface{}) *StepWebSocket {
+	s.step.WebSocket.BinaryMessage = message
+	return s
+}
+
+func (s *StepWebSocket) WithTimeout(timeout int64) *StepWebSocket {
+	s.step.WebSocket.Timeout = timeout
+	return s
+}
+
+func (s *StepWebSocket) WithCloseStatus(closeStatus int64) *StepWebSocket {
+	s.step.WebSocket.CloseStatusCode = closeStatus
+	return s
+}
+
+// Validate switches to step validation.
+func (s *StepWebSocket) Validate() *StepRequestValidation {
+	return &StepRequestValidation{
+		step: s.step,
+	}
+}
+
+// Extract switches to step extraction.
+func (s *StepWebSocket) Extract() *StepRequestExtraction {
+	s.step.Extract = make(map[string]string)
+	return &StepRequestExtraction{
+		step: s.step,
+	}
+}
+
+type WebSocketAction struct {
+	Type            ActionType             `json:"type" yaml:"type"`
+	URL             string                 `json:"url" yaml:"url"`
+	Params          map[string]interface{} `json:"params,omitempty" yaml:"params,omitempty"`
+	Headers         map[string]string      `json:"headers,omitempty" yaml:"headers,omitempty"`
+	NewConnection   bool                   `json:"new_connection,omitempty" yaml:"new_connection,omitempty"` // TODO support
+	TextMessage     interface{}            `json:"text,omitempty" yaml:"text,omitempty"`
+	BinaryMessage   interface{}            `json:"binary,omitempty" yaml:"binary,omitempty"`
+	CloseStatusCode int64                  `json:"close_status,omitempty" yaml:"close_status,omitempty"`
+	Timeout         int64                  `json:"timeout,omitempty" yaml:"timeout,omitempty"`
+}
+
+func runStepWebSocket(r *SessionRunner, step *TStep) (stepResult *StepResult, err error) {
+	stepResult = &StepResult{
+		Name:        step.Name,
+		StepType:    stepTypeWebSocket,
+		Success:     false,
+		ContentSize: 0,
+	}
+
+	defer func() {
+		// update testcase summary
+		if err != nil {
+			stepResult.Attachment = err.Error()
+		}
+	}()
+
+	// override step variables
+	stepVariables, err := r.MergeStepVariables(step.Variables)
+	if err != nil {
+		return
+	}
+
+	sessionData := newSessionData()
+	parser := r.GetParser()
+
+	// add request object to step variables, could be used in setup hooks
+	stepVariables["hrp_step_name"] = step.Name
+
+	// deal with setup hooks
+	for _, setupHook := range step.SetupHooks {
+		_, err = parser.Parse(setupHook, stepVariables)
+		if err != nil {
+			return stepResult, errors.Wrap(err, "run setup hooks failed")
+		}
+	}
+
+	// init websocket frame parameters
+	initStepWebsocket(step.WebSocket)
+	var resp interface{}
+	start := time.Now()
+
+	// do websocket action
+	if r.LogOn() {
+		fmt.Printf("-------------------- websocket action: %v --------------------\n", step.WebSocket.Type.toString())
+	}
+	switch step.WebSocket.Type {
+	case wsOpen:
+		log.Info().Int64("timeout(ms)", step.WebSocket.Timeout).Msg("open websocket connection")
+		// use the current websocket connection if existed
+		if r.wsConn != nil {
+			break
+		}
+		resp, err = openWithTimeout(sessionData, r, step, stepVariables)
+		if err != nil {
+			return stepResult, errors.Wrap(err, "open connection failed")
+		}
+	case wsPing:
+		log.Info().Int64("timeout(ms)", step.WebSocket.Timeout).Msg("send ping and expect pong")
+		if r.wsConn == nil {
+			return stepResult, errors.Errorf("try to use existing connection, but there is no connection")
+		}
+		err = writeWebSocket(r, step, stepVariables)
+		if err != nil {
+			return stepResult, errors.Wrap(err, "send ping message failed")
+		}
+		timer := time.NewTimer(time.Duration(step.WebSocket.Timeout) * time.Millisecond)
+		// asynchronous receiving pong message with timeout
+		go func() {
+			select {
+			case <-timer.C:
+				timer.Stop()
+				log.Warn().Msg("pong timeout")
+			case pongResponse := <-r.pongResponseChan:
+				resp = pongResponse
+				log.Info().Msg("pong message arrives")
+			}
+		}()
+	case wsWriteAndRead:
+		log.Info().Int64("timeout(ms)", step.WebSocket.Timeout).Msg("write a message and read response")
+		err = writeWebSocket(r, step, stepVariables)
+		if err != nil {
+			return stepResult, errors.Wrap(err, "write message failed")
+		}
+		resp, err = readMessageWithTimeout(r, step)
+		if err != nil {
+			return stepResult, errors.Wrap(err, "read message failed")
+		}
+	case wsRead:
+		log.Info().Int64("timeout(ms)", step.WebSocket.Timeout).Msg("read only")
+		resp, err = readMessageWithTimeout(r, step)
+		if err != nil {
+			return stepResult, errors.Wrap(err, "read message failed")
+		}
+	case wsWrite:
+		log.Info().Msg("write only")
+		err = writeWebSocket(r, step, stepVariables)
+		if err != nil {
+			return stepResult, errors.Wrap(err, "write message failed")
+		}
+	case wsClose:
+		log.Info().Int64("timeout(ms)", step.WebSocket.Timeout).Msg("close webSocket connection")
+		resp, err = closeWithTimeout(r, step, stepVariables)
+		if err != nil {
+			return stepResult, errors.Wrap(err, "close connection failed")
+		}
+	default:
+		return stepResult, errors.Errorf("unexpected websocket frame type: %v", step.WebSocket.Type)
+	}
+	if r.LogOn() {
+		err = printWebSocketResponse(resp)
+		if err != nil {
+			return stepResult, errors.Wrap(err, "print response failed")
+		}
+	}
+
+	stepResult.Elapsed = time.Since(start).Milliseconds()
+	respObj, err := getResponseObject(r.hrpRunner.t, r.parser, resp)
+	if err != nil {
+		err = errors.Wrap(err, "get response object error")
+		return
+	}
+
+	if respObj != nil {
+		// add response object to step variables, could be used in teardown hooks
+		stepVariables["hrp_step_response"] = respObj.respObjMeta
+	}
+
+	// deal with teardown hooks
+	for _, teardownHook := range step.TeardownHooks {
+		_, err = parser.Parse(teardownHook, stepVariables)
+		if err != nil {
+			return stepResult, errors.Wrap(err, "run teardown hooks failed")
+		}
+	}
+
+	if respObj != nil {
+		sessionData.ReqResps.Response = builtin.FormatResponse(respObj.respObjMeta)
+
+		// extract variables from response
+		extractors := step.Extract
+		extractMapping := respObj.Extract(extractors)
+		stepResult.ExportVars = extractMapping
+
+		// override step variables with extracted variables
+		stepVariables = mergeVariables(stepVariables, extractMapping)
+
+		// validate response
+		err = respObj.Validate(step.Validators, stepVariables)
+		sessionData.Validators = respObj.validationResults
+		if err == nil {
+			sessionData.Success = true
+			stepResult.Success = true
+		}
+		stepResult.ContentSize = getContentSize(resp)
+		stepResult.Data = sessionData
+	} else {
+		sessionData.Success = true
+		stepResult.Success = true
+	}
+
+	// update summary
+	r.summary.Records = append(r.summary.Records, stepResult)
+	r.summary.Stat.Total += 1
+	if stepResult.Success {
+		r.summary.Stat.Successes += 1
+	} else {
+		r.summary.Stat.Failures += 1
+		// update summary result to failed
+		r.summary.Success = false
+	}
+	return stepResult, nil
+}
+
+func printWebSocketResponse(resp interface{}) error {
+	if resp == nil {
+		fmt.Println("(response body is empty in this step)")
+		fmt.Println("----------------------------------------")
+		return nil
+	}
+	if httpResp, ok := resp.(*http.Response); ok {
+		return printResponse(httpResp)
+	}
+	fmt.Println("==================== response ====================")
+	switch r := resp.(type) {
+	case *wsReadRespObject:
+		if r.messageType == websocket.TextMessage {
+			fmt.Printf("message type: %v\r\nmessage: %s\r\n", MessageType(r.messageType).toString(), r.Message)
+		} else if r.messageType == websocket.BinaryMessage {
+			fmt.Printf("message type: %v\r\nmessage: %v\r\ncorresponding string: %s\r\n", MessageType(r.messageType).toString(), r.Message, r.Message)
+		} else {
+			return errors.New("unexpected response type")
+		}
+	case *wsCloseRespObject:
+		fmt.Printf("close status code: %v\r\nmessage: %v\r\n", r.StatusCode, r.Text)
+	case string:
+		fmt.Println(r)
+	default:
+		return errors.New("unexpected response type")
+	}
+	fmt.Println("----------------------------------------")
+	return nil
+}
+
+func initStepWebsocket(stepWebSocket *WebSocketAction) {
+	if stepWebSocket == nil {
+		return
+	}
+	if stepWebSocket.Timeout <= 0 {
+		stepWebSocket.Timeout = defaultTimeout
+	}
+	// close status code range: [1000, 4999]. ref: https://datatracker.ietf.org/doc/html/rfc6455#section-11.7
+	if stepWebSocket.CloseStatusCode < 1000 || stepWebSocket.CloseStatusCode > 4999 {
+		stepWebSocket.CloseStatusCode = defaultCloseStatus
+	}
+}
+
+// prepareDialInfo prepares url and headers before opening connection
+func prepareDialInfo(r *SessionRunner, step *TStep, stepVariables map[string]interface{}) (*requestBuilder, error) {
+	parser := r.GetParser()
+	config := r.GetConfig()
+
+	dummyReq := &Request{
+		URL:     step.WebSocket.URL,
+		Params:  step.WebSocket.Params,
+		Headers: step.WebSocket.Headers,
+	}
+	dummyBuilder := newRequestBuilder(parser, config, dummyReq)
+
+	err := dummyBuilder.prepareUrlParams(stepVariables)
+	if err != nil {
+		return nil, err
+	}
+
+	err = dummyBuilder.prepareHeaders(stepVariables)
+	if err != nil {
+		return nil, err
+	}
+
+	return dummyBuilder, nil
+}
+
+func openWithTimeout(d *SessionData, r *SessionRunner, step *TStep, stepVariables map[string]interface{}) (*http.Response, error) {
+	openResponseChan := make(chan *http.Response)
+	errorChan := make(chan error)
+	go func() {
+		// prepare request and dial
+		rb, err := prepareDialInfo(r, step, stepVariables)
+		if err != nil {
+			errorChan <- errors.Wrap(err, "prepare dail info failed")
+			return
+		}
+		d.ReqResps.Request = rb.requestMap
+		conn, resp, err := r.hrpRunner.wsDialer.Dial(rb.req.URL.String(), rb.req.Header)
+		if err != nil {
+			errorChan <- errors.Wrap(err, "dial tcp failed")
+			return
+		}
+		// handshake end here
+		defer resp.Body.Close()
+
+		// the following handlers handle and transport control message from server
+		conn.SetPongHandler(func(appData string) error {
+			r.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))
+			r.closeResponseChan <- &wsCloseRespObject{
+				StatusCode: code,
+				Text:       text,
+			}
+			return nil
+		})
+		r.wsConn = conn
+		openResponseChan <- resp
+	}()
+
+	timer := time.NewTimer(time.Duration(step.WebSocket.Timeout) * time.Millisecond)
+	select {
+	case <-timer.C:
+		timer.Stop()
+		return nil, errors.New("open timeout")
+	case err := <-errorChan:
+		return nil, err
+	case openResponse := <-openResponseChan:
+		return openResponse, nil
+	}
+}
+
+func readMessageWithTimeout(r *SessionRunner, step *TStep) (*wsReadRespObject, error) {
+	if r.wsConn == nil {
+		return nil, errors.New("try to use existing connection, but there is no connection")
+	}
+	readResponseChan := make(chan *wsReadRespObject)
+	errorChan := make(chan error)
+	go func() {
+		messageType, message, err := r.wsConn.ReadMessage()
+		if err != nil {
+			errorChan <- err
+		} else {
+			readResponseChan <- &wsReadRespObject{
+				messageType: messageType,
+				Message:     message,
+			}
+		}
+	}()
+	timer := time.NewTimer(time.Duration(step.WebSocket.Timeout) * time.Millisecond)
+	select {
+	case <-timer.C:
+		timer.Stop()
+		return nil, errors.New("read timeout")
+	case err := <-errorChan:
+		return nil, err
+	case readResult := <-readResponseChan:
+		return readResult, nil
+	}
+}
+
+func writeWebSocket(r *SessionRunner, step *TStep, stepVariables map[string]interface{}) error {
+	if r.wsConn == nil {
+		return errors.New("try to use existing connection, but there is no connection")
+	}
+	// TODO: only support writing one kind of message each step here?
+	// check priority: text message > binary message
+	if step.WebSocket.TextMessage != nil {
+		parsedMessage, parseErr := r.parser.Parse(step.WebSocket.TextMessage, stepVariables)
+		if parseErr != nil {
+			return parseErr
+		}
+		writeErr := writeWithType(r.wsConn, step, websocket.TextMessage, parsedMessage)
+		if writeErr != nil {
+			return writeErr
+		}
+	} else if step.WebSocket.BinaryMessage != nil {
+		parsedMessage, parseErr := r.parser.Parse(step.WebSocket.BinaryMessage, stepVariables)
+		if parseErr != nil {
+			return parseErr
+		}
+		writeErr := writeWithType(r.wsConn, step, websocket.BinaryMessage, parsedMessage)
+		if writeErr != nil {
+			return writeErr
+		}
+	} else {
+		log.Info().Msg("step with empty message")
+		err := writeWithAction(r.wsConn, step, websocket.BinaryMessage, []byte{})
+		if err != nil {
+			return err
+		}
+	}
+	return nil
+}
+
+func writeWithType(c *websocket.Conn, step *TStep, messageType int, message interface{}) error {
+	if message == nil {
+		return nil
+	}
+	if messageType != websocket.TextMessage && messageType != websocket.BinaryMessage {
+		return errors.New("unexpected message type")
+	}
+	switch msg := message.(type) {
+	case []byte:
+		return writeWithAction(c, step, messageType, msg)
+	case string:
+		return writeWithAction(c, step, messageType, []byte(msg))
+	case bytes.Buffer:
+		return writeWithAction(c, step, messageType, msg.Bytes())
+	default:
+		msgBytes, _ := json.Marshal(msg)
+		return writeWithAction(c, step, messageType, msgBytes)
+	}
+}
+
+func writeWithAction(c *websocket.Conn, step *TStep, messageType int, message []byte) error {
+	switch step.WebSocket.Type {
+	case wsPing:
+		return c.WriteControl(websocket.PingMessage, message, time.Now().Add(defaultWriteWait))
+	case wsClose:
+		closeMessage := websocket.FormatCloseMessage(int(step.WebSocket.CloseStatusCode), string(message))
+		return c.WriteControl(websocket.CloseMessage, closeMessage, time.Now().Add(defaultWriteWait))
+	default:
+		return c.WriteMessage(messageType, message)
+	}
+}
+
+func closeWithTimeout(r *SessionRunner, step *TStep, stepVariables map[string]interface{}) (*wsCloseRespObject, error) {
+	if r.wsConn == nil {
+		return nil, errors.New("no connection needs to be closed")
+	}
+	errorChan := make(chan error)
+	go func() {
+		err := writeWebSocket(r, step, stepVariables)
+		if err != nil {
+			errorChan <- errors.Wrap(err, "send close message failed")
+			return
+		}
+		// discard redundant message left in the connection before close
+		var mt int
+		var message []byte
+		var readErr error
+		for readErr == nil {
+			mt, message, readErr = r.wsConn.ReadMessage()
+			if readErr == nil {
+				log.Info().
+					Str("type", MessageType(mt).toString()).
+					Str("msg", string(message)).
+					Msg("discard redundant message")
+				continue
+			}
+			if e, ok := readErr.(*websocket.CloseError); !ok {
+				errorChan <- errors.Wrap(e, "read message failed")
+				return
+			}
+		}
+		// r.wsConn.Close() will be called at the end of current session, so no need to Close here
+		log.Info().Str("msg", readErr.Error()).Msg("connection closed")
+	}()
+	timer := time.NewTimer(time.Duration(step.WebSocket.Timeout) * time.Millisecond)
+	select {
+	case <-timer.C:
+		timer.Stop()
+		return nil, errors.New("close timeout")
+	case err := <-errorChan:
+		return nil, err
+	case closeResult := <-r.closeResponseChan:
+		return closeResult, nil
+	}
+}
+
+func getResponseObject(t *testing.T, parser *Parser, resp interface{}) (*responseObject, error) {
+	// response could be nil for ping and write case
+	if resp == nil {
+		return nil, nil
+	}
+	switch r := resp.(type) {
+	case *http.Response:
+		return newHttpResponseObject(t, parser, r)
+	case *wsReadRespObject:
+		return newWsReadResponseObject(t, parser, r)
+	case *wsCloseRespObject:
+		return newWsCloseResponseObject(t, parser, r)
+	default:
+		return nil, errors.New("unxexpected reponse type")
+	}
+}
+
+func getContentSize(resp interface{}) int64 {
+	switch r := resp.(type) {
+	case *http.Response:
+		return r.ContentLength
+	case *wsReadRespObject:
+		return int64(unsafe.Sizeof(r.Message))
+	case *wsCloseRespObject:
+		return int64(unsafe.Sizeof(r.Text))
+	default:
+		return -1
+	}
+}

hrp/testcase.go

@@ -132,6 +132,10 @@ func (path *TestCasePath) ToTestCase() (*TestCase, error) {
 			testCase.TestSteps = append(testCase.TestSteps, &StepRendezvous{
 				step: step,
 			})
+		} else if step.WebSocket != nil {
+			testCase.TestSteps = append(testCase.TestSteps, &StepWebSocket{
+				step: step,
+			})
 		} else {
 			log.Warn().Interface("step", step).Msg("[convertTestCase] unexpected step")
 		}

hrp/tests/demo_file_load_ws_message.txt

@@ -0,0 +1 @@
+demo file used for testing load_ws_message function

hrp/tests/protocol_test.go

@@ -6,9 +6,9 @@ import (
 	"github.com/httprunner/httprunner/hrp"
 )
 
-func TestProtocol(t *testing.T) {
+func TestHTTPProtocol(t *testing.T) {
 	testcase := &hrp.TestCase{
-		Config: hrp.NewConfig("run request with different protocol types").
+		Config: hrp.NewConfig("run request with HTTP/1.1 and HTTP/2").
 			SetBaseURL("https://postman-echo.com"),
 		TestSteps: []hrp.IStep{
 			hrp.NewStep("HTTP/1.1 get").
@@ -52,3 +52,76 @@ func TestProtocol(t *testing.T) {
 		t.Fatalf("run testcase error: %v", err)
 	}
 }
+
+func TestWebSocketProtocol(t *testing.T) {
+	testcase := &hrp.TestCase{
+		Config: hrp.NewConfig("run request with WebSocket protocol").
+			SetBaseURL("ws://echo.websocket.events").
+			WithVariables(map[string]interface{}{
+				"n":    5,
+				"a":    12.3,
+				"b":    3.45,
+				"file": "./demo_file_load_ws_message.txt",
+			}),
+		TestSteps: []hrp.IStep{
+			hrp.NewStep("open connection").
+				WebSocket().
+				OpenConnection("/").
+				WithHeaders(map[string]string{"User-Agent": "HttpRunnerPlus"}).
+				Validate().
+				AssertEqual("status_code", 101, "check open status code").
+				AssertEqual("headers.Connection", "Upgrade", "check headers"),
+			hrp.NewStep("ping pong test").
+				WebSocket().
+				PingPong("/").
+				WithTimeout(5000),
+			hrp.NewStep("read sponsor info").
+				WebSocket().
+				Read("/").
+				WithTimeout(5000).
+				Validate().
+				AssertContains("body", "Lob.com", "check sponsor message"),
+			hrp.NewStep("write json").
+				WebSocket().
+				Write("/").
+				WithTextMessage(map[string]interface{}{"foo1": "${gen_random_string($n)}", "foo2": "${max($a, $b)}"}),
+			hrp.NewStep("read json").
+				WebSocket().
+				Read("/").
+				Extract().
+				WithJmesPath("body.foo1", "varFoo1").
+				Validate().
+				AssertLengthEqual("body.foo1", 5, "check json foo1").
+				AssertEqual("body.foo2", 12.3, "check json foo2"),
+			hrp.NewStep("write and read text").
+				WebSocket().
+				WriteAndRead("/").
+				WithTextMessage("$varFoo1").
+				Validate().
+				AssertLengthEqual("body", 5, "check length equal"),
+			hrp.NewStep("write and read binary file").
+				WebSocket().
+				WriteAndRead("/").
+				WithBinaryMessage("${load_ws_message($file)}"),
+			hrp.NewStep("write something redundant").
+				WebSocket().
+				Write("/").
+				WithTextMessage("have a nice day!"),
+			hrp.NewStep("write something redundant").
+				WebSocket().
+				Write("/").
+				WithTextMessage("balabala ..."),
+			hrp.NewStep("close connection").
+				WebSocket().
+				CloseConnection("/").
+				WithTimeout(30000).
+				WithCloseStatus(1000).
+				Validate().
+				AssertEqual("status_code", 1000, "check close status code"),
+		},
+	}
+	err := hrp.NewRunner(t).Run(testcase)
+	if err != nil {
+		t.Fatalf("run testcase error: %v", err)
+	}
+}