Add services test cases

Signed-off-by: Hiroshi Hatake <[email protected]>

Author: cosmo0920

custom_jaeger_remote.go

@@ -34,6 +34,9 @@ type jaegerRemotePlugin struct {
 	config       *Config
 	server       *serverComponent // Server-related state
 	clientTracer *clientComponent // Client-related state
+
+	// newSamplerFn allows injecting a mock sampler factory for testing.
+	newSamplerFn func(context.Context, *Config) (*remoteSampler, error)
 }
 
 type serverComponent struct {
@@ -103,20 +106,21 @@ func (plug *jaegerRemotePlugin) Init(ctx context.Context, fbit *plugin.Fluentbit
 	}
 	plug.config = cfg
 
-	// Run client mode
+	// Default to the real sampler factory if none is injected for tests.
+	if plug.newSamplerFn == nil {
+		plug.newSamplerFn = newRemoteSampler
+	}
+
 	if cfg.Mode == "client" || cfg.Mode == "all" {
 		if err := plug.initClient(ctx); err != nil {
 			return fmt.Errorf("failed to initialize client mode: %w", err)
 		}
 	}
-
-	// Run server mode
 	if cfg.Mode == "server" || cfg.Mode == "all" {
 		if err := plug.initServer(ctx); err != nil {
 			return fmt.Errorf("failed to initialize server mode: %w", err)
 		}
 	}
-
 	plug.log.Info("plugin initialized successfully in mode: '%s'", cfg.Mode)
 	return nil
 }

jaeger_services.go

@@ -66,23 +66,17 @@ func (plug *jaegerRemotePlugin) initClient(ctx context.Context) error {
 	return nil
 }
 
-//
 // --- Server Mode Initialization ---
