common/metricsserver/metrics_test.go

package metricsserver

import (
	"context"
	"fmt"
	"io"
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"
	"time"

	"github.com/prometheus/client_golang/prometheus"
)

func TestNew(t *testing.T) {
	metrics := New()

	if metrics == nil {
		t.Fatal("New returned nil")
	}

	if metrics.r == nil {
		t.Error("metrics registry is nil")
	}
}

func TestRegistry(t *testing.T) {
	metrics := New()
	registry := metrics.Registry()

	if registry == nil {
		t.Fatal("Registry() returned nil")
	}

	if registry != metrics.r {
		t.Error("Registry() did not return the metrics registry")
	}

	// Test that we can register a metric
	counter := prometheus.NewCounter(prometheus.CounterOpts{
		Name: "test_counter",
		Help: "A test counter",
	})

	err := registry.Register(counter)
	if err != nil {
		t.Errorf("failed to register metric: %v", err)
	}

	// Test that the metric is registered
	metricFamilies, err := registry.Gather()
	if err != nil {
		t.Errorf("failed to gather metrics: %v", err)
	}

	found := false
	for _, mf := range metricFamilies {
		if mf.GetName() == "test_counter" {
			found = true
			break
		}
	}

	if !found {
		t.Error("registered metric not found in registry")
	}
}

func TestHandler(t *testing.T) {
	metrics := New()

	// Register a test metric
	counter := prometheus.NewCounterVec(
		prometheus.CounterOpts{
			Name: "test_requests_total",
			Help: "Total number of test requests",
		},
		[]string{"method"},
	)
	metrics.Registry().MustRegister(counter)
	counter.WithLabelValues("GET").Inc()

	// Test the handler
	handler := metrics.Handler()
	if handler == nil {
		t.Fatal("Handler() returned nil")
	}

	// Create a test request
	req := httptest.NewRequest("GET", "/metrics", nil)
	recorder := httptest.NewRecorder()

	// Call the handler
	handler.ServeHTTP(recorder, req)

	// Check response
	resp := recorder.Result()
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusOK {
		t.Errorf("expected status 200, got %d", resp.StatusCode)
	}

	body, err := io.ReadAll(resp.Body)
	if err != nil {
		t.Fatalf("failed to read response body: %v", err)
	}

	bodyStr := string(body)

	// Check for our test metric
	if !strings.Contains(bodyStr, "test_requests_total") {
		t.Error("test metric not found in metrics output")
	}

	// Check for OpenMetrics format indicators
	if !strings.Contains(bodyStr, "# TYPE") {
		t.Error("metrics output missing TYPE comments")
	}
}

func TestListenAndServe(t *testing.T) {
	metrics := New()

	// Register a test metric
	counter := prometheus.NewCounterVec(
		prometheus.CounterOpts{
			Name: "test_requests_total",
			Help: "Total number of test requests",
		},
		[]string{"method"},
	)
	metrics.Registry().MustRegister(counter)
	counter.WithLabelValues("GET").Inc()

	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	// Start server in a goroutine
	errCh := make(chan error, 1)
	go func() {
		// Use a high port number to avoid conflicts
		errCh <- metrics.ListenAndServe(ctx, 9999)
	}()

	// Give the server a moment to start
	time.Sleep(100 * time.Millisecond)

	// Test metrics endpoint
	resp, err := http.Get("http://localhost:9999/metrics")
	if err != nil {
		t.Fatalf("failed to GET /metrics: %v", err)
	}
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusOK {
		t.Errorf("expected status 200, got %d", resp.StatusCode)
	}

	body, err := io.ReadAll(resp.Body)
	if err != nil {
		t.Fatalf("failed to read response body: %v", err)
	}

	bodyStr := string(body)

	// Check for our test metric
	if !strings.Contains(bodyStr, "test_requests_total") {
		t.Error("test metric not found in metrics output")
	}

	// Cancel context to stop server
	cancel()

	// Wait for server to stop
	select {
	case err := <-errCh:
		if err != nil {
			t.Errorf("server returned error: %v", err)
		}
	case <-time.After(5 * time.Second):
		t.Error("server did not stop within timeout")
	}
}

func TestListenAndServeContextCancellation(t *testing.T) {
	metrics := New()

	ctx, cancel := context.WithCancel(context.Background())

	// Start server
	errCh := make(chan error, 1)
	go func() {
		errCh <- metrics.ListenAndServe(ctx, 9998)
	}()

	// Give server time to start
	time.Sleep(100 * time.Millisecond)

	// Cancel context
	cancel()

	// Server should stop gracefully
	select {
	case err := <-errCh:
		if err != nil {
			t.Errorf("server returned error on graceful shutdown: %v", err)
		}
	case <-time.After(5 * time.Second):
		t.Error("server did not stop within timeout after context cancellation")
	}
}

// Benchmark the metrics handler response time
func BenchmarkMetricsHandler(b *testing.B) {
	metrics := New()

	// Register some test metrics
	for i := 0; i < 10; i++ {
		counter := prometheus.NewCounter(prometheus.CounterOpts{
			Name: fmt.Sprintf("bench_counter_%d", i),
			Help: "A benchmark counter",
		})
		metrics.Registry().MustRegister(counter)
		counter.Add(float64(i * 100))
	}

	handler := metrics.Handler()

	b.ResetTimer()

	for i := 0; i < b.N; i++ {
		req := httptest.NewRequest("GET", "/metrics", nil)
		recorder := httptest.NewRecorder()
		handler.ServeHTTP(recorder, req)

		if recorder.Code != http.StatusOK {
			b.Fatalf("unexpected status code: %d", recorder.Code)
		}
	}
}