processor_test.go

// Copyright 2024 tailsampling-go contributors
// Licensed under the Apache License, Version 2.0

package tailsampling_test

import (
	"context"
	"strings"
	"sync"
	"testing"
	"time"

	tailsampling "github.com/lukashes/tailsampling-go"
	"github.com/lukashes/tailsampling-go/policies"
	"go.opentelemetry.io/otel"
	"go.opentelemetry.io/otel/codes"
	sdktrace "go.opentelemetry.io/otel/sdk/trace"
)

// testExporter collects exported spans for testing
type testExporter struct {
	mu    sync.Mutex
	spans []sdktrace.ReadOnlySpan
}

func (t *testExporter) ExportSpans(ctx context.Context, spans []sdktrace.ReadOnlySpan) error {
	t.mu.Lock()
	defer t.mu.Unlock()
	t.spans = append(t.spans, spans...)
	return nil
}

func (t *testExporter) Shutdown(ctx context.Context) error {
	return nil
}

func (t *testExporter) GetSpans() []sdktrace.ReadOnlySpan {
	t.mu.Lock()
	defer t.mu.Unlock()
	return t.spans
}

func TestLatencyPolicy(t *testing.T) {
	exporter := &testExporter{}
	latencyPolicy := policies.NewLatency("latency", 100, 0)

	processor, err := tailsampling.NewTailSamplingProcessor(
		tailsampling.WithDecisionWait(100*time.Millisecond),
		tailsampling.WithNumTraces(100),
		tailsampling.WithPolicies(latencyPolicy),
		tailsampling.WithExporter(exporter),
	)
	if err != nil {
		t.Fatalf("failed to create processor: %v", err)
	}

	tp := sdktrace.NewTracerProvider(
		sdktrace.WithSampler(sdktrace.AlwaysSample()),
		sdktrace.WithSpanProcessor(processor),
	)
	otel.SetTracerProvider(tp)

	tracer := tp.Tracer("test")

	// Fast trace - should NOT be sampled
	_, span1 := tracer.Start(context.Background(), "fast")
	time.Sleep(50 * time.Millisecond)
	span1.End()

	// Slow trace - SHOULD be sampled
	_, span2 := tracer.Start(context.Background(), "slow")
	time.Sleep(150 * time.Millisecond)
	span2.End()

	// Wait for decision (DecisionWait = 100ms, ticker = 50ms)
	time.Sleep(500 * time.Millisecond)

	if err := tp.Shutdown(context.Background()); err != nil {
		t.Fatalf("shutdown failed: %v", err)
	}

	// Check exported spans
	spans := exporter.GetSpans()
	if len(spans) != 1 {
		t.Fatalf("expected 1 exported span, got %d", len(spans))
	}
	if spans[0].Name() != "slow" {
		t.Errorf("expected 'slow' span, got '%s'", spans[0].Name())
	}
}

func TestStatusCodePolicy(t *testing.T) {
	exporter := &testExporter{}
	statusPolicy, _ := policies.NewStatusCode("errors", []string{"ERROR"})

	processor, err := tailsampling.NewTailSamplingProcessor(
		tailsampling.WithDecisionWait(100*time.Millisecond),
		tailsampling.WithNumTraces(100),
		tailsampling.WithPolicies(statusPolicy),
		tailsampling.WithExporter(exporter),
	)
	if err != nil {
		t.Fatalf("failed to create processor: %v", err)
	}

	tp := sdktrace.NewTracerProvider(
		sdktrace.WithSampler(sdktrace.AlwaysSample()),
		sdktrace.WithSpanProcessor(processor),
	)
	otel.SetTracerProvider(tp)

	tracer := tp.Tracer("test")

	// OK trace - should NOT be sampled
	_, span1 := tracer.Start(context.Background(), "ok-span")
	span1.End()

	// Error trace - SHOULD be sampled
	_, span2 := tracer.Start(context.Background(), "error-span")
	span2.SetStatus(codes.Error, "test error")
	span2.End()

	// Wait for decision (DecisionWait = 100ms, ticker = 50ms)
	time.Sleep(500 * time.Millisecond)

	if err := tp.Shutdown(context.Background()); err != nil {
		t.Fatalf("shutdown failed: %v", err)
	}

	// Check that error trace was sampled and OK trace was not
	metrics := processor.GetMetrics()
	if metrics.PolicySampled != 1 {
		t.Errorf("expected 1 sampled trace, got %d", metrics.PolicySampled)
	}
	if metrics.PolicyNotSampled != 1 {
		t.Errorf("expected 1 not-sampled trace, got %d", metrics.PolicyNotSampled)
	}
}