-//
-
 func (plug *jaegerRemotePlugin) initServer(ctx context.Context) error {
 	plug.log.Info("initializing server mode...")
-
 	plug.server = &serverComponent{}
-	sampler, err := newRemoteSampler(ctx, plug.config)
+	sampler, err := plug.newSamplerFn(ctx, plug.config)
 	if err != nil {
 		return fmt.Errorf("could not create remote sampler for server: %w", err)
 	}
 	plug.server.sampler = sampler
 	plug.server.cache = &samplingStrategyCache{strategies: make(map[string]*api_v2.SamplingStrategyResponse)}
 
-	go plug.pollStrategiesWithRetry(ctx)
-
 	if plug.config.ServerHttpListenAddr != "" {
 		plug.server.httpServer = plug.startHttpServer()
 	}
@@ -93,7 +87,7 @@ func (plug *jaegerRemotePlugin) initServer(ctx context.Context) error {
 			if plug.server.httpServer != nil {
 				shutdownCtx, cancel := context.WithTimeout(context.Background(), 1*time.Second)
 				defer cancel()
-				plug.server.httpServer.Shutdown(shutdownCtx)
+				_ = plug.server.httpServer.Shutdown(shutdownCtx)
 			}
 			return err
 		}
@@ -103,6 +97,8 @@ func (plug *jaegerRemotePlugin) initServer(ctx context.Context) error {
 		return errors.New("server mode is enabled, but neither 'server.http.listen_addr' nor 'server.grpc.listen_addr' are configured")
 	}
 
+	// All checks passed, now start background goroutines.
+	go plug.pollStrategiesWithRetry(ctx)
 	go func() {
 		<-ctx.Done()
 		plug.log.Info("shutting down server components...")
@@ -111,7 +107,7 @@ func (plug *jaegerRemotePlugin) initServer(ctx context.Context) error {
 			plug.log.Info("gRPC server stopped.")
 		}
 		if sampler.conn != nil {
-			sampler.conn.Close()
+			_ = sampler.conn.Close()
 			plug.log.Info("gRPC client connection to Jaeger Collector closed.")
 		}
 		if plug.server.httpServer != nil {

jaeger_services_test.go

@@ -0,0 +1,227 @@
+package main
+
+import (
+	"context"
+	"encoding/json"
+	"fmt"
+	"net"
+	"net/http"
+	"net/http/httptest"
+	"strings"
+	"sync"
+	"testing"
+	"time"
+
+	"github.com/alecthomas/assert/v2"
+	"github.com/calyptia/plugin"
+	api_v2 "github.com/jaegertracing/jaeger-idl/proto-gen/api_v2"
+	"google.golang.org/grpc"
+	"google.golang.org/grpc/credentials/insecure"
+	"google.golang.org/grpc/metadata"
+	"google.golang.org/grpc/test/bufconn"
+)
+
+type observedCall struct {
+	Context context.Context
+	Params  *api_v2.SamplingStrategyParameters
+}
+
+type samplingServer struct {
+	api_v2.UnimplementedSamplingManagerServer
+	mu            sync.Mutex
+	observedCalls []observedCall
+	strategy      *api_v2.SamplingStrategyResponse
+}
+
+func (s *samplingServer) GetSamplingStrategy(ctx context.Context, params *api_v2.SamplingStrategyParameters) (*api_v2.SamplingStrategyResponse, error) {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+	s.observedCalls = append(s.observedCalls, observedCall{Context: ctx, Params: params})
+	return s.strategy, nil
+}
+
+func (s *samplingServer) callCount() int {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+	return len(s.observedCalls)
+}
+func (s *samplingServer) lastCall() *observedCall {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+	if len(s.observedCalls) == 0 {
+		return nil
+	}
+	return &s.observedCalls[len(s.observedCalls)-1]
+}
+
+func startMockGrpcServer(t *testing.T, mock *samplingServer) (*grpc.Server, *bufconn.Listener) {
+	t.Helper()
+	lis := bufconn.Listen(1024 * 1024)
+	s := grpc.NewServer()
+	api_v2.RegisterSamplingManagerServer(s, mock)
+	go func() {
+		if err := s.Serve(lis); err != nil {
+			t.Logf("Mock gRPC server stopped: %v", err)
+		}
+	}()
+	return s, lis
+}
+
+func startMockHTTPSamplingServer(t *testing.T, strategy *api_v2.SamplingStrategyResponse) *httptest.Server {
+	t.Helper()
+	return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		w.Header().Set("Content-Type", "application/json")
+		err := json.NewEncoder(w).Encode(strategy)
+		assert.NoError(t, err)
+	}))
+}
+
+func Test_InitClient(t *testing.T) {
+	t.Run("successfully initializes in client mode", func(t *testing.T) {
+		ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+		defer cancel()
+
+		mockStrategy := &api_v2.SamplingStrategyResponse{
+			StrategyType:          api_v2.SamplingStrategyType_PROBABILISTIC,
+			ProbabilisticSampling: &api_v2.ProbabilisticSamplingStrategy{SamplingRate: 0.1},
+		}
+		mockSamplingSrv := startMockHTTPSamplingServer(t, mockStrategy)
+		defer mockSamplingSrv.Close()
+
+		mockOtlpSrv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			w.WriteHeader(http.StatusOK)
+		}))
+		defer mockOtlpSrv.Close()
+
+		fbit := &plugin.Fluentbit{
+			Logger: newTestLogger(t),
+			Conf: mapConfigLoader{
+				"mode":                "client",
+				"client.server_url":   mockOtlpSrv.URL,
+				"client.sampling_url": mockSamplingSrv.URL,
+			},
+		}
+		plug := &jaegerRemotePlugin{}
+
+		err := plug.Init(ctx, fbit)
+
+		assert.NoError(t, err)
+		assert.NotZero(t, plug.clientTracer, "clientTracer should be initialized")
+		assert.NotZero(t, plug.clientTracer.tracerProvider, "tracerProvider should be initialized")
+	})
+}
+
+func Test_InitServer_EndToEnd(t *testing.T) {
+	testCases := []struct {
+		name            string
+		configHeaders   string
+		expectedHeaders map[string]string
+	}{
+		{"no headers", "", nil},
+		{
+			"with headers",
+			"x-custom-header=fluent-bit,authorization=Bearer 12345",
+			map[string]string{"x-custom-header": "fluent-bit", "authorization": "Bearer 12345"},
+		},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.name, func(t *testing.T) {
+			ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+			defer cancel()
+
+			mockStrategy := &api_v2.SamplingStrategyResponse{StrategyType: api_v2.SamplingStrategyType_RATE_LIMITING, RateLimitingSampling: &api_v2.RateLimitingSamplingStrategy{MaxTracesPerSecond: 100}}
+			mockJaeger := &samplingServer{strategy: mockStrategy}
+			upstreamJaegerServer, lis := startMockGrpcServer(t, mockJaeger)
+			defer upstreamJaegerServer.Stop()
+
+			fbit := &plugin.Fluentbit{
+				Logger: newTestLogger(t),
+				Conf: mapConfigLoader{
+					"mode":                          "server",
+					"server.endpoint":               "bufnet",
+					"server.http.listen_addr":       getFreePort(t),
+					"server.service_names":          "test-service",
+					"server.retry.initial_interval": "10ms",
+					"server.headers":                tc.configHeaders,
+				},
+			}
+			plug := &jaegerRemotePlugin{}
+
+			plug.newSamplerFn = func(ctx context.Context, cfg *Config) (*remoteSampler, error) {
+				dialOpts := []grpc.DialOption{
+					grpc.WithContextDialer(func(context.Context, string) (net.Conn, error) { return lis.Dial() }),
+					grpc.WithTransportCredentials(insecure.NewCredentials()),
+				}
+				if len(cfg.ServerHeaders) > 0 {
+					dialOpts = append(dialOpts, grpc.WithUnaryInterceptor(func(ctx context.Context, method string, req, reply interface{}, cc *grpc.ClientConn, invoker grpc.UnaryInvoker, opts ...grpc.CallOption) error {
+						return invoker(metadata.NewOutgoingContext(ctx, metadata.New(cfg.ServerHeaders)), method, req, reply, cc, opts...)
+					}))
+				}
+				conn, err := grpc.DialContext(ctx, cfg.ServerEndpoint, dialOpts...)
+				if err != nil {
+					return nil, err
+				}
+				client := api_v2.NewSamplingManagerClient(conn)
+				return &remoteSampler{conn: conn, client: client}, nil
+			}
+
+			err := plug.Init(ctx, fbit)
+			assert.NoError(t, err)
+			assert.NotZero(t, plug.server)
+			if plug.server.httpServer != nil {
+				defer plug.server.httpServer.Close()
+			}
+
+			lastCall := mockJaeger.lastCall()
+			assert.NotZero(t, lastCall)
+			assert.Equal(t, "test-service", lastCall.Params.ServiceName)
+
+			md, ok := metadata.FromIncomingContext(lastCall.Context)
+			assert.True(t, ok)
+			for k, v := range tc.expectedHeaders {
+				assert.Equal(t, []string{v}, md.Get(k))
+			}
+		})
+	}
+}
+
+func Test_InitServer_Failure(t *testing.T) {
+	t.Run("fails if both http and grpc listen addresses are missing", func(t *testing.T) {
+		ctx, cancel := context.WithTimeout(context.Background(), 1*time.Second)
+		defer cancel()
+
+		fbit := &plugin.Fluentbit{
+			Logger: newTestLogger(t),
+			Conf: mapConfigLoader{
+				"mode":                 "server",
+				"server.endpoint":      "dummy:1234",
+				"server.service_names": "test-service",
+			},
+		}
+		plug := &jaegerRemotePlugin{}
+
+		err := plug.Init(ctx, fbit)
+
+		assert.Error(t, err)
+		assert.Contains(t, err.Error(), "neither 'server.http.listen_addr' nor 'server.grpc.listen_addr' are configured")
+	})
+}
+
+func getFreePort(t *testing.T) string {
+	t.Helper()
+	addr, err := net.ResolveTCPAddr("tcp", "localhost:0")
+	assert.NoError(t, err, "failed to resolve free port")
+	l, err := net.ListenTCP("tcp", addr)
+	assert.NoError(t, err, "failed to listen on free port")
+	defer l.Close()
+	return l.Addr().String()
+}
+
+func mapToString(m map[string]string) string {
+	var parts []string
+	for k, v := range m {
+		parts = append(parts, fmt.Sprintf("%s=%s", k, v))
+	}
+	return strings.Join(parts, ",")
+}