fetch_ml/tests/benchmarks/api_benchmark_test.go

package benchmarks

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

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

func newBenchmarkHTTPClient() *http.Client {
	transport := &http.Transport{
		Proxy: http.ProxyFromEnvironment,
		DialContext: (&net.Dialer{
			Timeout:   30 * time.Second,
			KeepAlive: 30 * time.Second,
		}).DialContext,
		ForceAttemptHTTP2:     true,
		MaxIdleConns:          256,
		MaxIdleConnsPerHost:   256,
		IdleConnTimeout:       90 * time.Second,
		TLSHandshakeTimeout:   10 * time.Second,
		ExpectContinueTimeout: 1 * time.Second,
	}

	return &http.Client{
		Timeout:   30 * time.Second,
		Transport: transport,
	}
}

// BenchmarkAPIServerCreateJob tests job creation performance
func BenchmarkAPIServerCreateJob(b *testing.B) {
	server := setupTestAPIServer(b)
	defer server.Close()
	client := newBenchmarkHTTPClient()
	b.ResetTimer()
	b.ReportAllocs()
	benchmarkCreateJob(b, server.URL, client)
}

// BenchmarkAPIServerCreateJobSimple tests job creation with simplified setup
func BenchmarkAPIServerCreateJobSimple(b *testing.B) {
	// Create a simple HTTP server without httptest
	mux := http.NewServeMux()

	mux.HandleFunc("/api/v1/jobs", func(w http.ResponseWriter, _ *http.Request) {
		w.Header().Set("Content-Type", "application/json")
		w.WriteHeader(http.StatusCreated)
		_ = json.NewEncoder(w).Encode(map[string]string{"id": "test-job-id"})
	})

	server := &http.Server{
		Addr:              "127.0.0.1:0", // Use random available port
		Handler:           mux,
		ReadHeaderTimeout: 5 * time.Second,
	}

	// Start server in goroutine
	go func() { _ = server.ListenAndServe() }()

	// Get the actual port
	addr := server.Addr
	if addr == "" {
		addr = "127.0.0.1:8080"
	}

	client := &http.Client{Timeout: 30 * time.Second}
	baseURL := "http://" + addr

	b.ResetTimer()
	b.ReportAllocs()

	for i := 0; i < b.N; i++ {
		jobData := map[string]interface{}{
			"job_name": fmt.Sprintf("benchmark-job-%d", i),
			"args": map[string]interface{}{
				"model": "test-model",
				"data":  generateTestPayload(1024),
			},
			"priority": 0,
		}

		jsonData, _ := json.Marshal(jobData)
		req, err := http.NewRequestWithContext(context.Background(), "POST",
			baseURL+"/api/v1/jobs", bytes.NewBuffer(jsonData))
		if err != nil {
			b.Fatalf("Failed to create request: %v", err)
		}
		req.Header.Set("Content-Type", "application/json")

		resp, err := client.Do(req)
		if err != nil {
			// Skip this iteration if server isn't ready
			continue
		}
		_, _ = io.Copy(io.Discard, resp.Body)
		_ = resp.Body.Close()
	}

	_ = server.Close()
}

func generateTestPayload(size int) string {
	data := make([]byte, size)
	for i := range data {
		data[i] = byte(i % 256)
	}
	return string(data)
}

// BenchmarkMetrics measures the performance impact of metrics collection
func BenchmarkMetricsCollection(b *testing.B) {
	registry := prometheus.NewRegistry()

	// Create test metrics
	counter := prometheus.NewCounter(prometheus.CounterOpts{
		Name: "test_operations_total",
		Help: "Total number of test operations",
	})

	histogram := prometheus.NewHistogram(prometheus.HistogramOpts{
		Name:    "test_duration_seconds",
		Help:    "Test operation duration",
		Buckets: prometheus.DefBuckets,
	})

	registry.MustRegister(counter, histogram)

	b.ResetTimer()
	b.ReportAllocs()

	for i := 0; i < b.N; i++ {
		counter.Inc()
		histogram.Observe(float64(i) * 0.001)
	}
}

// BenchmarkConcurrentRequests tests concurrent API performance
func BenchmarkConcurrentRequests(b *testing.B) {
	server := setupTestAPIServer(b)
	defer server.Close()

	client := newBenchmarkHTTPClient()

	b.ResetTimer()

	// Test different concurrency levels
	for _, concurrency := range []int{1, 5, 10, 25, 50} {
		b.Run(fmt.Sprintf("Concurrency-%d", concurrency), func(b *testing.B) {
			benchmarkConcurrentRequests(b, server.URL, client, concurrency)
		})
	}
}

