package influxdb_test import ( "bytes" "context" "fmt" "net" "net/url" "sync" "testing" "time" "github.com/influxdata/telegraf" "github.com/influxdata/telegraf/metric" "github.com/influxdata/telegraf/plugins/outputs/influxdb" "github.com/influxdata/telegraf/plugins/serializers/influx" "github.com/stretchr/testify/require" ) var ( metricString = "cpu value=42 0\n" ) func getMetric() telegraf.Metric { metric, err := metric.New( "cpu", map[string]string{}, map[string]interface{}{ "value": 42.0, }, time.Unix(0, 0), ) if err != nil { panic(err) } return metric } func getURL() *url.URL { u, err := url.Parse("udp://localhost:0") if err != nil { panic(err) } return u } type MockConn struct { WriteF func(b []byte) (n int, err error) CloseF func() error } func (c *MockConn) Write(b []byte) (n int, err error) { return c.WriteF(b) } func (c *MockConn) Close() error { return c.CloseF() } type MockDialer struct { DialContextF func(network, address string) (influxdb.Conn, error) } func (d *MockDialer) DialContext(ctx context.Context, network string, address string) (influxdb.Conn, error) { return d.DialContextF(network, address) } type MockSerializer struct { SerializeF func(metric telegraf.Metric) ([]byte, error) SerializeBatchF func(metrics []telegraf.Metric) ([]byte, error) } func (s *MockSerializer) Serialize(metric telegraf.Metric) ([]byte, error) { return s.SerializeF(metric) } func (s *MockSerializer) SerializeBatch(metrics []telegraf.Metric) ([]byte, error) { return s.SerializeBatchF(metrics) } func TestUDP_NewUDPClientNoURL(t *testing.T) { config := &influxdb.UDPConfig{} _, err := influxdb.NewUDPClient(config) require.Equal(t, err, influxdb.ErrMissingURL) } func TestUDP_URL(t *testing.T) { u := getURL() config := &influxdb.UDPConfig{ URL: u, } client, err := influxdb.NewUDPClient(config) require.NoError(t, err) require.Equal(t, u.String(), client.URL()) } func TestUDP_Simple(t *testing.T) { var buffer bytes.Buffer config := &influxdb.UDPConfig{ URL: getURL(), Dialer: &MockDialer{ DialContextF: func(network, address string) (influxdb.Conn, error) { conn := &MockConn{ WriteF: func(b []byte) (n int, err error) { buffer.Write(b) return 0, nil }, } return conn, nil }, }, } client, err := influxdb.NewUDPClient(config) require.NoError(t, err) ctx := context.Background() err = client.Write(ctx, []telegraf.Metric{ getMetric(), getMetric(), }) require.NoError(t, err) require.Equal(t, metricString+metricString, buffer.String()) } func TestUDP_DialError(t *testing.T) { u, err := url.Parse("invalid://127.0.0.1:9999") require.NoError(t, err) config := &influxdb.UDPConfig{ URL: u, Dialer: &MockDialer{ DialContextF: func(network, address string) (influxdb.Conn, error) { return nil, fmt.Errorf( `unsupported scheme [invalid://localhost:9999]: "invalid"`) }, }, } client, err := influxdb.NewUDPClient(config) require.NoError(t, err) ctx := context.Background() err = client.Write(ctx, []telegraf.Metric{getMetric()}) require.Error(t, err) } func TestUDP_WriteError(t *testing.T) { closed := false config := &influxdb.UDPConfig{ URL: getURL(), Dialer: &MockDialer{ DialContextF: func(network, address string) (influxdb.Conn, error) { conn := &MockConn{ WriteF: func(b []byte) (n int, err error) { return 0, fmt.Errorf( "write udp 127.0.0.1:52190->127.0.0.1:9999: write: connection refused") }, CloseF: func() error { closed = true return nil }, } return conn, nil }, }, } client, err := influxdb.NewUDPClient(config) require.NoError(t, err) ctx := context.Background() err = client.Write(ctx, []telegraf.Metric{getMetric()}) require.Error(t, err) require.True(t, closed) } func TestUDP_SerializeError(t *testing.T) { config := &influxdb.UDPConfig{ URL: getURL(), Dialer: &MockDialer{ DialContextF: func(network, address string) (influxdb.Conn, error) { conn := &MockConn{} return conn, nil }, }, Serializer: &MockSerializer{ SerializeF: func(metric telegraf.Metric) ([]byte, error) { return nil, influx.ErrNeedMoreSpace }, }, } client, err := influxdb.NewUDPClient(config) require.NoError(t, err) ctx := context.Background() err = client.Write(ctx, []telegraf.Metric{getMetric()}) require.Error(t, err) require.Contains(t, err.Error(), influx.ErrNeedMoreSpace.Error()) } func TestUDP_WriteWithRealConn(t *testing.T) { conn, err := net.ListenPacket("udp", "127.0.0.1:0") require.NoError(t, err) metrics := []telegraf.Metric{ getMetric(), getMetric(), } buf := make([]byte, 200) var wg sync.WaitGroup wg.Add(1) go func() { defer wg.Done() var total int for _, _ = range metrics { n, _, err := conn.ReadFrom(buf[total:]) if err != nil { break } total += n } buf = buf[:total] }() addr := conn.LocalAddr() u, err := url.Parse(fmt.Sprintf("%s://%s", addr.Network(), addr)) require.NoError(t, err) config := &influxdb.UDPConfig{ URL: u, } client, err := influxdb.NewUDPClient(config) require.NoError(t, err) ctx := context.Background() err = client.Write(ctx, metrics) require.NoError(t, err) wg.Wait() require.Equal(t, metricString+metricString, string(buf)) }