mirror of
https://github.com/hanzoai/traces.git
synced 2026-10-30 18:20:41 +00:00
Oleg V. KozlyukandGitHub
b1dc89c38b
* Add support for protobuf format in all endpoints * go fmt * Added CHANGELOG entry * Added explicit MarshallingFormat type * Finalize refactoring - use MarshallingFormat type everywhere * Take review notes into account * Reduce code duplication * More review notes
127 lines
3.5 KiB
Go
127 lines
3.5 KiB
Go
package combiner
|
|
|
|
import (
|
|
"bytes"
|
|
"io"
|
|
"net/http"
|
|
"testing"
|
|
|
|
"github.com/gogo/protobuf/jsonpb"
|
|
"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/assert"
|
|
"github.com/stretchr/testify/require"
|
|
)
|
|
|
|
type MockResponse struct {
|
|
resp *http.Response
|
|
}
|
|
|
|
func (m MockResponse) HTTPResponse() *http.Response {
|
|
return m.resp
|
|
}
|
|
|
|
func (m MockResponse) RequestData() any {
|
|
return nil
|
|
}
|
|
|
|
func (m MockResponse) IsMetadata() bool {
|
|
return false
|
|
}
|
|
|
|
func TestNewTraceByIdV2ReturnsAPartialTrace(t *testing.T) {
|
|
traceResponse := &tempopb.TraceByIDResponse{
|
|
Trace: test.MakeTrace(2, []byte{0x01, 0x02}),
|
|
Metrics: &tempopb.TraceByIDMetrics{},
|
|
}
|
|
resBytes, err := proto.Marshal(traceResponse)
|
|
require.NoError(t, err)
|
|
response := http.Response{
|
|
StatusCode: 200,
|
|
Header: map[string][]string{
|
|
"Content-Type": {"application/protobuf"},
|
|
},
|
|
Body: io.NopCloser(bytes.NewReader(resBytes)),
|
|
}
|
|
combiner := NewTraceByIDV2(10, api.HeaderAcceptJSON, nil)
|
|
err = combiner.AddResponse(MockResponse{&response})
|
|
require.NoError(t, err)
|
|
|
|
res, err := combiner.HTTPFinal()
|
|
require.NoError(t, err)
|
|
|
|
actualResp := &tempopb.TraceByIDResponse{}
|
|
err = new(jsonpb.Unmarshaler).Unmarshal(res.Body, actualResp)
|
|
require.NoError(t, err)
|
|
assert.Equal(t, actualResp.Status, tempopb.PartialStatus_PARTIAL)
|
|
}
|
|
|
|
func TestNewTraceByIdV2ReturnsAPartialTraceOnPartialTraceReturnedByQuerier(t *testing.T) {
|
|
traceResponse := &tempopb.TraceByIDResponse{
|
|
Trace: test.MakeTrace(2, []byte{0x01, 0x02}),
|
|
Status: tempopb.PartialStatus_PARTIAL,
|
|
Metrics: &tempopb.TraceByIDMetrics{},
|
|
}
|
|
resBytes, err := proto.Marshal(traceResponse)
|
|
require.NoError(t, err)
|
|
response := http.Response{
|
|
StatusCode: 200,
|
|
Header: map[string][]string{
|
|
"Content-Type": {"application/protobuf"},
|
|
},
|
|
Body: io.NopCloser(bytes.NewReader(resBytes)),
|
|
}
|
|
combiner := NewTraceByIDV2(10, api.HeaderAcceptJSON, nil)
|
|
err = combiner.AddResponse(MockResponse{&response})
|
|
require.NoError(t, err)
|
|
|
|
res, err := combiner.HTTPFinal()
|
|
require.NoError(t, err)
|
|
|
|
actualResp := &tempopb.TraceByIDResponse{}
|
|
err = new(jsonpb.Unmarshaler).Unmarshal(res.Body, actualResp)
|
|
require.NoError(t, err)
|
|
assert.Equal(t, actualResp.Status, tempopb.PartialStatus_PARTIAL)
|
|
}
|
|
|
|
func TestNewTraceByIDV2(t *testing.T) {
|
|
traceResponse := &tempopb.TraceByIDResponse{
|
|
Trace: test.MakeTrace(2, []byte{0x01, 0x02}),
|
|
Metrics: &tempopb.TraceByIDMetrics{},
|
|
}
|
|
resBytes, err := proto.Marshal(traceResponse)
|
|
require.NoError(t, err)
|
|
response := http.Response{
|
|
StatusCode: 200,
|
|
Header: map[string][]string{
|
|
"Content-Type": {"application/protobuf"},
|
|
},
|
|
Body: io.NopCloser(bytes.NewReader(resBytes)),
|
|
}
|
|
|
|
t.Run("returns a combined trace response as JSON", func(t *testing.T) {
|
|
combiner := NewTraceByIDV2(100_000, api.HeaderAcceptJSON, nil)
|
|
err = combiner.AddResponse(MockResponse{&response})
|
|
require.NoError(t, err)
|
|
|
|
res, err := combiner.HTTPFinal()
|
|
require.NoError(t, err)
|
|
assert.Equal(t, api.HeaderAcceptJSON, res.Header.Get(api.HeaderContentType))
|
|
|
|
actualResp := &tempopb.TraceByIDResponse{}
|
|
err = new(jsonpb.Unmarshaler).Unmarshal(res.Body, actualResp)
|
|
require.NoError(t, err)
|
|
})
|
|
t.Run("returns a combined trace response as protobuff", func(t *testing.T) {
|
|
combiner := NewTraceByIDV2(100_000, api.HeaderAcceptProtobuf, nil)
|
|
err = combiner.AddResponse(MockResponse{&response})
|
|
require.NoError(t, err)
|
|
|
|
res, err := combiner.HTTPFinal()
|
|
require.NoError(t, err)
|
|
require.NotNil(t, res)
|
|
})
|
|
}
|