func benchmarkConcurrentRequests(b *testing.B, baseURL string, client *http.Client, concurrency int) {
	b.SetParallelism(concurrency)
	b.RunParallel(func(pb *testing.PB) {
		i := 0
		for pb.Next() {
			req, _ := http.NewRequestWithContext(context.Background(), "GET", baseURL+"/api/v1/jobs?limit=10", nil)
			req.Header.Set("Authorization", "Bearer test-token")

			resp, err := client.Do(req)
			if err == nil && resp != nil {
				_, _ = io.Copy(io.Discard, resp.Body)
				_ = resp.Body.Close()
			}
			i++
		}
	})
}

// setupTestAPIServer creates a test HTTP server for benchmarking
func setupTestAPIServer(_ *testing.B) *httptest.Server {
	mux := http.NewServeMux()

	// Add basic API routes for benchmarking
	mux.HandleFunc("/api/v1/jobs", func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Type", "application/json")
		if r.Method == "POST" {
			w.WriteHeader(http.StatusCreated)
			_ = json.NewEncoder(w).Encode(map[string]string{"id": "test-job-id"})
		} else {
			w.WriteHeader(http.StatusOK)
			_ = json.NewEncoder(w).Encode([]map[string]string{{"id": "test-job-id", "status": "pending"}})
		}
	})
	mux.HandleFunc("/api/v1/jobs/", func(w http.ResponseWriter, _ *http.Request) {
		w.Header().Set("Content-Type", "application/json")
		w.WriteHeader(http.StatusOK)
		_ = json.NewEncoder(w).Encode(map[string]string{"status": "pending"})
	})
	mux.HandleFunc("/api/v1/metrics", func(w http.ResponseWriter, _ *http.Request) {
		w.Header().Set("Content-Type", "application/json")
		w.WriteHeader(http.StatusOK)
		_ = json.NewEncoder(w).Encode(map[string]string{"status": "ok"})
	})
	mux.HandleFunc("/ws", func(w http.ResponseWriter, _ *http.Request) {
		w.WriteHeader(http.StatusOK)
	})

	return httptest.NewServer(mux)
}

// benchmarkCreateJob tests job creation performance
func benchmarkCreateJob(b *testing.B, baseURL string, client *http.Client) {
	for i := 0; i < b.N; i++ {
		jobData := map[string]interface{}{
			"job_name": fmt.Sprintf("benchmark-job-%d", i),
			"args": map[string]interface{}{
				"model": "test-model",
				"data":  generateTestPayload(1024), // 1KB payload
			},
			"priority": 0,
		}

		jsonData, _ := json.Marshal(jobData)
		req, err := http.NewRequestWithContext(context.Background(), "POST",
			baseURL+"/api/v1/jobs", bytes.NewBuffer(jsonData))
		if err != nil {
			b.Fatalf("Failed to create request: %v", err)
		}
		req.Header.Set("Content-Type", "application/json")
		req.Header.Set("Authorization", "Bearer test-token")

		resp, err := client.Do(req)
		if err != nil {
			b.Fatalf("Request failed: %v", err)
		}
		_, _ = io.Copy(io.Discard, resp.Body)
		_ = resp.Body.Close()
	}
}

// benchmarkListJobs tests job listing performance
func benchmarkListJobs(b *testing.B, baseURL string, client *http.Client) {
	for i := 0; i < b.N; i++ {
		req, err := http.NewRequestWithContext(context.Background(), "GET", baseURL+"/api/v1/jobs", nil)
		if err != nil {
			b.Fatalf("Failed to create request: %v", err)
		}
		req.Header.Set("Authorization", "Bearer test-token")

		resp, err := client.Do(req)
		if err != nil {
			b.Fatalf("Request failed: %v", err)
		}
		_, _ = io.Copy(io.Discard, resp.Body)
		_ = resp.Body.Close()
	}
}

// BenchmarkAPIServerListJobs tests job listing performance
func BenchmarkAPIServerListJobs(b *testing.B) {
	server := setupTestAPIServer(b)
	defer server.Close()
	client := newBenchmarkHTTPClient()
	b.ResetTimer()
	b.ReportAllocs()
	benchmarkListJobs(b, server.URL, client)
}

// BenchmarkWebSocketConnection tests WebSocket connection performance
func BenchmarkWebSocketConnection(b *testing.B) {
	server := setupTestAPIServer(b)
	defer server.Close()

	for i := 0; i < b.N; i++ {
		// Convert HTTP URL to WebSocket URL
		wsURL := strings.Replace(server.URL, "http://", "ws://", 1)
		wsURL += "/ws"

		conn, resp, err := websocket.DefaultDialer.Dial(wsURL, nil)
		if resp != nil && resp.Body != nil {
			_ = resp.Body.Close()
		}
		if err != nil {
			// Skip iteration if WebSocket server isn't available
			continue
		}
		_ = conn.Close()
	}
}