76 lines
1.6 KiB
Go
76 lines
1.6 KiB
Go
package load
|
|
|
|
import (
|
|
"bytes"
|
|
"io"
|
|
"net"
|
|
"testing"
|
|
"time"
|
|
)
|
|
|
|
// 极简 MQTT broker：读 CONNECT，回 CONNACK accepted，保持连接。
|
|
func startFakeMQTTBroker(t *testing.T) (addr string, closeFn func()) {
|
|
t.Helper()
|
|
ln, err := net.Listen("tcp", "127.0.0.1:0")
|
|
if err != nil {
|
|
t.Fatalf("listen: %v", err)
|
|
}
|
|
done := make(chan struct{})
|
|
go func() {
|
|
for {
|
|
conn, err := ln.Accept()
|
|
if err != nil {
|
|
select {
|
|
case <-done:
|
|
return
|
|
default:
|
|
return
|
|
}
|
|
}
|
|
go func(c net.Conn) {
|
|
defer func() { _ = c.Close() }()
|
|
_ = c.SetDeadline(time.Now().Add(5 * time.Second))
|
|
buf := make([]byte, 256)
|
|
n, err := c.Read(buf)
|
|
if err != nil || n < 2 || buf[0] != 0x10 {
|
|
return
|
|
}
|
|
// CONNACK: type 0x20, remaining 2, flags 0, rc 0
|
|
_, _ = c.Write([]byte{0x20, 0x02, 0x00, 0x00})
|
|
_ = c.SetDeadline(time.Time{})
|
|
_, _ = io.Copy(io.Discard, c)
|
|
}(conn)
|
|
}
|
|
}()
|
|
return ln.Addr().String(), func() {
|
|
close(done)
|
|
_ = ln.Close()
|
|
}
|
|
}
|
|
|
|
func TestMQTTBenchConnectN(t *testing.T) {
|
|
addr, stop := startFakeMQTTBroker(t)
|
|
defer stop()
|
|
|
|
var out bytes.Buffer
|
|
b := &MQTTBench{Addr: addr, ClientIDPrefix: "q-t-", Timeout: 2 * time.Second}
|
|
defer b.Close()
|
|
if err := b.ConnectN(3, &out); err != nil {
|
|
t.Fatalf("ConnectN: %v", err)
|
|
}
|
|
if b.Alive() != 3 {
|
|
t.Fatalf("alive=%d", b.Alive())
|
|
}
|
|
got := out.String()
|
|
if !bytes.Contains(out.Bytes(), []byte("mqtt connections: 3")) {
|
|
t.Fatalf("output missing count: %q", got)
|
|
}
|
|
}
|
|
|
|
func TestMQTTBenchRequiresAddr(t *testing.T) {
|
|
b := &MQTTBench{}
|
|
if err := b.ConnectN(1, nil); err == nil {
|
|
t.Fatal("expected error")
|
|
}
|
|
}
|