Function TestForwarder

modules/distributor/forwarder_test.go:25–73

(t *testing.T)

const tenantID = "tenant-id"

func TestForwarder(t *testing.T) {
	oCfg := overrides.Config{}
	oCfg.RegisterFlagsAndApplyDefaults(&flag.FlagSet{})

	id, err := util.HexStringToTraceID("1234567890abcdef")
	require.NoError(t, err)

	b := test.MakeBatch(10, id)
	keys, rebatchedTraces, _, _, err := requestsByTraceID([]*v1.ResourceSpans{b}, tenantID, 10, 1000)
	require.NoError(t, err)

	o, err := overrides.NewOverrides(oCfg, nil, prometheus.DefaultRegisterer)
	require.NoError(t, err)

	noGenerateMetricsRequestCount := 0

	wg := sync.WaitGroup{}
	f := newGeneratorForwarder(
		log.NewNopLogger(),
		func(_ context.Context, userID string, k []uint32, traces []*rebatchedTrace, noGenerateMetrics bool) error {
			assert.Equal(t, tenantID, userID)
			assert.Equal(t, keys, k)
			assert.Equal(t, rebatchedTraces, traces)
			if noGenerateMetrics {
				noGenerateMetricsRequestCount++
			}
			wg.Done()
			return nil
		},
		o,
	)
	require.NoError(t, f.start(context.Background()))
	defer func() {
		require.NoError(t, f.stop(nil))
	}()

	wg.Add(1)
	// Mark this request as "to-be-ignored" for metrics generation.
	ctx = metadata.NewIncomingContext(ctx, metadata.Pairs(generator.NoGenerateMetricsContextKey, ""))
	f.SendTraces(ctx, tenantID, keys, rebatchedTraces)
	wg.Wait()

	wg.Add(1)
	f.SendTraces(context.Background(), tenantID, keys, rebatchedTraces)
	wg.Wait()

	// Expect to receive one request for which no metrics should be generated in metrics-generator.
	require.Equal(t, 1, noGenerateMetricsRequestCount)
}

func TestForwarder_shutdown(t *testing.T) {
	oCfg := overrides.Config{}