func TestDropPendingTracesOnShutdown(t *testing.T) {
	exporter := &testExporter{}
	statusPolicy, _ := policies.NewStatusCode("all", []string{"ERROR", "OK", "UNSET"})

	// DropPendingTracesOnShutdown = false
	processor, err := tailsampling.NewTailSamplingProcessor(
		tailsampling.WithDecisionWait(10*time.Second), // Long wait
		tailsampling.WithNumTraces(100),
		tailsampling.WithPolicies(statusPolicy),
		tailsampling.WithExporter(exporter),
		tailsampling.WithDropPendingTracesOnShutdown(false),
	)
	if err != nil {
		t.Fatalf("failed to create processor: %v", err)
	}

	tp := sdktrace.NewTracerProvider(
		sdktrace.WithSampler(sdktrace.AlwaysSample()),
		sdktrace.WithSpanProcessor(processor),
	)
	otel.SetTracerProvider(tp)

	tracer := tp.Tracer("test")

	// Create span
	_, span := tracer.Start(context.Background(), "pending-span")
	span.End()

	// Shutdown immediately (before DecisionWait)
	if err := tp.Shutdown(context.Background()); err != nil {
		t.Fatalf("shutdown failed: %v", err)
	}

	// Span should be exported despite not waiting
	spans := exporter.GetSpans()
	if len(spans) != 1 {
		t.Errorf("expected 1 span (processed on shutdown), got %d", len(spans))
	}
}

func TestConfigValidation(t *testing.T) {
	exporter := &testExporter{}
	policy := policies.NewLatency("latency", 100, 0)

	tests := []struct {
		name string
		opts []tailsampling.Option
		want string
	}{
		{
			name: "no exporter",
			opts: []tailsampling.Option{
				tailsampling.WithPolicies(policy),
			},
			want: "no exporter configured",
		},
		{
			name: "no policies",
			opts: []tailsampling.Option{
				tailsampling.WithExporter(exporter),
			},
			want: "no policies configured",
		},
		{
			name: "zero DecisionWait",
			opts: []tailsampling.Option{
				tailsampling.WithExporter(exporter),
				tailsampling.WithPolicies(policy),
				tailsampling.WithDecisionWait(0),
			},
			want: "DecisionWait must be positive",
		},
		{
			name: "DecisionWait too large",
			opts: []tailsampling.Option{
				tailsampling.WithExporter(exporter),
				tailsampling.WithPolicies(policy),
				tailsampling.WithDecisionWait(48 * time.Hour),
			},
			want: "DecisionWait too large",
		},
		{
			name: "zero NumTraces",
			opts: []tailsampling.Option{
				tailsampling.WithExporter(exporter),
				tailsampling.WithPolicies(policy),
				tailsampling.WithNumTraces(0),
			},
			want: "NumTraces must be positive",
		},
		{
			name: "NumTraces too large",
			opts: []tailsampling.Option{
				tailsampling.WithExporter(exporter),
				tailsampling.WithPolicies(policy),
				tailsampling.WithNumTraces(20_000_000),
			},
			want: "NumTraces too large",
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			_, err := tailsampling.NewTailSamplingProcessor(tt.opts...)
			if err == nil {
				t.Fatal("expected error, got nil")
			}
			if !strings.Contains(err.Error(), tt.want) {
				t.Errorf("expected error containing %q, got %q", tt.want, err.Error())
			}
		})
	}
}