github.com/yankunsam/loki/v2@v2.6.3-0.20220817130409-389df5235c27/clients/pkg/promtail/targets/kafka/consumer_test.go (about)

     1  package kafka
     2  
     3  import (
     4  	"context"
     5  	"errors"
     6  	"testing"
     7  	"time"
     8  
     9  	"github.com/Shopify/sarama"
    10  	"github.com/go-kit/log"
    11  	"github.com/prometheus/common/model"
    12  	"github.com/stretchr/testify/require"
    13  
    14  	"github.com/grafana/loki/clients/pkg/promtail/targets/target"
    15  )
    16  
    17  type DiscovererFn func(sarama.ConsumerGroupSession, sarama.ConsumerGroupClaim) (RunnableTarget, error)
    18  
    19  func (d DiscovererFn) NewTarget(session sarama.ConsumerGroupSession, claim sarama.ConsumerGroupClaim) (RunnableTarget, error) {
    20  	return d(session, claim)
    21  }
    22  
    23  type fakeTarget struct {
    24  	ctx context.Context
    25  	lbs model.LabelSet
    26  }
    27  
    28  func (f *fakeTarget) run()                             { <-f.ctx.Done() }
    29  func (f *fakeTarget) Type() target.TargetType          { return "" }
    30  func (f *fakeTarget) DiscoveredLabels() model.LabelSet { return nil }
    31  func (f *fakeTarget) Labels() model.LabelSet           { return f.lbs }
    32  func (f *fakeTarget) Ready() bool                      { return true }
    33  func (f *fakeTarget) Details() interface{}             { return nil }
    34  
    35  func Test_ComsumerConsume(t *testing.T) {
    36  	var (
    37  		group       = &testConsumerGroupHandler{}
    38  		session     = &testSession{}
    39  		ctx, cancel = context.WithCancel(context.Background())
    40  		c           = &consumer{
    41  			logger:        log.NewNopLogger(),
    42  			ctx:           context.Background(),
    43  			cancel:        func() {},
    44  			ConsumerGroup: group,
    45  			discoverer: DiscovererFn(func(session sarama.ConsumerGroupSession, claim sarama.ConsumerGroupClaim) (RunnableTarget, error) {
    46  				if claim.Topic() != "dropped" {
    47  					return &fakeTarget{
    48  						ctx: ctx,
    49  						lbs: model.LabelSet{"topic": model.LabelValue(claim.Topic())},
    50  					}, nil
    51  				}
    52  				return &fakeTarget{
    53  					ctx: ctx,
    54  				}, nil
    55  			}),
    56  		}
    57  	)
    58  
    59  	c.start(ctx, []string{"foo"})
    60  	require.Eventually(t, group.consuming.Load, 5*time.Second, 100*time.Microsecond)
    61  	require.NoError(t, group.handler.Setup(session))
    62  	go func() {
    63  		err := group.handler.ConsumeClaim(session, newTestClaim("foo", 1, 2))
    64  		require.NoError(t, err)
    65  	}()
    66  	go func() {
    67  		err := group.handler.ConsumeClaim(session, newTestClaim("bar", 1, 2))
    68  		require.NoError(t, err)
    69  	}()
    70  	go func() {
    71  		err := group.handler.ConsumeClaim(session, newTestClaim("dropped", 1, 2))
    72  		require.NoError(t, err)
    73  	}()
    74  	require.Eventually(t, func() bool {
    75  		return len(c.getActiveTargets()) == 2
    76  	}, 2*time.Second, 100*time.Millisecond)
    77  	require.Eventually(t, func() bool {
    78  		return len(c.getDroppedTargets()) == 1
    79  	}, 2*time.Second, 100*time.Millisecond)
    80  	err := group.handler.Cleanup(session)
    81  	require.NoError(t, err)
    82  	cancel()
    83  	c.stop()
    84  }
    85  
    86  func Test_ComsumerRetry(t *testing.T) {
    87  	var (
    88  		group = &testConsumerGroupHandler{
    89  			returnErr: errors.New("foo"),
    90  		}
    91  		ctx, cancel = context.WithCancel(context.Background())
    92  		c           = &consumer{
    93  			logger:        log.NewNopLogger(),
    94  			ctx:           context.Background(),
    95  			cancel:        func() {},
    96  			ConsumerGroup: group,
    97  		}
    98  	)
    99  	defer cancel()
   100  	c.start(ctx, []string{"foo"})
   101  	<-time.After(2 * time.Second)
   102  	c.stop()
   103  }