github.com/yankunsam/loki/v2@v2.6.3-0.20220817130409-389df5235c27/clients/cmd/fluent-bit/loki_test.go (about)

     1  package main
     2  
     3  import (
     4  	"errors"
     5  	"reflect"
     6  	"testing"
     7  	"time"
     8  
     9  	jsoniter "github.com/json-iterator/go"
    10  	"github.com/prometheus/common/model"
    11  
    12  	"github.com/grafana/loki/clients/pkg/promtail/api"
    13  	"github.com/grafana/loki/clients/pkg/promtail/client/fake"
    14  
    15  	"github.com/grafana/loki/pkg/logproto"
    16  )
    17  
    18  var now = time.Now()
    19  
    20  func Test_loki_sendRecord(t *testing.T) {
    21  	simpleRecordFixture := map[interface{}]interface{}{
    22  		"foo":   "bar",
    23  		"bar":   500,
    24  		"error": make(chan struct{}),
    25  	}
    26  	mapRecordFixture := map[interface{}]interface{}{
    27  		// lots of key/value pairs in map to increase chances of test hitting in case of unsorted map marshalling
    28  		"A": "A",
    29  		"B": "B",
    30  		"C": "C",
    31  		"D": "D",
    32  		"E": "E",
    33  		"F": "F",
    34  		"G": "G",
    35  		"H": "H",
    36  	}
    37  	byteArrayRecordFixture := map[interface{}]interface{}{
    38  		"label": "label",
    39  		"outer": []byte("foo"),
    40  		"map": map[interface{}]interface{}{
    41  			"inner": []byte("bar"),
    42  		},
    43  	}
    44  	mixedTypesRecordFixture := map[interface{}]interface{}{
    45  		"label": "label",
    46  		"int":   42,
    47  		"float": 42.42,
    48  		"array": []interface{}{42, 42.42, "foo"},
    49  		"map": map[interface{}]interface{}{
    50  			"nested": map[interface{}]interface{}{
    51  				"foo":     "bar",
    52  				"invalid": []byte("a\xc5z"),
    53  			},
    54  		},
    55  	}
    56  	nestedJSONFixture := map[interface{}]interface{}{
    57  		"kubernetes": map[interface{}]interface{}{
    58  			"annotations": map[interface{}]interface{}{
    59  				"kubernetes.io/psp":  "test",
    60  				"prometheus.io/port": "8085",
    61  			},
    62  		},
    63  		"log": "\tstatus code: 403, request id: b41c1ffa-c586-4359-a7da-457dd8da4bad\n",
    64  	}
    65  
    66  	tests := []struct {
    67  		name    string
    68  		cfg     *config
    69  		record  map[interface{}]interface{}
    70  		want    []api.Entry
    71  		wantErr bool
    72  	}{
    73  		{"map to JSON", &config{labelKeys: []string{"A"}, lineFormat: jsonFormat}, mapRecordFixture, []api.Entry{{Labels: model.LabelSet{"A": "A"}, Entry: logproto.Entry{Line: `{"B":"B","C":"C","D":"D","E":"E","F":"F","G":"G","H":"H"}`, Timestamp: now}}}, false},
    74  		{"map to kvPairFormat", &config{labelKeys: []string{"A"}, lineFormat: kvPairFormat}, mapRecordFixture, []api.Entry{{Labels: model.LabelSet{"A": "A"}, Entry: logproto.Entry{Line: `B=B C=C D=D E=E F=F G=G H=H`, Timestamp: now}}}, false},
    75  		{"not enough records", &config{labelKeys: []string{"foo"}, lineFormat: jsonFormat, removeKeys: []string{"bar", "error"}}, simpleRecordFixture, []api.Entry{}, false},
    76  		{"labels", &config{labelKeys: []string{"bar", "fake"}, lineFormat: jsonFormat, removeKeys: []string{"fuzz", "error"}}, simpleRecordFixture, []api.Entry{{Labels: model.LabelSet{"bar": "500"}, Entry: logproto.Entry{Line: `{"foo":"bar"}`, Timestamp: now}}}, false},
    77  		{"remove key", &config{labelKeys: []string{"fake"}, lineFormat: jsonFormat, removeKeys: []string{"foo", "error", "fake"}}, simpleRecordFixture, []api.Entry{{Labels: model.LabelSet{}, Entry: logproto.Entry{Line: `{"bar":500}`, Timestamp: now}}}, false},
    78  		{"error", &config{labelKeys: []string{"fake"}, lineFormat: jsonFormat, removeKeys: []string{"foo"}}, simpleRecordFixture, []api.Entry{}, true},
    79  		{"key value", &config{labelKeys: []string{"fake"}, lineFormat: kvPairFormat, removeKeys: []string{"foo", "error", "fake"}}, simpleRecordFixture, []api.Entry{{Labels: model.LabelSet{}, Entry: logproto.Entry{Line: `bar=500`, Timestamp: now}}}, false},
    80  		{"single", &config{labelKeys: []string{"fake"}, dropSingleKey: true, lineFormat: kvPairFormat, removeKeys: []string{"foo", "error", "fake"}}, simpleRecordFixture, []api.Entry{{Labels: model.LabelSet{}, Entry: logproto.Entry{Line: `500`, Timestamp: now}}}, false},
    81  		{"labelmap", &config{labelMap: map[string]interface{}{"bar": "other"}, lineFormat: jsonFormat, removeKeys: []string{"bar", "error"}}, simpleRecordFixture, []api.Entry{{Labels: model.LabelSet{"other": "500"}, Entry: logproto.Entry{Line: `{"foo":"bar"}`, Timestamp: now}}}, false},
    82  		{"byte array", &config{labelKeys: []string{"label"}, lineFormat: jsonFormat}, byteArrayRecordFixture, []api.Entry{{Labels: model.LabelSet{"label": "label"}, Entry: logproto.Entry{Line: `{"map":{"inner":"bar"},"outer":"foo"}`, Timestamp: now}}}, false},
    83  		{"mixed types", &config{labelKeys: []string{"label"}, lineFormat: jsonFormat}, mixedTypesRecordFixture, []api.Entry{{Labels: model.LabelSet{"label": "label"}, Entry: logproto.Entry{Line: `{"array":[42,42.42,"foo"],"float":42.42,"int":42,"map":{"nested":{"foo":"bar","invalid":"a\ufffdz"}}}`, Timestamp: now}}}, false},
    84  		{"JSON inner string escaping", &config{removeKeys: []string{"kubernetes"}, labelMap: map[string]interface{}{"kubernetes": map[string]interface{}{"annotations": map[string]interface{}{"kubernetes.io/psp": "label"}}}, lineFormat: jsonFormat}, nestedJSONFixture, []api.Entry{{Labels: model.LabelSet{"label": "test"}, Entry: logproto.Entry{Line: `{"log":"\tstatus code: 403, request id: b41c1ffa-c586-4359-a7da-457dd8da4bad\n"}`, Timestamp: now}}}, false},
    85  	}
    86  	for _, tt := range tests {
    87  		t.Run(tt.name, func(t *testing.T) {
    88  			rec := fake.New(func() {})
    89  			l := &loki{
    90  				cfg:    tt.cfg,
    91  				client: rec,
    92  				logger: logger,
    93  			}
    94  			err := l.sendRecord(tt.record, now)
    95  			if (err != nil) != tt.wantErr {
    96  				t.Errorf("sendRecord() error = %v, wantErr %v", err, tt.wantErr)
    97  				return
    98  			}
    99  			rec.Stop()
   100  			got := rec.Received()
   101  			if !reflect.DeepEqual(got, tt.want) {
   102  				t.Errorf("sendRecord() want:%v got:%v", tt.want, got)
   103  			}
   104  		})
   105  	}
   106  }
   107  
   108  func Test_createLine(t *testing.T) {
   109  	tests := []struct {
   110  		name    string
   111  		records map[string]interface{}
   112  		f       format
   113  		want    string
   114  		wantErr bool
   115  	}{
   116  
   117  		{"json", map[string]interface{}{"foo": "bar", "bar": map[string]interface{}{"bizz": "bazz"}}, jsonFormat, `{"foo":"bar","bar":{"bizz":"bazz"}}`, false},
   118  		{"json with number", map[string]interface{}{"foo": "bar", "bar": map[string]interface{}{"bizz": 20}}, jsonFormat, `{"foo":"bar","bar":{"bizz":20}}`, false},
   119  		{"bad json", map[string]interface{}{"foo": make(chan interface{})}, jsonFormat, "", true},
   120  		{"kv with space", map[string]interface{}{"foo": "bar", "bar": "foo foo"}, kvPairFormat, `bar="foo foo" foo=bar`, false},
   121  		{"kv with number", map[string]interface{}{"foo": "bar foo", "decimal": 12.2}, kvPairFormat, `decimal=12.2 foo="bar foo"`, false},
   122  		{"kv with nil", map[string]interface{}{"foo": "bar", "null": nil}, kvPairFormat, `foo=bar null=null`, false},
   123  		{"kv with array", map[string]interface{}{"foo": "bar", "array": []string{"foo", "bar"}}, kvPairFormat, `array="[foo bar]" foo=bar`, false},
   124  		{"kv with map", map[string]interface{}{"foo": "bar", "map": map[string]interface{}{"foo": "bar", "bar ": "foo "}}, kvPairFormat, `foo=bar map="map[bar :foo  foo:bar]"`, false},
   125  		{"kv empty", map[string]interface{}{}, kvPairFormat, ``, false},
   126  		{"bad format", nil, format(3), "", true},
   127  		{"nested json", map[string]interface{}{"log": `{"level":"error"}`}, jsonFormat, `{"log":{"level":"error"}}`, false},
   128  		{"nested json", map[string]interface{}{"log": `["level","error"]`}, jsonFormat, `{"log":["level","error"]}`, false},
   129  	}
   130  	for _, tt := range tests {
   131  		t.Run(tt.name, func(t *testing.T) {
   132  			l := &loki{
   133  				logger: logger,
   134  			}
   135  			got, err := l.createLine(tt.records, tt.f)
   136  			if (err != nil) != tt.wantErr {
   137  				t.Errorf("createLine() error = %v, wantErr %v", err, tt.wantErr)
   138  				return
   139  			}
   140  			if err != nil {
   141  				return
   142  			}
   143  			if tt.f == jsonFormat {
   144  				compareJSON(t, got, tt.want)
   145  			} else {
   146  				if got != tt.want {
   147  					t.Errorf("createLine() = %v, want %v", got, tt.want)
   148  				}
   149  			}
   150  		})
   151  	}
   152  }
   153  
   154  // compareJson unmarshal both string to map[string]interface compare json result.
   155  // we can't compare string to string as jsoniter doesn't ensure field ordering.
   156  func compareJSON(t *testing.T, got, want string) {
   157  	var w map[string]interface{}
   158  	err := jsoniter.Unmarshal([]byte(want), &w)
   159  	if err != nil {
   160  		t.Errorf("failed to unmarshal string: %s", err)
   161  	}
   162  	var g map[string]interface{}
   163  	err = jsoniter.Unmarshal([]byte(got), &g)
   164  	if err != nil {
   165  		t.Errorf("failed to unmarshal string: %s", err)
   166  	}
   167  	if !reflect.DeepEqual(g, w) {
   168  		t.Errorf("compareJson() = %v, want %v", g, w)
   169  	}
   170  }
   171  
   172  func Test_removeKeys(t *testing.T) {
   173  	tests := []struct {
   174  		name     string
   175  		records  map[string]interface{}
   176  		expected map[string]interface{}
   177  		keys     []string
   178  	}{
   179  		{"remove all keys", map[string]interface{}{"foo": "bar", "bar": map[string]interface{}{"bizz": "bazz"}}, map[string]interface{}{}, []string{"foo", "bar"}},
   180  		{"remove none", map[string]interface{}{"foo": "bar"}, map[string]interface{}{"foo": "bar"}, []string{}},
   181  		{"remove not existing", map[string]interface{}{"foo": "bar"}, map[string]interface{}{"foo": "bar"}, []string{"bar"}},
   182  		{"remove one", map[string]interface{}{"foo": "bar", "bazz": "buzz"}, map[string]interface{}{"foo": "bar"}, []string{"bazz"}},
   183  	}
   184  	for _, tt := range tests {
   185  		t.Run(tt.name, func(t *testing.T) {
   186  			removeKeys(tt.records, tt.keys)
   187  			if !reflect.DeepEqual(tt.expected, tt.records) {
   188  				t.Errorf("removeKeys() = %v, want %v", tt.records, tt.expected)
   189  			}
   190  		})
   191  	}
   192  }
   193  
   194  func Test_extractLabels(t *testing.T) {
   195  	tests := []struct {
   196  		name    string
   197  		records map[string]interface{}
   198  		keys    []string
   199  		want    model.LabelSet
   200  	}{
   201  		{"single string", map[string]interface{}{"foo": "bar", "bar": map[string]interface{}{"bizz": "bazz"}}, []string{"foo"}, model.LabelSet{"foo": "bar"}},
   202  		{"multiple", map[string]interface{}{"foo": "bar", "bar": map[string]interface{}{"bizz": "bazz"}}, []string{"foo", "bar"}, model.LabelSet{"foo": "bar", "bar": "map[bizz:bazz]"}},
   203  		{"nil", map[string]interface{}{"foo": nil}, []string{"foo"}, model.LabelSet{"foo": "<nil>"}},
   204  		{"none", map[string]interface{}{"foo": nil}, []string{}, model.LabelSet{}},
   205  		{"missing", map[string]interface{}{"foo": "bar"}, []string{"foo", "buzz"}, model.LabelSet{"foo": "bar"}},
   206  		{"skip invalid", map[string]interface{}{"foo.blah": "bar", "bar": "a\xc5z"}, []string{"foo.blah", "bar"}, model.LabelSet{}},
   207  	}
   208  	for _, tt := range tests {
   209  		t.Run(tt.name, func(t *testing.T) {
   210  			if got := extractLabels(tt.records, tt.keys); !reflect.DeepEqual(got, tt.want) {
   211  				t.Errorf("extractLabels() = %v, want %v", got, tt.want)
   212  			}
   213  		})
   214  	}
   215  }
   216  
   217  func Test_toStringMap(t *testing.T) {
   218  	tests := []struct {
   219  		name   string
   220  		record map[interface{}]interface{}
   221  		want   map[string]interface{}
   222  	}{
   223  		{"already string", map[interface{}]interface{}{"string": "foo", "bar": []byte("buzz")}, map[string]interface{}{"string": "foo", "bar": "buzz"}},
   224  		{"skip non string", map[interface{}]interface{}{"string": "foo", 1.0: []byte("buzz")}, map[string]interface{}{"string": "foo"}},
   225  		{
   226  			"byteslice in array",
   227  			map[interface{}]interface{}{"string": "foo", "bar": []interface{}{map[interface{}]interface{}{"baz": []byte("quux")}}},
   228  			map[string]interface{}{"string": "foo", "bar": []interface{}{map[string]interface{}{"baz": "quux"}}},
   229  		},
   230  	}
   231  	for _, tt := range tests {
   232  		t.Run(tt.name, func(t *testing.T) {
   233  			if got := toStringMap(tt.record); !reflect.DeepEqual(got, tt.want) {
   234  				t.Errorf("toStringMap() = %v, want %v", got, tt.want)
   235  			}
   236  		})
   237  	}
   238  }
   239  
   240  func Test_labelMapping(t *testing.T) {
   241  	tests := []struct {
   242  		name    string
   243  		records map[string]interface{}
   244  		mapping map[string]interface{}
   245  		want    model.LabelSet
   246  	}{
   247  		{
   248  			"empty record",
   249  			map[string]interface{}{},
   250  			map[string]interface{}{},
   251  			model.LabelSet{},
   252  		},
   253  		{
   254  			"empty subrecord",
   255  			map[string]interface{}{
   256  				"kubernetes": map[interface{}]interface{}{
   257  					"foo": []byte("buzz"),
   258  				},
   259  			},
   260  			map[string]interface{}{},
   261  			model.LabelSet{},
   262  		},
   263  		{
   264  			"deep string",
   265  			map[string]interface{}{
   266  				"int":   "42",
   267  				"float": "42.42",
   268  				"array": `[42,42.42,"foo"]`,
   269  				"kubernetes": map[string]interface{}{
   270  					"label": map[string]interface{}{
   271  						"component": map[string]interface{}{
   272  							"buzz": "value",
   273  						},
   274  					},
   275  				},
   276  			},
   277  			map[string]interface{}{
   278  				"int":   "int",
   279  				"float": "float",
   280  				"array": "array",
   281  				"kubernetes": map[string]interface{}{
   282  					"label": map[string]interface{}{
   283  						"component": map[string]interface{}{
   284  							"buzz": "label",
   285  						},
   286  					},
   287  				},
   288  				"stream": "output",
   289  				"nope":   "nope",
   290  			},
   291  			model.LabelSet{
   292  				"int":   "42",
   293  				"float": "42.42",
   294  				"array": `[42,42.42,"foo"]`,
   295  				"label": "value",
   296  			},
   297  		},
   298  	}
   299  	for _, tt := range tests {
   300  		t.Run(tt.name, func(t *testing.T) {
   301  			got := model.LabelSet{}
   302  			if mapLabels(tt.records, tt.mapping, got); !reflect.DeepEqual(got, tt.want) {
   303  				t.Errorf("mapLabels() = %v, want %v", got, tt.want)
   304  			}
   305  		})
   306  	}
   307  }
   308  
   309  func Test_AutoKubernetesLabels(t *testing.T) {
   310  	tests := []struct {
   311  		name    string
   312  		records map[interface{}]interface{}
   313  		want    model.LabelSet
   314  		err     error
   315  	}{
   316  		{
   317  			"records without labels",
   318  			map[interface{}]interface{}{
   319  				"kubernetes": map[interface{}]interface{}{
   320  					"foo": []byte("buzz"),
   321  				},
   322  			},
   323  			model.LabelSet{
   324  				"foo": "buzz",
   325  			},
   326  			nil,
   327  		},
   328  		{
   329  			"records with labels",
   330  			map[interface{}]interface{}{
   331  				"kubernetes": map[string]interface{}{
   332  					"labels": map[string]interface{}{
   333  						"foo":  "bar",
   334  						"buzz": "value",
   335  					},
   336  				},
   337  			},
   338  			model.LabelSet{
   339  				"foo":  "bar",
   340  				"buzz": "value",
   341  			},
   342  			nil,
   343  		},
   344  		{
   345  			"records without kubernetes labels",
   346  			map[interface{}]interface{}{
   347  				"foo":   "bar",
   348  				"label": "value",
   349  			},
   350  			model.LabelSet{},
   351  			errors.New("kubernetes labels not found, no labels will be added"),
   352  		},
   353  	}
   354  	for _, tt := range tests {
   355  		t.Run(tt.name, func(t *testing.T) {
   356  			m := toStringMap(tt.records)
   357  			lbs := model.LabelSet{}
   358  			err := autoLabels(m, lbs)
   359  			if err != nil && err.Error() != tt.err.Error() {
   360  				t.Errorf("error in autolabels, error = %v", err)
   361  				return
   362  			}
   363  			if !reflect.DeepEqual(lbs, tt.want) {
   364  				t.Errorf("mapLabels() = %v, want %v", lbs, tt.want)
   365  			}
   366  		})
   367  	}
   368  }