query-orchestration/internal/serviceconfig/observability/prometheus/prometheus_test.go

package prometheus

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

	"github.com/stretchr/testify/require"
)

func GetFreePortForTesting() (int, error) {
	// Create a TCP listener on port 0 (lets OS choose a free port)
	addr, err := net.ResolveTCPAddr("tcp", "localhost:0")
	if err != nil {
		return 0, fmt.Errorf("failed to resolve TCP address: %v", err)
	}

	l, err := net.ListenTCP("tcp", addr)
	if err != nil {
		return 0, fmt.Errorf("failed to create TCP listener: %v", err)
	}
	defer l.Close()

	// Get the port that was assigned by the OS
	addr, ok := l.Addr().(*net.TCPAddr)
	if !ok {
		return 0, fmt.Errorf("failed to parse TCP address")
	}

	return addr.Port, nil
}

func TestMetricsIntegration(t *testing.T) {
	// Create a context with cancellation for the test
	ctx, cancel := context.WithCancel(t.Context())
	defer cancel()

	testingPort, errorGettingPort := GetFreePortForTesting()
	require.NoError(t, errorGettingPort)

	// Step 1: Define our test metrics configuration
	config := MetricsConfig{
		Namespace: "test",
		Subsystem: "metrics",
		Port:      testingPort, // Using 9090 for tests
		ServeHttp: true,        // Enable the built-in HTTP server
		Definitions: []MetricDefinition{
			{
				// Counter example: counting total events
				Type:   Counter,
				Name:   "events_total",
				Help:   "Total number of events processed",
				Labels: []string{"type", "status"},
			},
			{
				// Gauge example: tracking current value
				Type:   Gauge,
				Name:   "queue_size",
				Help:   "Current number of items in queue",
				Labels: []string{"queue_name"},
			},
			{
				// Histogram example: measuring distribution of values
				Type:    Histogram,
				Name:    "response_time_seconds",
				Help:    "Response time distribution",
				Labels:  []string{"endpoint"},
				Buckets: []float64{0.1, 0.5, 1, 2, 5}, // Define bucket boundaries
			},
			{
				// Summary example: calculating quantiles
				Type:   Summary,
				Name:   "request_size_bytes",
				Help:   "Request size distribution",
				Labels: []string{"method"},
				Objectives: map[float64]float64{
					0.5:  0.05,  // 50th percentile with 5% error
					0.9:  0.01,  // 90th percentile with 1% error
					0.99: 0.001, // 99th percentile with 0.1% error
				},
			},
		},
	}

	// Step 2: Create new metrics instance with context
	metrics, err := New(ctx, config)
	require.NoError(t, err)

	// Ensure cleanup after test
	defer func() {
		shutdownCtx, shutdownCancel := context.WithTimeout(t.Context(), 5*time.Second)
		defer shutdownCancel()

		if err := metrics.Shutdown(shutdownCtx); err != nil {
			t.Errorf("Failed to shutdown metrics server: %v", err)
		}
	}()

	// Step 3: Record some sample metrics
	// Record a counter metric
	err = metrics.RecordMetric("events_total", 42, map[string]string{
		"type":   "user_login",
		"status": "success",
	})
	require.NoError(t, err)

	// Record a gauge metric
	err = metrics.RecordMetric("queue_size", 15, map[string]string{
		"queue_name": "default",
	})
	require.NoError(t, err)

	// Record a histogram metric
	err = metrics.RecordMetric("response_time_seconds", 0.7, map[string]string{
		"endpoint": "/api/users",
	})
	require.NoError(t, err)

	// Record a summary metric
	err = metrics.RecordMetric("request_size_bytes", 1024, map[string]string{
		"method": "POST",
	})
	require.NoError(t, err)

	// Step 4: Wait for HTTP server to start
	// Give the server a moment to start up
	time.Sleep(time.Second)

	// Step 5: Fetch metrics from the HTTP endpoint
	URLForMetrics := fmt.Sprintf("http://localhost:%d/metrics", testingPort)
	resp, err := http.Get(URLForMetrics)
	require.NoError(t, err)
	defer resp.Body.Close()

	// Read the response body
	body, err := io.ReadAll(resp.Body)
	require.NoError(t, err)

	metricsOutput := string(body)
	// if needed for debugging.
	//log.Infof("Metrics output: %s", metricsOutput)

	// Step 6: Verify that our custom metrics and go metrics are present
	expectedMetrics := []string{
		// Our custom metrics (with namespace and subsystem)
		"test_metrics_events_total",
		"test_metrics_queue_size",
		"test_metrics_response_time_seconds",
		"test_metrics_request_size_bytes",
		// Standard Go metrics that should be present
		"go_memstats_alloc_bytes_total",
	}

	// Check for each expected metric
	for _, metric := range expectedMetrics {
		if !strings.Contains(metricsOutput, metric) {
			t.Errorf("Expected metric %s not found in output", metric)
			// Print a snippet of the output around where we expect to find the metric
			fmt.Printf("Searching for: %s\nOutput snippet: %s\n",
				metric,
				getOutputSnippet(metricsOutput, metric))
		}
	}

	// Test cancellation
	t.Run("Context cancellation", func(t *testing.T) {
		// Create a new metrics instance with a separate context
		testingPort, errorGettingPort := GetFreePortForTesting()
		require.NoError(t, errorGettingPort)

		cancelCtx, cancelFunc := context.WithCancel(t.Context())
		_, err := New(cancelCtx, MetricsConfig{
			Namespace: "test_cancel",
			Port:      testingPort,
			ServeHttp: true,
		})
		require.NoError(t, err)

		// Wait for server to start
		time.Sleep(time.Second)

		// Cancel the context
		cancelFunc()

		// Try to connect to the server - it should fail after a short delay
		time.Sleep(time.Second)
		testingUrl := fmt.Sprintf("http://localhost:%d/metrics", testingPort)
		_, err = http.Get(testingUrl)
		if err == nil {
			t.Error("Expected server to be shutdown after context cancellation")
		}
	})
}

// Helper function to get a snippet of the output around where we expect to find the metric
func getOutputSnippet(output string, searchTerm string) string {
	index := strings.Index(output, searchTerm)
	if index == -1 {
		return "Metric not found in output"
	}

	// Get up to 200 characters before and after the search term
	start := max(0, index-200)
	end := min(len(output), index+200)
	return output[start:end]
}

func max(a, b int) int {
	if a > b {
		return a
	}
	return b
}

func min(a, b int) int {
	if a < b {
		return a
	}
	return b
}