68 lines
1.7 KiB
Go
68 lines
1.7 KiB
Go
package main
|
|
|
|
import (
|
|
"context"
|
|
"net"
|
|
"testing"
|
|
"time"
|
|
)
|
|
|
|
func TestUDPListenerRoundTrip(t *testing.T) {
|
|
ch := make(chan LogRecord, 4)
|
|
ps := NewPromStore()
|
|
|
|
// Bind to an ephemeral port on loopback.
|
|
pc, err := net.ListenPacket("udp", "127.0.0.1:0")
|
|
if err != nil {
|
|
t.Fatalf("listen probe: %v", err)
|
|
}
|
|
addr := pc.LocalAddr().String()
|
|
pc.Close() // release; listener will re-bind
|
|
|
|
ctx, cancel := context.WithCancel(context.Background())
|
|
defer cancel()
|
|
|
|
u := NewUDPListener(addr, 24, 48, ch)
|
|
u.SetProm(ps)
|
|
go u.Run(ctx)
|
|
|
|
// Dial the listener and send one valid and one malformed packet.
|
|
conn, err := net.Dial("udp", addr)
|
|
if err != nil {
|
|
t.Fatalf("dial: %v", err)
|
|
}
|
|
defer conn.Close()
|
|
|
|
// The listener is started asynchronously; retry for up to 1s.
|
|
good := "www.example.com\t1.2.3.4\tGET\t/\t200\t42\t0.010\t0\t12345\tdirect\t10.0.0.1\thttps"
|
|
bad := "not enough\tfields"
|
|
deadline := time.Now().Add(time.Second)
|
|
for time.Now().Before(deadline) {
|
|
conn.Write([]byte(good))
|
|
conn.Write([]byte(bad))
|
|
select {
|
|
case rec := <-ch:
|
|
if rec.Website != "www.example.com" || rec.SourceTag != "direct" {
|
|
t.Fatalf("bad record: %+v", rec)
|
|
}
|
|
// Give the listener a moment to process the malformed packet too.
|
|
time.Sleep(50 * time.Millisecond)
|
|
ps.udpMu.Lock()
|
|
pkt, suc, con := ps.udpPacketsReceived, ps.udpLoglinesSuccess, ps.udpLoglinesConsumed
|
|
ps.udpMu.Unlock()
|
|
if pkt < 2 {
|
|
t.Errorf("udpPacketsReceived=%d, want >=2", pkt)
|
|
}
|
|
if suc < 1 {
|
|
t.Errorf("udpLoglinesSuccess=%d, want >=1", suc)
|
|
}
|
|
if con < 1 {
|
|
t.Errorf("udpLoglinesConsumed=%d, want >=1", con)
|
|
}
|
|
return
|
|
case <-time.After(50 * time.Millisecond):
|
|
}
|
|
}
|
|
t.Fatal("no record received within 1s")
|
|
}
|