fix test

Signed-off-by: rim99 <[email protected]>
jaegertracing · Nov 21, 2024 · 93825bb · 93825bb
1 parent b3fc472
commit 93825bb
Show file tree

Hide file tree

Showing 4 changed files with 52 additions and 21 deletions.
diff --git a/cmd/query/app/grpc_handler.go b/cmd/query/app/grpc_handler.go
@@ -115,9 +115,13 @@ func (g *GRPCHandler) ArchiveTrace(ctx context.Context, r *api_v2.ArchiveTraceRe
 		return nil, errUninitializedTraceID
 	}
 	query := spanstore.TraceGetParameters{
-		TraceID:   r.TraceID,
-		StartTime: *r.StartTime,
-		EndTime:   *r.EndTime,
+		TraceID: r.TraceID,
+	}
+	if r.StartTime != nil {
+		query.StartTime = *r.StartTime
+	}
+	if r.EndTime != nil {
+		query.EndTime = *r.EndTime
 	}
 	err := g.queryService.ArchiveTrace(ctx, query)
 	if errors.Is(err, spanstore.ErrTraceNotFound) {

diff --git a/plugin/storage/grpc/shared/archive_test.go b/plugin/storage/grpc/shared/archive_test.go
@@ -7,6 +7,7 @@ import (
 	"context"
 	"io"
 	"testing"
+	"time"
 
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/mock"
@@ -55,7 +56,9 @@ func TestArchiveReader_GetTrace(t *testing.T) {
 
 	archiveSpanReader := new(mocks.ArchiveSpanReaderPluginClient)
 	archiveSpanReader.On("GetArchiveTrace", mock.Anything, &storage_v1.GetTraceRequest{
-		TraceID: mockTraceID,
+		TraceID:   mockTraceID,
+		StartTime: &time.Time{},
+		EndTime:   &time.Time{},
 	}).Return(traceClient, nil)
 	reader := &archiveReader{client: archiveSpanReader}
 
@@ -71,7 +74,9 @@ func TestArchiveReaderGetTrace_NoTrace(t *testing.T) {
 
 	archiveSpanReader := new(mocks.ArchiveSpanReaderPluginClient)
 	archiveSpanReader.On("GetArchiveTrace", mock.Anything, &storage_v1.GetTraceRequest{
-		TraceID: mockTraceID,
+		TraceID:   mockTraceID,
+		StartTime: &time.Time{},
+		EndTime:   &time.Time{},
 	}).Return(nil, status.Errorf(codes.NotFound, ""))
 	reader := &archiveReader{client: archiveSpanReader}
 

diff --git a/plugin/storage/grpc/shared/grpc_client_test.go b/plugin/storage/grpc/shared/grpc_client_test.go
@@ -192,7 +192,9 @@ func TestGRPCClientGetTrace_StreamError(t *testing.T) {
 		traceClient := new(grpcMocks.SpanReaderPlugin_GetTraceClient)
 		traceClient.On("Recv").Return(nil, errors.New("an error"))
 		r.spanReader.On("GetTrace", mock.Anything, &storage_v1.GetTraceRequest{
-			TraceID: mockTraceID,
+			TraceID:   mockTraceID,
+			StartTime: &time.Time{},
+			EndTime:   &time.Time{},
 		}).Return(traceClient, nil)
 
 		s, err := r.client.GetTrace(context.Background(), spanstore.TraceGetParameters{TraceID: mockTraceID})
@@ -204,7 +206,9 @@ func TestGRPCClientGetTrace_StreamError(t *testing.T) {
 func TestGRPCClientGetTrace_NoTrace(t *testing.T) {
 	withGRPCClient(func(r *grpcClientTest) {
 		r.spanReader.On("GetTrace", mock.Anything, &storage_v1.GetTraceRequest{
-			TraceID: mockTraceID,
+			TraceID:   mockTraceID,
+			StartTime: &time.Time{},
+			EndTime:   &time.Time{},
 		}).Return(nil, status.Errorf(codes.NotFound, ""))
 
 		s, err := r.client.GetTrace(context.Background(), spanstore.TraceGetParameters{TraceID: mockTraceID})
@@ -220,7 +224,9 @@ func TestGRPCClientGetTrace_StreamErrorTraceNotFound(t *testing.T) {
 		traceClient := new(grpcMocks.SpanReaderPlugin_GetTraceClient)
 		traceClient.On("Recv").Return(nil, s.Err())
 		r.spanReader.On("GetTrace", mock.Anything, &storage_v1.GetTraceRequest{
-			TraceID: mockTraceID,
+			TraceID:   mockTraceID,
+			StartTime: &time.Time{},
+			EndTime:   &time.Time{},
 		}).Return(traceClient, nil)
 
 		s, err := r.client.GetTrace(context.Background(), spanstore.TraceGetParameters{TraceID: mockTraceID})
@@ -407,7 +413,9 @@ func TestGrpcClientGetArchiveTrace_StreamError(t *testing.T) {
 		traceClient := new(grpcMocks.ArchiveSpanReaderPlugin_GetArchiveTraceClient)
 		traceClient.On("Recv").Return(nil, errors.New("an error"))
 		r.archiveReader.On("GetArchiveTrace", mock.Anything, &storage_v1.GetTraceRequest{
-			TraceID: mockTraceID,
+			TraceID:   mockTraceID,
+			StartTime: &time.Time{},
+			EndTime:   &time.Time{},
 		}).Return(traceClient, nil)
 
 		s, err := r.client.ArchiveSpanReader().GetTrace(context.Background(), spanstore.TraceGetParameters{
@@ -421,7 +429,9 @@ func TestGrpcClientGetArchiveTrace_StreamError(t *testing.T) {
 func TestGrpcClientGetArchiveTrace_NoTrace(t *testing.T) {
 	withGRPCClient(func(r *grpcClientTest) {
 		r.archiveReader.On("GetArchiveTrace", mock.Anything, &storage_v1.GetTraceRequest{
-			TraceID: mockTraceID,
+			TraceID:   mockTraceID,
+			StartTime: &time.Time{},
+			EndTime:   &time.Time{},
 		}).Return(nil, spanstore.ErrTraceNotFound)
 
 		s, err := r.client.ArchiveSpanReader().GetTrace(context.Background(), spanstore.TraceGetParameters{
@@ -437,7 +447,9 @@ func TestGrpcClientGetArchiveTrace_StreamErrorTraceNotFound(t *testing.T) {
 		traceClient := new(grpcMocks.ArchiveSpanReaderPlugin_GetArchiveTraceClient)
 		traceClient.On("Recv").Return(nil, spanstore.ErrTraceNotFound)
 		r.archiveReader.On("GetArchiveTrace", mock.Anything, &storage_v1.GetTraceRequest{
-			TraceID: mockTraceID,
+			TraceID:   mockTraceID,
+			StartTime: &time.Time{},
+			EndTime:   &time.Time{},
 		}).Return(traceClient, nil)
 
 		s, err := r.client.ArchiveSpanReader().GetTrace(context.Background(), spanstore.TraceGetParameters{

diff --git a/plugin/storage/grpc/shared/grpc_handler.go b/plugin/storage/grpc/shared/grpc_handler.go
@@ -151,11 +151,16 @@ func (s *GRPCHandler) Close(context.Context, *storage_v1.CloseWriterRequest) (*s
 
 // GetTrace takes a traceID and streams a Trace associated with that traceID
 func (s *GRPCHandler) GetTrace(r *storage_v1.GetTraceRequest, stream storage_v1.SpanReaderPlugin_GetTraceServer) error {
-	trace, err := s.impl.SpanReader().GetTrace(stream.Context(), spanstore.TraceGetParameters{
-		TraceID:   r.TraceID,
-		StartTime: *r.StartTime,
-		EndTime:   *r.EndTime,
-	})
+	query := spanstore.TraceGetParameters{
+		TraceID: r.TraceID,
+	}
+	if r.StartTime != nil {
+		query.StartTime = *r.StartTime
+	}
+	if r.EndTime != nil {
+		query.EndTime = *r.EndTime
+	}
+	trace, err := s.impl.SpanReader().GetTrace(stream.Context(), query)
 	if errors.Is(err, spanstore.ErrTraceNotFound) {
 		return status.Error(codes.NotFound, spanstore.ErrTraceNotFound.Error())
 	}
@@ -280,11 +285,16 @@ func (s *GRPCHandler) GetArchiveTrace(r *storage_v1.GetTraceRequest, stream stor
 	if reader == nil {
 		return status.Error(codes.Unimplemented, "not implemented")
 	}
-	trace, err := reader.GetTrace(stream.Context(), spanstore.TraceGetParameters{
-		TraceID:   r.TraceID,
-		StartTime: *r.StartTime,
-		EndTime:   *r.EndTime,
-	})
+	query := spanstore.TraceGetParameters{
+		TraceID: r.TraceID,
+	}
+	if r.StartTime != nil {
+		query.StartTime = *r.StartTime
+	}
+	if r.EndTime != nil {
+		query.EndTime = *r.EndTime
+	}
+	trace, err := reader.GetTrace(stream.Context(), query)
 	if errors.Is(err, spanstore.ErrTraceNotFound) {
 		return status.Error(codes.NotFound, spanstore.ErrTraceNotFound.Error())
 	}