diff --git a/cmd/query/app/grpc_handler_test.go b/cmd/query/app/grpc_handler_test.go index 52a989873168..abba3a3628c3 100644 --- a/cmd/query/app/grpc_handler_test.go +++ b/cmd/query/app/grpc_handler_test.go @@ -196,47 +196,41 @@ func withServerAndClient(t *testing.T, actualTest func(server *grpcServer, clien } func TestGetTraceSuccessGRPC(t *testing.T) { - withServerAndClient(t, func(server *grpcServer, client *grpcClient) { - expectedTraceGetParameters := spanstore.GetTraceParameters{ - TraceID: mockTraceID, - } - server.spanReader.On("GetTrace", mock.AnythingOfType("*context.valueCtx"), expectedTraceGetParameters). - Return(mockTrace, nil).Once() - - res, err := client.GetTrace(context.Background(), &api_v2.GetTraceRequest{ - TraceID: mockTraceID, - }) + inputs := []struct { + expectedQuery spanstore.GetTraceParameters + request api_v2.GetTraceRequest + }{ + { + spanstore.GetTraceParameters{TraceID: mockTraceID}, + api_v2.GetTraceRequest{TraceID: mockTraceID}, + }, + { + spanstore.GetTraceParameters{ + TraceID: mockTraceID, + StartTime: startTime, + EndTime: endTime, + }, + api_v2.GetTraceRequest{ + TraceID: mockTraceID, + StartTime: startTime, + EndTime: endTime, + }, + }, + } - spanResChunk, _ := res.Recv() + for _, input := range inputs { + withServerAndClient(t, func(server *grpcServer, client *grpcClient) { + server.spanReader.On("GetTrace", mock.AnythingOfType("*context.valueCtx"), input.expectedQuery). + Return(mockTrace, nil).Once() - require.NoError(t, err) - assert.Equal(t, spanResChunk.Spans[0].TraceID, mockTraceID) - }) -} + res, err := client.GetTrace(context.Background(), &input.request) -func TestGetTraceWithTimeWindowSuccessGRPC(t *testing.T) { - withServerAndClient(t, func(server *grpcServer, client *grpcClient) { - startTime := time.Date(2020, time.January, 1, 13, 0, 0, 0, time.UTC) - endTime := time.Date(2020, time.January, 1, 14, 0, 0, 0, time.UTC) - expectedTraceGetParameters := spanstore.GetTraceParameters{ - TraceID: mockTraceID, - StartTime: startTime, - EndTime: endTime, - } - server.spanReader.On("GetTrace", mock.AnythingOfType("*context.valueCtx"), expectedTraceGetParameters). - Return(mockTrace, nil).Once() + spanResChunk, _ := res.Recv() - res, err := client.GetTrace(context.Background(), &api_v2.GetTraceRequest{ - TraceID: mockTraceID, - StartTime: startTime, - EndTime: endTime, + require.NoError(t, err) + assert.Equal(t, spanResChunk.Spans[0].TraceID, mockTraceID) }) - - spanResChunk, _ := res.Recv() - - require.NoError(t, err) - assert.Equal(t, spanResChunk.Spans[0].TraceID, mockTraceID) - }) + } } func assertGRPCError(t *testing.T, err error, code codes.Code, msg string) { @@ -305,42 +299,39 @@ func TestGetTraceNilRequestOnHandlerGRPC(t *testing.T) { } func TestArchiveTraceSuccessGRPC(t *testing.T) { - withServerAndClient(t, func(server *grpcServer, client *grpcClient) { - server.spanReader.On("GetTrace", mock.AnythingOfType("*context.valueCtx"), mock.AnythingOfType("spanstore.GetTraceParameters")). - Return(mockTrace, nil).Once() - server.archiveSpanWriter.On("WriteSpan", mock.AnythingOfType("*context.valueCtx"), mock.AnythingOfType("*model.Span")). - Return(nil).Times(2) - - _, err := client.ArchiveTrace(context.Background(), &api_v2.ArchiveTraceRequest{ - TraceID: mockTraceID, - }) - - require.NoError(t, err) - }) -} + inputs := []struct { + expectedQuery spanstore.GetTraceParameters + request api_v2.ArchiveTraceRequest + }{ + { + spanstore.GetTraceParameters{TraceID: mockTraceID}, + api_v2.ArchiveTraceRequest{TraceID: mockTraceID}, + }, + { + spanstore.GetTraceParameters{ + TraceID: mockTraceID, + StartTime: startTime, + EndTime: endTime, + }, + api_v2.ArchiveTraceRequest{ + TraceID: mockTraceID, + StartTime: &startTime, + EndTime: &endTime, + }, + }, + } + for _, input := range inputs { + withServerAndClient(t, func(server *grpcServer, client *grpcClient) { + server.spanReader.On("GetTrace", mock.AnythingOfType("*context.valueCtx"), input.expectedQuery). + Return(mockTrace, nil).Once() + server.archiveSpanWriter.On("WriteSpan", mock.AnythingOfType("*context.valueCtx"), mock.AnythingOfType("*model.Span")). + Return(nil).Times(2) -func TestArchiveTraceWithTimeWindowSuccessGRPC(t *testing.T) { - startTime := time.Unix(1, 2).UTC() - endTime := time.Unix(3, 4).UTC() - withServerAndClient(t, func(server *grpcServer, client *grpcClient) { - expectedGetTraceParameters := spanstore.GetTraceParameters{ - TraceID: mockTraceID, - StartTime: startTime, - EndTime: endTime, - } - server.spanReader.On("GetTrace", mock.AnythingOfType("*context.valueCtx"), expectedGetTraceParameters). - Return(mockTrace, nil).Once() - server.archiveSpanWriter.On("WriteSpan", mock.AnythingOfType("*context.valueCtx"), mock.AnythingOfType("*model.Span")). - Return(nil).Times(2) + _, err := client.ArchiveTrace(context.Background(), &input.request) - _, err := client.ArchiveTrace(context.Background(), &api_v2.ArchiveTraceRequest{ - TraceID: mockTraceID, - StartTime: &startTime, - EndTime: &endTime, + require.NoError(t, err) }) - - require.NoError(t, err) - }) + } } func TestArchiveTraceNotFoundGRPC(t *testing.T) { diff --git a/cmd/query/app/http_handler_test.go b/cmd/query/app/http_handler_test.go index 32c15d6aa91b..3fc33f6bebff 100644 --- a/cmd/query/app/http_handler_test.go +++ b/cmd/query/app/http_handler_test.go @@ -65,6 +65,8 @@ var ( } mockTraceID = model.NewTraceID(0, 123456) + startTime = time.Date(2020, time.January, 1, 13, 0, 0, 0, time.UTC) + endTime = time.Date(2020, time.January, 1, 14, 0, 0, 0, time.UTC) mockTrace = &model.Trace{ Spans: []*model.Span{ { diff --git a/plugin/storage/cassandra/spanstore/reader_test.go b/plugin/storage/cassandra/spanstore/reader_test.go index a55ced3e7ddd..afbcf48f0fbc 100644 --- a/plugin/storage/cassandra/spanstore/reader_test.go +++ b/plugin/storage/cassandra/spanstore/reader_test.go @@ -171,7 +171,7 @@ func TestSpanReaderGetTrace(t *testing.T) { r.session.On("Query", mock.AnythingOfType("string"), matchEverything()).Return(query) - trace, err := r.reader.GetTrace(context.Background(), spanstore.GetTraceParameters{TraceID: model.TraceID{}}) + trace, err := r.reader.GetTrace(context.Background(), spanstore.GetTraceParameters{}) if testCase.expectedErr == "" { require.NotEmpty(t, r.traceBuffer.GetSpans(), "Spans recorded") require.NoError(t, err) @@ -197,7 +197,7 @@ func TestSpanReaderGetTrace_TraceNotFound(t *testing.T) { r.session.On("Query", mock.AnythingOfType("string"), matchEverything()).Return(query) - trace, err := r.reader.GetTrace(context.Background(), spanstore.GetTraceParameters{TraceID: model.TraceID{}}) + trace, err := r.reader.GetTrace(context.Background(), spanstore.GetTraceParameters{}) require.NotEmpty(t, r.traceBuffer.GetSpans(), "Spans recorded") assert.Nil(t, trace) require.EqualError(t, err, "trace not found") diff --git a/plugin/storage/es/spanstore/reader_test.go b/plugin/storage/es/spanstore/reader_test.go index a91baafc2bd6..3cc8292e6405 100644 --- a/plugin/storage/es/spanstore/reader_test.go +++ b/plugin/storage/es/spanstore/reader_test.go @@ -1277,7 +1277,7 @@ func TestSpanReader_ArchiveTraces(t *testing.T) { Return(&elastic.MultiSearchResult{ Responses: []*elastic.SearchResult{}, }, nil) - query := spanstore.GetTraceParameters{TraceID: model.TraceID{}} + query := spanstore.GetTraceParameters{} trace, err := r.reader.GetTrace(context.Background(), query) require.NotEmpty(t, r.traceBuffer.GetSpans(), "Spans recorded") require.Nil(t, trace) @@ -1294,7 +1294,7 @@ func TestSpanReader_ArchiveTraces_ReadAlias(t *testing.T) { Responses: []*elastic.SearchResult{}, }, nil) - query := spanstore.GetTraceParameters{TraceID: model.TraceID{}} + query := spanstore.GetTraceParameters{} trace, err := r.reader.GetTrace(context.Background(), query) require.NotEmpty(t, r.traceBuffer.GetSpans(), "Spans recorded") require.Nil(t, trace) diff --git a/plugin/storage/memory/memory_test.go b/plugin/storage/memory/memory_test.go index 84de6cbd9178..af130edd5662 100644 --- a/plugin/storage/memory/memory_test.go +++ b/plugin/storage/memory/memory_test.go @@ -227,7 +227,7 @@ func TestStoreGetTraceError(t *testing.T) { func TestStoreGetTraceFailure(t *testing.T) { withPopulatedMemoryStore(func(store *Store) { - query := spanstore.GetTraceParameters{TraceID: model.TraceID{}} + query := spanstore.GetTraceParameters{} trace, err := store.GetTrace(context.Background(), query) require.EqualError(t, err, spanstore.ErrTraceNotFound.Error()) assert.Nil(t, trace) diff --git a/storage/spanstore/metrics/decorator.go b/storage/spanstore/metrics/decorator.go index 73b88b3d2cdf..ab7e076aa97d 100644 --- a/storage/spanstore/metrics/decorator.go +++ b/storage/spanstore/metrics/decorator.go @@ -80,7 +80,7 @@ func (m *ReadMetricsDecorator) FindTraceIDs(ctx context.Context, traceQuery *spa // GetTrace implements spanstore.Reader#GetTrace func (m *ReadMetricsDecorator) GetTrace(ctx context.Context, query spanstore.GetTraceParameters) (*model.Trace, error) { start := time.Now() - retMe, err := m.spanReader.GetTrace(ctx, traceGet) + retMe, err := m.spanReader.GetTrace(ctx, query) m.getTraceMetrics.emit(err, time.Since(start), 1) return retMe, err } diff --git a/storage/spanstore/metrics/decorator_test.go b/storage/spanstore/metrics/decorator_test.go index 48e64381ae34..7dfee3385c07 100644 --- a/storage/spanstore/metrics/decorator_test.go +++ b/storage/spanstore/metrics/decorator_test.go @@ -29,7 +29,7 @@ func TestSuccessfulUnderlyingCalls(t *testing.T) { mockReader.On("GetOperations", context.Background(), operationQuery). Return([]spanstore.Operation{}, nil) mrs.GetOperations(context.Background(), operationQuery) - mockReader.On("GetTrace", context.Background(), spanstore.GetTraceParameters{TraceID: model.TraceID{}}).Return(&model.Trace{}, nil) + mockReader.On("GetTrace", context.Background(), spanstore.GetTraceParameters{}).Return(&model.Trace{}, nil) mrs.GetTrace(context.Background(), spanstore.GetTraceParameters{}) mockReader.On("FindTraces", context.Background(), &spanstore.TraceQueryParameters{}). Return([]*model.Trace{}, nil) @@ -97,9 +97,9 @@ func TestFailingUnderlyingCalls(t *testing.T) { mockReader.On("GetOperations", context.Background(), operationQuery). Return(nil, errors.New("Failure")) mrs.GetOperations(context.Background(), operationQuery) - mockReader.On("GetTrace", context.Background(), spanstore.GetTraceParameters{TraceID: model.TraceID{}}). + mockReader.On("GetTrace", context.Background(), spanstore.GetTraceParameters{}). Return(nil, errors.New("Failure")) - mrs.GetTrace(context.Background(), spanstore.GetTraceParameters{TraceID: model.TraceID{}}) + mrs.GetTrace(context.Background(), spanstore.GetTraceParameters{}) mockReader.On("FindTraces", context.Background(), &spanstore.TraceQueryParameters{}). Return(nil, errors.New("Failure")) mrs.FindTraces(context.Background(), &spanstore.TraceQueryParameters{})