test: 补齐 Q2 验收与本机 Q3 弱网崩溃压测
This commit is contained in:
Nixevol
@@ -0,0 +1,342 @@
|
||||
package accept
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"encoding/json"
|
||||
"io"
|
||||
"sync"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"git.asio.asia/nixevol/NixMsg/internal/protocol"
|
||||
"git.asio.asia/nixevol/NixMsg/test/harness"
|
||||
"github.com/mochi-mqtt/server/v2/packets"
|
||||
)
|
||||
|
||||
// MQTTSession 是验收/弱网用的端侧 MQTT 会话（WebSocket + hello + 应用帧）。
|
||||
type MQTTSession struct {
|
||||
t *testing.T
|
||||
mc harness.MQTTClient
|
||||
EndpointID string
|
||||
pktID uint16
|
||||
mu sync.Mutex
|
||||
inbox []map[string]any
|
||||
closed bool
|
||||
done chan struct{}
|
||||
}
|
||||
|
||||
// AppResp 是 type=resp 的解析结果。
|
||||
type AppResp struct {
|
||||
OK bool
|
||||
Error map[string]any
|
||||
Data any
|
||||
Raw map[string]any
|
||||
}
|
||||
|
||||
// MQTTLogin 用密码连上 /mqtt、订阅 down、完成 hello。
|
||||
func MQTTLogin(t *testing.T, httpBase, endpointID, password string) *MQTTSession {
|
||||
t.Helper()
|
||||
mc, err := harness.DialMQTTWebSocket(httpBase, 10*time.Second)
|
||||
if err != nil {
|
||||
t.Fatalf("dial mqtt: %v", err)
|
||||
}
|
||||
s := &MQTTSession{t: t, mc: mc, EndpointID: endpointID, pktID: 10, done: make(chan struct{})}
|
||||
s.connectSubscribeHello(password)
|
||||
go s.readLoop()
|
||||
return s
|
||||
}
|
||||
|
||||
// Close 关闭底层连接。
|
||||
func (s *MQTTSession) Close() {
|
||||
s.mu.Lock()
|
||||
if s.closed {
|
||||
s.mu.Unlock()
|
||||
return
|
||||
}
|
||||
s.closed = true
|
||||
s.mu.Unlock()
|
||||
_ = s.mc.Close()
|
||||
select {
|
||||
case <-s.done:
|
||||
case <-time.After(3 * time.Second):
|
||||
}
|
||||
}
|
||||
|
||||
func (s *MQTTSession) nextPkt() uint16 {
|
||||
s.pktID++
|
||||
if s.pktID == 0 {
|
||||
s.pktID = 1
|
||||
}
|
||||
return s.pktID
|
||||
}
|
||||
|
||||
func (s *MQTTSession) connectSubscribeHello(password string) {
|
||||
t := s.t
|
||||
pk := packets.Packet{
|
||||
FixedHeader: packets.FixedHeader{Type: packets.Connect},
|
||||
ProtocolVersion: 5,
|
||||
Connect: packets.ConnectParams{
|
||||
ProtocolName: []byte("MQTT"),
|
||||
Clean: true,
|
||||
ClientIdentifier: s.EndpointID,
|
||||
Keepalive: 30,
|
||||
UsernameFlag: true,
|
||||
Username: []byte(s.EndpointID),
|
||||
PasswordFlag: true,
|
||||
Password: []byte(password),
|
||||
},
|
||||
}
|
||||
var buf bytes.Buffer
|
||||
if err := pk.ConnectEncode(&buf); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := s.mc.Send(buf.Bytes()); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
ack, err := s.mc.Recv()
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(ack) < 4 || ack[0]>>4 != packets.Connack || ack[3] != 0 {
|
||||
t.Fatalf("connack %x", ack)
|
||||
}
|
||||
|
||||
sub := packets.Packet{
|
||||
FixedHeader: packets.FixedHeader{Type: packets.Subscribe, Qos: 1},
|
||||
ProtocolVersion: 5,
|
||||
PacketID: s.nextPkt(),
|
||||
Filters: packets.Subscriptions{
|
||||
{Filter: "nix/c/" + s.EndpointID + "/down", Qos: 1},
|
||||
},
|
||||
}
|
||||
buf.Reset()
|
||||
if err := sub.SubscribeEncode(&buf); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := s.mc.Send(buf.Bytes()); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if _, err := s.mc.Recv(); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
hello, _ := protocol.Marshal(protocol.Hello{V: protocol.Version, Type: protocol.TypeHello, RID: "h0"})
|
||||
s.publishRaw(hello)
|
||||
deadline := time.Now().Add(10 * time.Second)
|
||||
for time.Now().Before(deadline) {
|
||||
raw, err := s.mc.Recv()
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
m := s.handlePacket(raw)
|
||||
if m == nil {
|
||||
continue
|
||||
}
|
||||
if m["type"] == "resp" && m["ok"] == true {
|
||||
return
|
||||
}
|
||||
if m["type"] == "resp" {
|
||||
t.Fatalf("hello failed: %v", m)
|
||||
}
|
||||
s.push(m)
|
||||
}
|
||||
t.Fatal("hello timeout")
|
||||
}
|
||||
|
||||
func (s *MQTTSession) publishRaw(payload []byte) {
|
||||
t := s.t
|
||||
pub := packets.Packet{
|
||||
FixedHeader: packets.FixedHeader{Type: packets.Publish, Qos: 1},
|
||||
ProtocolVersion: 5,
|
||||
TopicName: "nix/c/" + s.EndpointID + "/up",
|
||||
PacketID: s.nextPkt(),
|
||||
Payload: payload,
|
||||
}
|
||||
var buf bytes.Buffer
|
||||
if err := pub.PublishEncode(&buf); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := s.mc.Send(buf.Bytes()); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
|
||||
func (s *MQTTSession) readLoop() {
|
||||
defer close(s.done)
|
||||
for {
|
||||
raw, err := s.mc.Recv()
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
m := s.handlePacket(raw)
|
||||
if m != nil {
|
||||
s.push(m)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (s *MQTTSession) handlePacket(raw []byte) map[string]any {
|
||||
if len(raw) < 2 {
|
||||
return nil
|
||||
}
|
||||
typ := raw[0] >> 4
|
||||
qos := (raw[0] >> 1) & 0x3
|
||||
switch typ {
|
||||
case packets.Puback, packets.Pingresp, packets.Suback:
|
||||
return nil
|
||||
case packets.Publish:
|
||||
payload, err := decodePublishPayload(raw)
|
||||
if err != nil {
|
||||
return nil
|
||||
}
|
||||
if qos == 1 {
|
||||
rem, n, _ := decodeRemainingLength(raw[1:])
|
||||
body := raw[1+n:]
|
||||
pk := packets.Packet{ProtocolVersion: 5, FixedHeader: packets.FixedHeader{Type: packets.Publish, Remaining: rem, Qos: qos}}
|
||||
if decErr := pk.PublishDecode(body); decErr == nil && pk.PacketID != 0 {
|
||||
ack := packets.Packet{
|
||||
FixedHeader: packets.FixedHeader{Type: packets.Puback},
|
||||
ProtocolVersion: 5,
|
||||
PacketID: pk.PacketID,
|
||||
}
|
||||
var buf bytes.Buffer
|
||||
if encErr := ack.PubackEncode(&buf); encErr == nil {
|
||||
_ = s.mc.Send(buf.Bytes())
|
||||
}
|
||||
}
|
||||
}
|
||||
var m map[string]any
|
||||
if json.Unmarshal(payload, &m) != nil {
|
||||
return nil
|
||||
}
|
||||
return m
|
||||
default:
|
||||
return nil
|
||||
}
|
||||
}
|
||||
|
||||
func (s *MQTTSession) push(m map[string]any) {
|
||||
s.mu.Lock()
|
||||
s.inbox = append(s.inbox, m)
|
||||
s.mu.Unlock()
|
||||
}
|
||||
|
||||
// Request 发上行帧并等同 rid 的 resp。
|
||||
func (s *MQTTSession) Request(t *testing.T, frame map[string]any) AppResp {
|
||||
t.Helper()
|
||||
rid, _ := frame["rid"].(string)
|
||||
payload, err := protocol.Marshal(frame)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
s.publishRaw(payload)
|
||||
deadline := time.Now().Add(15 * time.Second)
|
||||
for time.Now().Before(deadline) {
|
||||
m := s.takeMatching(func(x map[string]any) bool {
|
||||
return x["type"] == "resp" && x["rid"] == rid
|
||||
})
|
||||
if m != nil {
|
||||
r := AppResp{OK: m["ok"] == true, Raw: m, Data: m["data"]}
|
||||
if e, ok := m["error"].(map[string]any); ok {
|
||||
r.Error = e
|
||||
}
|
||||
return r
|
||||
}
|
||||
time.Sleep(5 * time.Millisecond)
|
||||
}
|
||||
t.Fatalf("timeout waiting resp rid=%s", rid)
|
||||
return AppResp{}
|
||||
}
|
||||
|
||||
// WaitType 等到指定 type 的下行帧。
|
||||
func (s *MQTTSession) WaitType(t *testing.T, typ string, timeout time.Duration) map[string]any {
|
||||
t.Helper()
|
||||
deadline := time.Now().Add(timeout)
|
||||
for time.Now().Before(deadline) {
|
||||
m := s.takeMatching(func(x map[string]any) bool { return x["type"] == typ })
|
||||
if m != nil {
|
||||
return m
|
||||
}
|
||||
time.Sleep(5 * time.Millisecond)
|
||||
}
|
||||
t.Fatalf("timeout waiting type=%s", typ)
|
||||
return nil
|
||||
}
|
||||
|
||||
// TryType 在超时内尝试取指定 type；超时返回 nil。
|
||||
func (s *MQTTSession) TryType(typ string, timeout time.Duration) map[string]any {
|
||||
deadline := time.Now().Add(timeout)
|
||||
for time.Now().Before(deadline) {
|
||||
m := s.takeMatching(func(x map[string]any) bool { return x["type"] == typ })
|
||||
if m != nil {
|
||||
return m
|
||||
}
|
||||
time.Sleep(10 * time.Millisecond)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *MQTTSession) takeMatching(pred func(map[string]any) bool) map[string]any {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
for i, m := range s.inbox {
|
||||
if pred(m) {
|
||||
s.inbox = append(s.inbox[:i], s.inbox[i+1:]...)
|
||||
return m
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// DrainEvents 排空 group_event / presence，避免干扰断言。
|
||||
func DrainEvents(t *testing.T, s *MQTTSession, d time.Duration) {
|
||||
t.Helper()
|
||||
deadline := time.Now().Add(d)
|
||||
for time.Now().Before(deadline) {
|
||||
_ = s.takeMatching(func(x map[string]any) bool {
|
||||
typ, _ := x["type"].(string)
|
||||
return typ == "group_event" || typ == "presence"
|
||||
})
|
||||
time.Sleep(20 * time.Millisecond)
|
||||
}
|
||||
}
|
||||
|
||||
func decodePublishPayload(raw []byte) ([]byte, error) {
|
||||
if len(raw) < 2 {
|
||||
return nil, io.ErrUnexpectedEOF
|
||||
}
|
||||
rem, n, err := decodeRemainingLength(raw[1:])
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
body := raw[1+n:]
|
||||
if len(body) != rem {
|
||||
return nil, io.ErrUnexpectedEOF
|
||||
}
|
||||
pk := packets.Packet{
|
||||
ProtocolVersion: 5,
|
||||
FixedHeader: packets.FixedHeader{
|
||||
Type: packets.Publish,
|
||||
Remaining: rem,
|
||||
Qos: (raw[0] >> 1) & 0x3,
|
||||
},
|
||||
}
|
||||
if err := pk.PublishDecode(body); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return pk.Payload, nil
|
||||
}
|
||||
|
||||
func decodeRemainingLength(b []byte) (value int, n int, err error) {
|
||||
var mul uint32 = 1
|
||||
var v uint32
|
||||
for i := 0; i < len(b) && i < 4; i++ {
|
||||
v += uint32(b[i]&127) * mul
|
||||
n++
|
||||
if b[i]&128 == 0 {
|
||||
return int(v), n, nil
|
||||
}
|
||||
mul *= 128
|
||||
}
|
||||
return 0, 0, io.ErrUnexpectedEOF
|
||||
}
|
||||
Reference in New Issue
Block a user