update unit tests

Signed-off-by: rim99 <[email protected]>
jaegertracing · Nov 24, 2024 · 3eddefd · 3eddefd
1 parent fe42be6
commit 3eddefd
Show file tree

Hide file tree

Showing 7 changed files with 71 additions and 78 deletions.
diff --git 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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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{})