diff --git a/plugins/processors/s2geo/s2geo_test.go b/plugins/processors/s2geo/s2geo_test.go index 78b0d4d2f..bab79908c 100644 --- a/plugins/processors/s2geo/s2geo_test.go +++ b/plugins/processors/s2geo/s2geo_test.go @@ -1,10 +1,12 @@ package s2geo import ( + "sync" "testing" "time" "github.com/influxdata/telegraf" + "github.com/influxdata/telegraf/metric" "github.com/influxdata/telegraf/testutil" "github.com/stretchr/testify/require" ) @@ -24,7 +26,7 @@ func TestGeo(t *testing.T) { err := plugin.Init() require.NoError(t, err) - metric := testutil.MustMetric( + m := testutil.MustMetric( "mta", map[string]string{}, map[string]interface{}{ @@ -48,8 +50,74 @@ func TestGeo(t *testing.T) { ), } - actual := plugin.Apply(metric) + actual := plugin.Apply(m) testutil.RequireMetricsEqual(t, expected, actual) - actual = pluginMostlyDefault.Apply(metric) + actual = pluginMostlyDefault.Apply(m) testutil.RequireMetricsEqual(t, expected, actual) } + +func TestTracking(t *testing.T) { + inputRaw := []telegraf.Metric{ + metric.New("foo", map[string]string{}, map[string]interface{}{"lat": 40.878738, "lon": 72.517572}, time.Unix(0, 0)), + metric.New("bar", map[string]string{}, map[string]interface{}{"lat": 42.842451, "lon": 74.211361}, time.Unix(0, 0)), + metric.New("baz", map[string]string{}, map[string]interface{}{"lat": 32.300963, "lon": 14.123442}, time.Unix(0, 0)), + } + + var mu sync.Mutex + delivered := make([]telegraf.DeliveryInfo, 0, len(inputRaw)) + notify := func(di telegraf.DeliveryInfo) { + mu.Lock() + defer mu.Unlock() + delivered = append(delivered, di) + } + + input := make([]telegraf.Metric, 0, len(inputRaw)) + for _, m := range inputRaw { + tm, _ := metric.WithTracking(m, notify) + input = append(input, tm) + } + + expected := []telegraf.Metric{ + metric.New( + "foo", + map[string]string{"s2_cell_id": "3"}, + map[string]interface{}{"lat": 40.878738, "lon": 72.517572}, + time.Unix(0, 0), + ), + metric.New( + "bar", + map[string]string{"s2_cell_id": "3"}, + map[string]interface{}{"lat": 42.842451, "lon": 74.211361}, + time.Unix(0, 0), + ), + metric.New( + "baz", + map[string]string{"s2_cell_id": "1"}, + map[string]interface{}{"lat": 32.300963, "lon": 14.123442}, + time.Unix(0, 0), + ), + } + + plugin := &Geo{ + LatField: "lat", + LonField: "lon", + TagKey: "s2_cell_id", + } + require.NoError(t, plugin.Init()) + + // Process expected metrics and compare with resulting metrics + actual := plugin.Apply(input...) + testutil.RequireMetricsEqual(t, expected, actual) + + // Simulate output acknowledging delivery + for _, m := range actual { + m.Accept() + } + + // Check delivery + require.Eventuallyf(t, func() bool { + mu.Lock() + defer mu.Unlock() + return len(expected) == len(delivered) + }, time.Second, 100*time.Millisecond, "%d delivered but %d expected", len(delivered), len(expected)) +}