+单元测试代码

antlabs · Aug 3, 2023 · 7e89161 · 7e89161
1 parent 44440f3
commit 7e89161
Show file tree

Hide file tree

Showing 5 changed files with 98 additions and 7 deletions.
diff --git a/README.md b/README.md
@@ -125,7 +125,7 @@ func main() {
 #### 配置header
 ```go
 func main() {
-	quickws.Dial("ws://127.0.0.1:12345/test", quickws.WithHTTPHeader(http.Header{
+	quickws.Dial("ws://127.0.0.1:12345/test", quickws.WithClientHTTPHeader(http.Header{
 		"h1": "v1",
 		"h2":"v2", 
 	}))
@@ -134,14 +134,14 @@ func main() {
 #### 配置握手时的超时时间
 ```go
 func main() {
-	quickws.Dial("ws://127.0.0.1:12345/test", quickws.WithDialTimeout(2 * time.Second))
+	quickws.Dial("ws://127.0.0.1:12345/test", quickws.WithClientDialTimeout(2 * time.Second))
 }
 ```
 
 #### 配置自动回复ping消息
 ```go
 func main() {
-	quickws.Dial("ws://127.0.0.1:12345/test", quickws.WithReplyPing())
+	quickws.Dial("ws://127.0.0.1:12345/test", quickws.WithClientReplyPing())
 }
 ```
 ### 服务端配置参数

diff --git a/conn.go b/conn.go
@@ -361,7 +361,7 @@ func (c *Conn) WriteMessage(op Opcode, writeBuf []byte) (err error) {
 		maskValue = rand.Uint32()
 	}
 
-	return frame.WriteFrame(&c.fw, c.c, writeBuf, rsv1, c.client, op, maskValue)
+	return frame.WriteFrame(&c.fw, c.c, writeBuf, true, rsv1, c.client, op, maskValue)
 }
 
 // TODO 这是一个不安全的方法, writeBuf的格式必须是 14个字节的空白长度+需要写的payload组成
@@ -430,6 +430,54 @@ func (c *Conn) WriteControl(op Opcode, data []byte) (err error) {
 	return c.WriteMessage(op, data)
 }
 
+// 写分段数据, 目前主要是单元测试使用
+func (c *Conn) writeFragment(op Opcode, writeBuf []byte, maxFragment int /*单个段最大size*/) (err error) {
+	if len(writeBuf) < maxFragment {
+		c.WriteMessage(op, writeBuf)
+		return
+	}
+
+	if op == opcode.Text {
+		if !c.utf8Check(writeBuf) {
+			return ErrTextNotUTF8
+		}
+	}
+
+	rsv1 := c.compression && (op == opcode.Text || op == opcode.Binary)
+	if rsv1 {
+		var out wrapBuffer
+		w := compressNoContextTakeover(&out, defaultCompressionLevel)
+		if _, err = io.Copy(w, bytes.NewReader(writeBuf)); err != nil {
+			return
+		}
+
+		if err = w.Close(); err != nil {
+			return
+		}
+		writeBuf = out.Bytes()
+	}
+
+	// f.Opcode = op
+	// f.PayloadLen = int64(len(writeBuf))
+	maskValue := uint32(0)
+	if c.client {
+		maskValue = rand.Uint32()
+	}
+
+	for len(writeBuf) > 0 {
+		if len(writeBuf) > maxFragment {
+			if err := frame.WriteFrame(&c.fw, c.c, writeBuf[:maxFragment], false, rsv1, c.client, op, maskValue); err != nil {
+				return err
+			}
+			writeBuf = writeBuf[maxFragment:]
+			op = Continuation
+			continue
+		}
+		return frame.WriteFrame(&c.fw, c.c, writeBuf, true, rsv1, c.client, op, maskValue)
+	}
+	return nil
+}
+
 func (c *Conn) Close() (err error) {
 	c.once.Do(func() {
 		c.bp.Free()

diff --git a/conn_test.go b/conn_test.go
@@ -183,3 +183,46 @@ func Test_ReadMessage(t *testing.T) {
 		}
 	})
 }
+
+// 测试分段frame
+func TestFragmentFrame(t *testing.T) {
+	t.Run("FragmentFrame10", func(t *testing.T) {
+		run := int32(0)
+		data := make(chan string, 1)
+		upgrade := NewUpgrade(WithServerBufioParseMode(), WithServerOnMessageFunc(func(c *Conn, op Opcode, payload []byte) {
+			c.WriteMessage(op, payload)
+		}))
+		ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			c, err := upgrade.Upgrade(w, r)
+			if err != nil {
+				t.Error(err)
+			}
+			c.StartReadLoop()
+		}))
+
+		defer ts.Close()
+
+		url := strings.ReplaceAll(ts.URL, "http", "ws")
+		con, err := Dial(url, WithClientDisableBufioClearHack(), WithClientOnMessageFunc(func(c *Conn, mt Opcode, payload []byte) {
+			atomic.AddInt32(&run, int32(1))
+			data <- string(payload)
+		}))
+		if err != nil {
+			t.Error(err)
+		}
+		defer con.Close()
+
+		con.writeFragment(Binary, []byte("hello"), 1)
+		con.StartReadLoop()
+		select {
+		case d := <-data:
+			if d != "hello" {
+				t.Errorf("write message or read message fail:got:%s, need:hello\n", d)
+			}
+		case <-time.After(1000 * time.Millisecond):
+		}
+		if atomic.LoadInt32(&run) != 1 {
+			t.Error("not run server:method fail")
+		}
+	})
+}
diff --git a/go.mod b/go.mod
@@ -2,4 +2,4 @@ module github.com/antlabs/quickws
 
 go 1.20
 
-require github.com/antlabs/wsutil v0.0.11
+require github.com/antlabs/wsutil v0.0.12
diff --git a/go.sum b/go.sum
@@ -1,2 +1,2 @@
-github.com/antlabs/wsutil v0.0.11 h1:T4pTMroxn3VXrDZyj5LvvidXbabaARzQYHQj3rySD3o=
-github.com/antlabs/wsutil v0.0.11/go.mod h1:7ec5eUM7nmKW+Oi6F1I58iatOeL9k+yIsfOh1zh910g=
+github.com/antlabs/wsutil v0.0.12 h1:YHn8uuwhxYyy0tgpCWgDHNnKouaEhtQUkrX0QNgoDqc=
+github.com/antlabs/wsutil v0.0.12/go.mod h1:7ec5eUM7nmKW+Oi6F1I58iatOeL9k+yIsfOh1zh910g=