udpconn_test.go

package statsd

import (
	"errors"
	"net"
	"runtime"
	"testing"
	"time"
	"unsafe"
)

func checkNumGoroutines() func(t *testing.T) {
	n := runtime.NumGoroutine()
	return func(t *testing.T) {
		t.Helper()
		for i := 0; i < 100; i++ {
			if n := runtime.NumGoroutine() - n; n <= 0 {
				return
			}
			time.Sleep(time.Millisecond * 10)
		}
		if n := runtime.NumGoroutine() - n; n > 0 {
			t.Error(n)
		}
	}
}

func TestNewUDPConn_invalidRate(t *testing.T) {
	defer checkNumGoroutines()(t)

	if v, err := NewUDPConn(`udp`, testAddr, 0); err == nil || v != nil {
		t.Error(v, err)
	}
}

func TestNewUDPConn_resolveUDPAddressErr(t *testing.T) {
	defer checkNumGoroutines()(t)

	e := errors.New(`some error`)
	defer func() func() {
		old := resolveUDPAddress
		resolveUDPAddress = func(network, address string) (*net.UDPAddr, error) {
			if network != `udp` || address != testAddr {
				t.Error(network, address)
			}
			return nil, e
		}
		return func() { resolveUDPAddress = old }
	}()()
	if v, err := NewUDPConn(`udp`, testAddr, time.Minute); err != e || v != nil {
		t.Error(v, err)
	}
}

func TestNewUDPConn_listenUDPErr(t *testing.T) {
	defer checkNumGoroutines()(t)

	e := errors.New(`some error`)
	defer func() func() {
		old := resolveUDPAddress
		resolveUDPAddress = func(network, address string) (*net.UDPAddr, error) { return nil, nil }
		return func() { resolveUDPAddress = old }
	}()()
	defer func() func() {
		old := listenUDP
		listenUDP = func(network string, laddr *net.UDPAddr) (*net.UDPConn, error) {
			if network != `udp` || laddr != nil {
				t.Error(network, laddr)
			}
			return nil, e
		}
		return func() { listenUDP = old }
	}()()
	if v, err := NewUDPConn(`udp`, testAddr, time.Minute); err != e || v != nil {
		t.Error(v, err)
	}
}

func TestUDPConn_stats(t *testing.T) {
	defer checkNumGoroutines()(t)

	var conn *UDPConn
	defer func() {
		_ = conn.Close()
		_ = conn.Close()
	}()
	testNetwork(t, "udp", func(srv *server) []Option {
		var err error
		conn, err = NewUDPConn(`udp`, srv.addr, time.Minute)
		if err != nil {
			t.Fatal(err)
		}
		return []Option{
			WriteCloser(conn),
			TrimTrailingNewline(true),
			UDPCheck(true),
		}
	})
}

func TestUDPConn_worker(t *testing.T) {
	defer checkNumGoroutines()(t)

	type (
		ResolveUDPAddressIn struct {
			Network string
			Address string
		}
		ResolveUDPAddressOut struct {
			Addr *net.UDPAddr
			Err  error
		}
	)
	var (
		resolveUDPAddressIn  = make(chan ResolveUDPAddressIn)
		resolveUDPAddressOut = make(chan ResolveUDPAddressOut)
	)
	defer func() func() {
		old := resolveUDPAddress
		resolveUDPAddress = func(network, address string) (*net.UDPAddr, error) {
			resolveUDPAddressIn <- ResolveUDPAddressIn{Network: network, Address: address}
			v := <-resolveUDPAddressOut
			return v.Addr, v.Err
		}
		return func() { resolveUDPAddress = old }
	}()()

	const (
		network = `udp`
		address = `some-address:1245`
	)
	conn := UDPConn{
		addr: unsafe.Pointer(&net.UDPAddr{
			IP:   net.ParseIP(`127.0.0.1`),
			Port: 184823,
		}),
		stop: make(chan struct{}, 1),
		done: make(chan struct{}),
	}
	workerIsDone := make(chan struct{})
	go func() {
		conn.worker(network, address, 1)
		close(workerIsDone)
	}()

	// set the value
	addr := &net.UDPAddr{
		IP:   net.ParseIP(`127.0.0.2`),
		Port: 1245,
	}
	if v := <-resolveUDPAddressIn; v.Network != network || v.Address != address {
		t.Fatal(v)
	}
	resolveUDPAddressOut <- ResolveUDPAddressOut{Addr: addr}

	// set the value to the same, but different pointer (won't update)
	if v := <-resolveUDPAddressIn; v.Network != network || v.Address != address {
		t.Fatal(v)
	}
	resolveUDPAddressOut <- ResolveUDPAddressOut{Addr: &net.UDPAddr{
		IP:   net.ParseIP(`127.0.0.2`),
		Port: 1245,
	}}

	// doesn't change anything (also, ensure the value is stored)
	for i := 0; i < 2; i++ {
		if v := <-resolveUDPAddressIn; v.Network != network || v.Address != address {
			t.Fatal(v)
		}
		resolveUDPAddressOut <- ResolveUDPAddressOut{Err: errors.New(`some error`)}
	}

	if v := (*net.UDPAddr)(conn.addr); v != addr {
		t.Error(v)
	}

	close(conn.stop)
	go func() {
		for {
			select {
			case <-conn.done:
				return
			case <-resolveUDPAddressIn:
				resolveUDPAddressOut <- ResolveUDPAddressOut{Err: errors.New(`closing`)}
			}
		}
	}()
	<-conn.done
	<-workerIsDone
}