nginx-logtail/cmd/collector/udp_test.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 := "v1\twww.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")
}