tempo/modules/frontend/combiner/trace_by_id_test.go at 2de43d0ffb269e4e0299f48e0581014193bb5d3a · carles-grafana/tempo · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
package combiner

import (
	"bytes"
	"io"
	"net/http"
	"strings"
	"testing"

	"github.com/gogo/protobuf/proto"
	"github.com/grafana/tempo/pkg/api"
	"github.com/grafana/tempo/pkg/tempopb"
	"github.com/grafana/tempo/pkg/util/test"
	"github.com/stretchr/testify/require"
)

func TestTraceByIDShouldQuit(t *testing.T) {
	// new combiner should not quit
	c := NewTraceByID(0, api.HeaderAcceptJSON)
	should := c.ShouldQuit()
	require.False(t, should)

	// 500 response should quit
	c = NewTraceByID(0, api.HeaderAcceptJSON)
	err := c.AddResponse(toHTTPResponse(t, &tempopb.SearchResponse{}, 500))
	require.NoError(t, err)
	should = c.ShouldQuit()
	require.True(t, should)

	// 429 response should quit
	c = NewTraceByID(0, api.HeaderAcceptJSON)
	err = c.AddResponse(toHTTPProtoResponse(t, &tempopb.SearchResponse{}, 429))
	require.NoError(t, err)
	should = c.ShouldQuit()
	require.True(t, should)

	// 404 response should not quit
	c = NewTraceByID(0, api.HeaderAcceptJSON)
	err = c.AddResponse(toHTTPProtoResponse(t, &tempopb.SearchResponse{}, 404))
	require.NoError(t, err)
	should = c.ShouldQuit()
	require.False(t, should)

	// unparseable body should not quit, but should return an error
	c = NewTraceByID(0, api.HeaderAcceptJSON)
	err = c.AddResponse(&testPipelineResponse{r: &http.Response{Body: io.NopCloser(strings.NewReader("foo")), StatusCode: 200}})
	require.Error(t, err)
	should = c.ShouldQuit()
	require.False(t, should)

	// trace too large, should quit and should not return an error
	c = NewTraceByID(1, api.HeaderAcceptJSON)
	err = c.AddResponse(toHTTPProtoResponse(t, &tempopb.TraceByIDResponse{
		Trace:   test.MakeTrace(1, nil),
		Metrics: &tempopb.TraceByIDMetrics{},
	}, 200))
	require.NoError(t, err)
	should = c.ShouldQuit()
	require.True(t, should)
}

func TestTraceByIDHonorsContentType(t *testing.T) {
	expected := test.MakeTrace(2, nil)

	// json
	c := NewTraceByID(0, api.HeaderAcceptJSON)
	err := c.AddResponse(toHTTPProtoResponse(t, &tempopb.TraceByIDResponse{Trace: expected, Metrics: &tempopb.TraceByIDMetrics{InspectedBytes: 100}}, 200))
	require.NoError(t, err)

	resp, err := c.HTTPFinal()
	require.NoError(t, err)

	actual := &tempopb.Trace{}
	bodyBytes, _ := io.ReadAll(resp.Body)
	err = tempopb.UnmarshalFromJSONV1(bodyBytes, actual)
	require.NoError(t, err)
	require.Equal(t, expected, actual)

	// proto
	c = NewTraceByID(0, api.HeaderAcceptProtobuf)
	err = c.AddResponse(toHTTPProtoResponse(t, &tempopb.TraceByIDResponse{Trace: expected, Metrics: &tempopb.TraceByIDMetrics{InspectedBytes: 100}}, 200))
	require.NoError(t, err)

	resp, err = c.HTTPFinal()
	require.NoError(t, err)

	actual = &tempopb.Trace{}
	buff, err := io.ReadAll(resp.Body)
	require.NoError(t, err)
	err = proto.Unmarshal(buff, actual)
	require.NoError(t, err)
	require.Equal(t, expected, actual)
}

func toHTTPProtoResponse(t *testing.T, pb proto.Message, statusCode int) PipelineResponse {
	var body []byte

	if pb != nil {
		var err error
		body, err = proto.Marshal(pb)
		require.NoError(t, err)
	}

	return &testPipelineResponse{r: &http.Response{
		Body:       io.NopCloser(bytes.NewReader(body)),
		StatusCode: statusCode,
	}}
}