Add tests for anonymizer/app/query.

Signed-off-by: Pratham <prathms007@gmail.com>

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_test.go

@@ -0,0 +1,145 @@
+// 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"
+	"google.golang.org/grpc/credentials/insecure"
+
+	"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 newQuery(t *testing.T, addr string) *Query {
+	conn, err := grpc.NewClient(
+		addr,
+		grpc.WithTransportCredentials(insecure.NewCredentials()),
+	)
+	require.NoError(t, err)
+	t.Cleanup(func() {
+		conn.Close()
+	})
+
+	return &Query{
+		client: api_v2.NewQueryServiceClient(conn),
+		conn:   conn,
+	}
+}
+
+func newTestServer(t *testing.T) *testServer {
+	spanReader := &spanstoremocks.Reader{}
+	disabledReader, err := disabled.NewMetricsReader()
+	require.NoError(t, err)
+
+	q := querysvc.NewQueryService(
+		spanReader,
+		&dependencyStoreMocks.Reader{},
+		querysvc.QueryServiceOptions{},
+	)
+	h := app.NewGRPCHandler(q, disabledReader, 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.conn.Close()
+
+	assert.NotNil(t, query)
+}
+
+func TestQueryTrace(t *testing.T) {
+	s := newTestServer(t)
+	q := newQuery(t, s.address.String())
+
+	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)
+		require.NotNil(t, spans)
+
+		spansArr := make([]model.Span, 0, len(mockTraceGRPC.Spans))
+		for _, span := range mockTraceGRPC.Spans {
+			spansArr = append(spansArr, *span)
+		}
+		assert.Equal(t, spansArr, 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)
+	})
+}