From 93825bb6b6abb5003d56348f49e674640c20f881 Mon Sep 17 00:00:00 2001 From: rim99 Date: Thu, 21 Nov 2024 22:41:19 +0800 Subject: [PATCH] fix test Signed-off-by: rim99 --- cmd/query/app/grpc_handler.go | 10 +++++-- plugin/storage/grpc/shared/archive_test.go | 9 ++++-- .../storage/grpc/shared/grpc_client_test.go | 24 +++++++++++---- plugin/storage/grpc/shared/grpc_handler.go | 30 ++++++++++++------- 4 files changed, 52 insertions(+), 21 deletions(-) diff --git a/cmd/query/app/grpc_handler.go b/cmd/query/app/grpc_handler.go index 5fbf922c55e..6c1acba3dcd 100644 --- 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 index e1005b7153a..5ef2987a415 100644 --- 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 index 4649e7a5daa..4014407bf68 100644 --- 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 index a3dd0b84de1..66fb6312dbb 100644 --- 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()) }