Add tests for anonymizer/app/query.

cmd/anonymizer/app/query/package_test.go

@@ -0,0 +1,14 @@
+// Copyright (c) 2024 The Jaeger Authors.
+// SPDX-License-Identifier: Apache-2.0
+
+package query
+
+import (
+	"testing"
+
+	"github.com/jaegertracing/jaeger/pkg/testutils"
+)
+
+func TestMain(m *testing.M) {
+	testutils.VerifyGoLeaks(m)
+}

cmd/anonymizer/app/query/query.go

@@ -19,6 +19,7 @@ import (
 	"errors"
 	"fmt"
 	"io"
+	"strings"
 
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/credentials/insecure"
@@ -51,7 +52,7 @@ func New(addr string) (*Query, error) {
 // unwrapNotFoundErr is a conversion function
 func unwrapNotFoundErr(err error) error {
 	if s, _ := status.FromError(err); s != nil {
-		if s.Message() == spanstore.ErrTraceNotFound.Error() {
+		if strings.Contains(s.Message(), spanstore.ErrTraceNotFound.Error()) {
 			return spanstore.ErrTraceNotFound
 		}
 	}
@@ -82,3 +83,8 @@ func (q *Query) QueryTrace(traceID string) ([]model.Span, error) {
 
 	return spans, nil
 }
+
+// Close closes the grpc client connection
+func (q *Query) Close() error {
+	return q.conn.Close()
+}

cmd/anonymizer/app/query/query_test.go

@@ -0,0 +1,124 @@
+// Copyright (c) 2024 The Jaeger Authors.
+// SPDX-License-Identifier: Apache-2.0
+
+package query
+
+import (
+	"net"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/mock"
+	"github.com/stretchr/testify/require"
+	"google.golang.org/grpc"
+
+	"github.com/jaegertracing/jaeger/cmd/query/app"
+	"github.com/jaegertracing/jaeger/cmd/query/app/querysvc"
+	"github.com/jaegertracing/jaeger/model"
+	"github.com/jaegertracing/jaeger/plugin/metrics/disabled"
+	"github.com/jaegertracing/jaeger/proto-gen/api_v2"
+	dependencyStoreMocks "github.com/jaegertracing/jaeger/storage/dependencystore/mocks"
+	"github.com/jaegertracing/jaeger/storage/spanstore"
+	spanstoremocks "github.com/jaegertracing/jaeger/storage/spanstore/mocks"
+)
+
+var (
+	matchContext = mock.AnythingOfType("*context.valueCtx")
+	matchTraceID = mock.AnythingOfType("model.TraceID")
+
+	mockInvalidTraceID = "xyz"
+	mockTraceID        = model.NewTraceID(0, 123456)
+
+	mockTraceGRPC = &model.Trace{
+		Spans: []*model.Span{
+			{
+				TraceID: mockTraceID,
+				SpanID:  model.NewSpanID(1),
+				Process: &model.Process{},
+			},
+			{
+				TraceID: mockTraceID,
+				SpanID:  model.NewSpanID(2),
+				Process: &model.Process{},
+			},
+		},
+		Warnings: []string{},
+	}
+)
+
+type testServer struct {
+	address    net.Addr
+	server     *grpc.Server
+	spanReader *spanstoremocks.Reader
+}
+
+func newTestServer(t *testing.T) *testServer {
+	spanReader := &spanstoremocks.Reader{}
+	metricsReader, err := disabled.NewMetricsReader()
+	require.NoError(t, err)
+
+	q := querysvc.NewQueryService(
+		spanReader,
+		&dependencyStoreMocks.Reader{},
+		querysvc.QueryServiceOptions{},
+	)
+	h := app.NewGRPCHandler(q, metricsReader, app.GRPCHandlerOptions{})
+
+	server := grpc.NewServer()
+	api_v2.RegisterQueryServiceServer(server, h)
+
+	lis, err := net.Listen("tcp", ":0")
+	require.NoError(t, err)
+
+	go func() {
+		err := server.Serve(lis)
+		require.NoError(t, err)
+	}()
+	t.Cleanup(func() { server.Stop() })
+
+	return &testServer{
+		server:     server,
+		address:    lis.Addr(),
+		spanReader: spanReader,
+	}
+}
+
+func TestNew(t *testing.T) {
+	server := newTestServer(t)
+
+	query, err := New(server.address.String())
+	require.NoError(t, err)
+	defer query.Close()
+
+	assert.NotNil(t, query)
+}
+
+func TestQueryTrace(t *testing.T) {
+	s := newTestServer(t)
+	q, err := New(s.address.String())
+	require.NoError(t, err)
+	defer q.Close()
+
+	t.Run("No error", func(t *testing.T) {
+		s.spanReader.On("GetTrace", matchContext, matchTraceID).Return(
+			mockTraceGRPC, nil).Once()
+
+		spans, err := q.QueryTrace(mockTraceID.String())
+		require.NoError(t, err)
+		assert.Equal(t, len(spans), len(mockTraceGRPC.Spans))
+	})
+
+	t.Run("Invalid TraceID", func(t *testing.T) {
+		_, err := q.QueryTrace(mockInvalidTraceID)
+		assert.ErrorContains(t, err, "failed to convert the provided trace id")
+	})
+
+	t.Run("Trace not found", func(t *testing.T) {
+		s.spanReader.On("GetTrace", matchContext, matchTraceID).Return(
+			nil, spanstore.ErrTraceNotFound).Once()
+
+		spans, err := q.QueryTrace(mockTraceID.String())
+		assert.Nil(t, spans)
+		assert.ErrorIs(t, err, spanstore.ErrTraceNotFound)
+	})
+}

cmd/anonymizer/main.go

@@ -68,6 +68,9 @@ func main() {
 			if err != nil {
 				logger.Fatal("error while querying for trace", zap.Error(err))
 			}
+			if err := query.Close(); err != nil {
+				logger.Error("Failed to close grpc client connection", zap.Error(err))
+			}
 
 			for _, span := range spans {
 				if err := w.WriteSpan(&span); err != nil {