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

     1  package stdin
     2  
     3  import (
     4  	"bytes"
     5  	"io"
     6  	"os"
     7  	"strings"
     8  	"testing"
     9  
    10  	"github.com/prometheus/client_golang/prometheus"
    11  	"github.com/prometheus/common/model"
    12  	"github.com/stretchr/testify/require"
    13  	"gopkg.in/yaml.v2"
    14  
    15  	"github.com/grafana/loki/clients/pkg/logentry/stages"
    16  	"github.com/grafana/loki/clients/pkg/promtail/api"
    17  	"github.com/grafana/loki/clients/pkg/promtail/client/fake"
    18  	"github.com/grafana/loki/clients/pkg/promtail/scrapeconfig"
    19  
    20  	"github.com/grafana/loki/pkg/logproto"
    21  	util_log "github.com/grafana/loki/pkg/util/log"
    22  )
    23  
    24  func Test_newReaderTarget(t *testing.T) {
    25  	tests := []struct {
    26  		name    string
    27  		in      io.Reader
    28  		cfg     scrapeconfig.Config
    29  		want    []api.Entry
    30  		wantErr bool
    31  	}{
    32  		{
    33  			"no newlines",
    34  			bytes.NewReader([]byte("bar")),
    35  			scrapeconfig.Config{},
    36  			[]api.Entry{
    37  				{Labels: model.LabelSet{}, Entry: logproto.Entry{Line: "bar"}},
    38  			},
    39  			false,
    40  		},
    41  		{
    42  			"empty",
    43  			bytes.NewReader([]byte("")),
    44  			scrapeconfig.Config{},
    45  			nil,
    46  			false,
    47  		},
    48  		{
    49  			"newlines",
    50  			bytes.NewReader([]byte("\nfoo\r\nbar")),
    51  			scrapeconfig.Config{},
    52  			[]api.Entry{
    53  				{Labels: model.LabelSet{}, Entry: logproto.Entry{Line: "foo"}},
    54  				{Labels: model.LabelSet{}, Entry: logproto.Entry{Line: "bar"}},
    55  			},
    56  			false,
    57  		},
    58  		{
    59  			"pipeline",
    60  			bytes.NewReader([]byte("\nfoo\r\nbar")),
    61  			scrapeconfig.Config{
    62  				PipelineStages: loadConfig(stagesConfig),
    63  			},
    64  			[]api.Entry{
    65  				{Labels: model.LabelSet{"new_key": "hello world!"}, Entry: logproto.Entry{Line: "foo"}},
    66  				{Labels: model.LabelSet{"new_key": "hello world!"}, Entry: logproto.Entry{Line: "bar"}},
    67  			},
    68  			false,
    69  		},
    70  		{
    71  			"default config",
    72  			bytes.NewReader([]byte("\nfoo\r\nbar")),
    73  			defaultStdInCfg,
    74  			[]api.Entry{
    75  				{Labels: model.LabelSet{"job": "stdin", "hostname": model.LabelValue(hostName)}, Entry: logproto.Entry{Line: "foo"}},
    76  				{Labels: model.LabelSet{"job": "stdin", "hostname": model.LabelValue(hostName)}, Entry: logproto.Entry{Line: "bar"}},
    77  			},
    78  			false,
    79  		},
    80  	}
    81  	for _, tt := range tests {
    82  		t.Run(tt.name, func(t *testing.T) {
    83  			c := fake.New(func() {})
    84  			got, err := newReaderTarget(prometheus.DefaultRegisterer, util_log.Logger, tt.in, c, tt.cfg)
    85  			if (err != nil) != tt.wantErr {
    86  				t.Errorf("newReaderTarget() error = %v, wantErr %v", err, tt.wantErr)
    87  				return
    88  			}
    89  			if err != nil {
    90  				return
    91  			}
    92  			<-got.ctx.Done()
    93  			c.Stop()
    94  			compareEntries(t, tt.want, c.Received())
    95  		})
    96  	}
    97  }
    98  
    99  type mockShutdownable struct {
   100  	called chan bool
   101  }
   102  
   103  func (m *mockShutdownable) Shutdown() {
   104  	m.called <- true
   105  }
   106  
   107  type fakeStdin struct {
   108  	io.Reader
   109  	os.FileInfo
   110  }
   111  
   112  func newFakeStdin(data string) *fakeStdin {
   113  	return &fakeStdin{
   114  		Reader: strings.NewReader(data),
   115  	}
   116  }
   117  
   118  func (f fakeStdin) Stat() (os.FileInfo, error) { return f.FileInfo, nil }
   119  
   120  func Test_Shutdown(t *testing.T) {
   121  	stdIn = newFakeStdin("line")
   122  	appMock := &mockShutdownable{called: make(chan bool, 1)}
   123  	recorder := fake.New(func() {})
   124  	manager, err := NewStdinTargetManager(prometheus.DefaultRegisterer, util_log.Logger, appMock, recorder, []scrapeconfig.Config{{}})
   125  	require.NoError(t, err)
   126  	require.NotNil(t, manager)
   127  	require.Equal(t, true, <-appMock.called)
   128  	recorder.Stop()
   129  	compareEntries(t, []api.Entry{{Labels: model.LabelSet{}, Entry: logproto.Entry{Line: "line"}}}, recorder.Received())
   130  }
   131  
   132  func compareEntries(t *testing.T, expected, actual []api.Entry) {
   133  	t.Helper()
   134  	require.Equal(t, len(expected), len(actual))
   135  	for i := range expected {
   136  		require.Equal(t, expected[i].Entry.Line, actual[i].Entry.Line)
   137  		require.Equal(t, expected[i].Labels, actual[i].Labels)
   138  	}
   139  }
   140  
   141  func Test_StdinConfigs(t *testing.T) {
   142  
   143  	// should take the first config
   144  	require.Equal(t, scrapeconfig.DefaultScrapeConfig, getStdinConfig(util_log.Logger, []scrapeconfig.Config{
   145  		scrapeconfig.DefaultScrapeConfig,
   146  		{},
   147  	}))
   148  	// or use the default if none if provided
   149  	require.Equal(t, defaultStdInCfg, getStdinConfig(util_log.Logger, []scrapeconfig.Config{}))
   150  }
   151  
   152  var stagesConfig = `
   153  pipeline_stages:
   154  - template:
   155      source: new_key
   156      template: 'hello world!'
   157  - labels:
   158      new_key:
   159  `
   160  
   161  func loadConfig(yml string) stages.PipelineStages {
   162  	var config map[string]interface{}
   163  	err := yaml.Unmarshal([]byte(yml), &config)
   164  	if err != nil {
   165  		panic(err)
   166  	}
   167  	return config["pipeline_stages"].([]interface{})
   168  }