Add unit test for OTLP receiver in jaeger-v1

cmd/collector/app/handler/grpc_handler.go

@@ -7,6 +7,7 @@
 	"context"
 	"errors"
 
+	"go.opentelemetry.io/collector/client"
 	"go.uber.org/zap"
 	"google.golang.org/grpc/codes"
 	_ "google.golang.org/grpc/encoding/gzip" // register zip encoding
@@ -98,12 +99,19 @@
 		return "", nil
 	}
 
+	httpMd := client.FromContext(ctx)
+	tenants := httpMd.Metadata.Get(c.tenancyMgr.Header)
+
+	if len(tenants) > 0 {
+		return tenants[0], nil
+	}
+
 	md, ok := metadata.FromIncomingContext(ctx)
 	if !ok {
 		return "", status.Errorf(codes.PermissionDenied, "missing tenant header")
 	}
 
-	tenants := md.Get(c.tenancyMgr.Header)
+	tenants = md.Get(c.tenancyMgr.Header)
 	if len(tenants) < 1 {
 		return "", status.Errorf(codes.PermissionDenied, "missing tenant header")
 	} else if len(tenants) > 1 {

cmd/collector/app/span_processor_test.go

@@ -10,26 +10,41 @@ import (
 	"errors"
 	"fmt"
 	"io"
+	"net/http"
 	"reflect"
 	"slices"
+	"strings"
 	"sync"
 	"sync/atomic"
 	"testing"
 	"time"
 
 	"github.com/gogo/protobuf/jsonpb"
 	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/mock"
 	"github.com/stretchr/testify/require"
+	"go.opentelemetry.io/collector/config/configgrpc"
+	"go.opentelemetry.io/collector/config/confighttp"
+	"go.opentelemetry.io/collector/config/confignet"
+	"go.opentelemetry.io/collector/pdata/ptrace"
+	otlptrace "go.opentelemetry.io/proto/otlp/collector/trace/v1"
+	tracepb "go.opentelemetry.io/proto/otlp/trace/v1"
 	"go.uber.org/zap"
+	"go.uber.org/zap/zaptest"
+	"google.golang.org/grpc"
+	"google.golang.org/grpc/credentials/insecure"
+	"google.golang.org/grpc/metadata"
 
 	"github.com/jaegertracing/jaeger-idl/model/v1"
+	cFlags "github.com/jaegertracing/jaeger/cmd/collector/app/flags"
 	"github.com/jaegertracing/jaeger/cmd/collector/app/handler"
 	"github.com/jaegertracing/jaeger/cmd/collector/app/processor"
 	zipkinsanitizer "github.com/jaegertracing/jaeger/cmd/collector/app/sanitizer/zipkin"
 	"github.com/jaegertracing/jaeger/internal/metricstest"
 	"github.com/jaegertracing/jaeger/pkg/metrics"
 	"github.com/jaegertracing/jaeger/pkg/tenancy"
 	"github.com/jaegertracing/jaeger/pkg/testutils"
+	"github.com/jaegertracing/jaeger/storage_v2/tracestore/mocks"
 	"github.com/jaegertracing/jaeger/storage_v2/v1adapter"
 	"github.com/jaegertracing/jaeger/thrift-gen/jaeger"
 	zc "github.com/jaegertracing/jaeger/thrift-gen/zipkincore"
@@ -807,3 +822,182 @@ func TestSpanProcessorWithOnDroppedSpanOption(t *testing.T) {
 	require.EqualError(t, err, processor.ErrBusy.Error())
 	assert.Equal(t, []string{"op3"}, droppedOperations)
 }
+
+func optionsWithPorts(portHttp string, portGrpc string) *cFlags.CollectorOptions {
+	opts := &cFlags.CollectorOptions{
+		OTLP: struct {
+			Enabled bool
+			GRPC    configgrpc.ServerConfig
+			HTTP    confighttp.ServerConfig
+		}{
+			Enabled: true,
+			HTTP: confighttp.ServerConfig{
+				Endpoint:        portHttp,
+				IncludeMetadata: true,
+			},
+			GRPC: configgrpc.ServerConfig{
+				NetAddr: confignet.AddrConfig{
+					Endpoint:  portGrpc,
+					Transport: confignet.TransportTypeTCP,
+				},
+			},
+		},
+	}
+	return opts
+}
+
+func TestOTLPReceiverWithV2Storage(t *testing.T) {
+	// Setup mock writer and expectations
+	mockWriter := mocks.NewWriter(t)
+
+	var receivedTraces atomic.Pointer[ptrace.Traces]
+	var receivedCtx atomic.Pointer[context.Context]
+	mockWriter.On("WriteTraces", mock.Anything, mock.Anything).
+		Run(func(args mock.Arguments) {
+			storeContext := args.Get(0).(context.Context)
+			storeTrace := args.Get(1).(ptrace.Traces)
+			receivedTraces.Store(&storeTrace)
+			receivedCtx.Store(&storeContext)
+		}).Return(nil)
+
+	spanProcessor, err := NewSpanProcessor(
+		mockWriter,
+		nil,
+		Options.NumWorkers(1),
+		Options.QueueSize(1),
+		Options.ReportBusy(true),
+	)
+	require.NoError(t, err)
+	defer spanProcessor.Close()
+	logger := zaptest.NewLogger(t)
+
+	portHttp := "4317"
+	portGrpc := "4318"
+
+	// Send trace via HTTP
+	t.Run("Send trace data using HTTP connection", func(t *testing.T) {
+		// Can't send tenancy headers with http request to OTLP receiver
+		tenancyMgr := &tenancy.Manager{
+			Enabled: true,
+			Header:  "x-tenant",
+		}
+
+		rec, err := handler.StartOTLPReceiver(
+			optionsWithPorts(fmt.Sprintf("localhost:%v", portHttp), fmt.Sprintf("localhost:%v", portGrpc)),
+			logger,
+			spanProcessor,
+			tenancyMgr,
+		)
+		require.NoError(t, err)
+		ctx := context.Background()
+		defer rec.Shutdown(ctx)
+
+		url := fmt.Sprintf("http://localhost:%v/v1/traces", portHttp)
+
+		traceJSON := `{
+			"resourceSpans": [{
+				"scopeSpans": [{
+					"spans": [{
+						"name": "test-trace"
+					}]
+				}]
+			}]
+		}`
+
+		httpClient := http.Client{}
+
+		req, err := http.NewRequestWithContext(ctx, http.MethodPost, url, strings.NewReader(traceJSON))
+		require.NoError(t, err)
+		req.Header.Set("Content-Type", "application/json")
+		req.Header.Set("x-tenant", "test-tenant")
+		resp, err := httpClient.Do(req)
+		require.NoError(t, err)
+		defer resp.Body.Close()
+		require.Equal(t, http.StatusOK, resp.StatusCode)
+
+		assert.Eventually(t, func() bool {
+			storedTraces := receivedTraces.Load()
+			storedCtx := receivedCtx.Load()
+			if storedTraces == nil || storedCtx == nil {
+				return false
+			}
+			receivedSpan := storedTraces.ResourceSpans().At(0).
+				ScopeSpans().At(0).
+				Spans().At(0)
+			receivedTenant := tenancy.GetTenant(*storedCtx)
+			return receivedSpan.Name() == "test-trace" && receivedTenant == "test-tenant"
+		}, 1*time.Second, 1*time.Millisecond)
+
+		mockWriter.AssertExpectations(t)
+	})
+
+	t.Run("Send trace data using GRPC connection", func(t *testing.T) {
+		// Setup proper tenancy manager using NewManager
+		tenancyMgr := tenancy.NewManager(&tenancy.Options{
+			Enabled: true,
+			Header:  "x-tenant",
+			Tenants: []string{"test-tenant"},
+		})
+
+		// Create and start receiver
+		rec, err := handler.StartOTLPReceiver(
+			optionsWithPorts(fmt.Sprintf("localhost:%v", portHttp), fmt.Sprintf("localhost:%v", portGrpc)),
+			logger,
+			spanProcessor,
+			tenancyMgr,
+		)
+		require.NoError(t, err)
+		ctx := context.Background()
+		defer rec.Shutdown(ctx)
+
+		// Create gRPC client connection
+		conn, err := grpc.NewClient(
+			fmt.Sprintf("localhost:%v", portGrpc),
+			grpc.WithTransportCredentials(insecure.NewCredentials()),
+		)
+		require.NoError(t, err)
+		defer conn.Close()
+
+		// Create metadata with tenant information
+		md := metadata.New(map[string]string{
+			tenancyMgr.Header: "test-tenant",
+		})
+		ctxWithMD := metadata.NewOutgoingContext(ctx, md)
+
+		client := otlptrace.NewTraceServiceClient(conn)
+		req := &otlptrace.ExportTraceServiceRequest{
+			ResourceSpans: []*tracepb.ResourceSpans{
+				{
+					ScopeSpans: []*tracepb.ScopeSpans{
+						{
+							Spans: []*tracepb.Span{
+								{
+									Name: "test-trace",
+								},
+							},
+						},
+					},
+				},
+			},
+		}
+
+		// Send traces via gRPC
+		_, err = client.Export(ctxWithMD, req)
+		require.NoError(t, err)
+
+		assert.Eventually(t, func() bool {
+			storedTraces := receivedTraces.Load()
+			storedCtx := receivedCtx.Load()
+			if storedTraces == nil || storedCtx == nil {
+				return false
+			}
+			receivedSpan := storedTraces.ResourceSpans().At(0).
+				ScopeSpans().At(0).
+				Spans().At(0)
+			receivedTenant := tenancy.GetTenant(*storedCtx)
+			return receivedSpan.Name() == "test-trace" && receivedTenant == "test-tenant"
+		}, 1*time.Second, 1*time.Millisecond)
+
+		mockWriter.AssertExpectations(t)
+	})
+}