proxy-pool/internal/controller/runtime/runtime_test.go

package runtime

import (
	"context"
	"errors"
	"io"
	"net"
	"net/http"
	"sync/atomic"
	"testing"
	"time"

	"proxy-pool/internal/config"
	"proxy-pool/internal/controller/admin"
	controllerExtraction "proxy-pool/internal/controller/extraction"
	"proxy-pool/internal/platform/httpserver"
)

func TestRuntimeServesDistributionAndAdminOnIndependentListeners(t *testing.T) {
	t.Parallel()
	cfg := runtimeConfig()
	cfg.Metrics = config.Metrics{Enabled: true, Listen: "127.0.0.1:0"}
	adminService := &stubAdminService{status: admin.Status{ConfigVersion: "cfg-7", SnapshotVersion: 11}}
	runtime, err := New(cfg, Dependencies{
		Extractor:    stubExtractor{},
		Readiness:    stubReadiness{},
		AdminService: adminService,
		MetricsHandler: http.HandlerFunc(func(response http.ResponseWriter, request *http.Request) {
			if request.URL.Path != "/readyz" {
				http.NotFound(response, request)
				return
			}
			response.WriteHeader(http.StatusOK)
		}),
	}, Options{HTTP: testHTTPOptions()})
	if err != nil {
		t.Fatalf("New() error = %v", err)
	}
	distributionListener := mustListen(t)
	adminListener := mustListen(t)
	metricsListener := mustListen(t)
	ctx, cancel := context.WithCancel(context.Background())
	result := make(chan error, 1)
	go func() {
		result <- runtime.Serve(ctx, Listeners{
			Distribution: distributionListener,
			Admin:        adminListener,
			Metrics:      metricsListener,
		})
	}()

	distributionURL := "http://" + distributionListener.Addr().String()
	adminURL := "http://" + adminListener.Addr().String()
	metricsURL := "http://" + metricsListener.Addr().String()
	assertStatus(t, http.MethodGet, distributionURL+"/health/live", nil, http.StatusOK)
	assertStatus(t, http.MethodGet, distributionURL+"/api/v1/status", nil, http.StatusNotFound)
	assertStatus(t, http.MethodGet, adminURL+"/api/v1/status", nil, http.StatusUnauthorized)
	adminHeaders := http.Header{"Authorization": []string{"Bearer admin-token"}}
	assertStatus(t, http.MethodGet, adminURL+"/api/v1/status", adminHeaders, http.StatusOK)
	assertStatus(t, http.MethodGet, adminURL+"/api/v1/audit", adminHeaders, http.StatusOK)
	assertStatus(t, http.MethodGet, adminURL+"/health/live", adminHeaders, http.StatusNotFound)
	assertStatus(t, http.MethodGet, metricsURL+"/readyz", nil, http.StatusOK)
	assertStatus(t, http.MethodGet, metricsURL+"/api/v1/status", nil, http.StatusNotFound)
	if adminService.statusCalls.Load() != 1 {
		t.Fatalf("admin status calls = %d, want 1", adminService.statusCalls.Load())
	}
	if adminService.auditCalls.Load() != 1 {
		t.Fatalf("admin audit calls = %d, want 1", adminService.auditCalls.Load())
	}

	cancel()
	select {
	case err := <-result:
		if err != nil {
			t.Fatalf("Serve() error = %v", err)
		}
	case <-time.After(2 * time.Second):
		t.Fatal("Serve() did not stop after cancellation")
	}
}

func TestNewRejectsInvalidRuntimeDependencies(t *testing.T) {
	t.Parallel()
	validDependencies := Dependencies{
		Extractor: stubExtractor{}, Readiness: stubReadiness{}, AdminService: &stubAdminService{},
	}
	tests := []struct {
		name    string
		config  *config.Config
		deps    Dependencies
		options Options
	}{
		{name: "nil config", deps: validDependencies},
		{name: "no HTTP listener", config: &config.Config{}, deps: validDependencies},
		{name: "distribution extractor", config: runtimeConfig(), deps: Dependencies{Readiness: stubReadiness{}, AdminService: &stubAdminService{}}},
		{name: "distribution readiness", config: runtimeConfig(), deps: Dependencies{Extractor: stubExtractor{}, AdminService: &stubAdminService{}}},
		{name: "admin service", config: runtimeConfig(), deps: Dependencies{Extractor: stubExtractor{}, Readiness: stubReadiness{}}},
		{name: "negative distribution body limit", config: runtimeConfig(), deps: validDependencies, options: Options{DistributionBodyLimitBytes: -1}},
		{name: "negative admin body limit", config: runtimeConfig(), deps: validDependencies, options: Options{AdminBodyLimitBytes: -1}},
	}
	for _, test := range tests {
		test := test
		t.Run(test.name, func(t *testing.T) {
			if _, err := New(test.config, test.deps, test.options); !errors.Is(err, ErrInvalidRuntime) {
				t.Fatalf("New() error = %v, want %v", err, ErrInvalidRuntime)
			}
		})
	}
}

func TestServeRequiresExactlyTheEnabledListeners(t *testing.T) {
	t.Parallel()
	runtime, err := New(runtimeConfig(), Dependencies{
		Extractor: stubExtractor{}, Readiness: stubReadiness{}, AdminService: &stubAdminService{},
	}, Options{})
	if err != nil {
		t.Fatalf("New() error = %v", err)
	}
	listener := mustListen(t)
	defer listener.Close()

	if err := runtime.Serve(context.Background(), Listeners{Distribution: listener}); !errors.Is(err, ErrInvalidRuntime) {
		t.Fatalf("Serve() error = %v, want %v", err, ErrInvalidRuntime)
	}
}

func TestRuntimeServesMetricsOnIndependentListener(t *testing.T) {
	t.Parallel()
	cfg := &config.Config{Metrics: config.Metrics{Enabled: true, Listen: "127.0.0.1:0"}}
	metricsHandler := http.HandlerFunc(func(response http.ResponseWriter, request *http.Request) {
		if request.URL.Path != "/readyz" {
			http.NotFound(response, request)
			return
		}
		response.WriteHeader(http.StatusOK)
	})
	runtime, err := New(cfg, Dependencies{MetricsHandler: metricsHandler}, Options{HTTP: testHTTPOptions()})
	if err != nil {
		t.Fatalf("New(metrics-only) error = %v", err)
	}
	listener := mustListen(t)
	ctx, cancel := context.WithCancel(context.Background())
	result := make(chan error, 1)
	go func() { result <- runtime.Serve(ctx, Listeners{Metrics: listener}) }()
	assertStatus(t, http.MethodGet, "http://"+listener.Addr().String()+"/readyz", nil, http.StatusOK)
	cancel()
	select {
	case err := <-result:
		if err != nil {
			t.Fatalf("Serve(metrics-only) error = %v", err)
		}
	case <-time.After(2 * time.Second):
		t.Fatal("Serve(metrics-only) did not stop")
	}
}

func TestNewRequiresMetricsHandlerWhenEnabled(t *testing.T) {
	t.Parallel()
	cfg := &config.Config{Metrics: config.Metrics{Enabled: true, Listen: "127.0.0.1:0"}}
	if _, err := New(cfg, Dependencies{}, Options{}); !errors.Is(err, ErrInvalidRuntime) {
		t.Fatalf("New(metrics without handler) error = %v", err)
	}
}

func runtimeConfig() *config.Config {
	return &config.Config{
		Distribution: config.Distribution{
			Listener: config.Listener{
				Enabled: true,
				Listen:  "127.0.0.1:0",
				Auth: config.Auth{
					Mode: "apiKey", Header: "X-API-Key", Token: "distribution-token",
				},
			},
			ClientIdentification: config.ClientIdentification{Mode: "authenticatedClient"},
		},
		Admin: config.Listener{
			Enabled: true,
			Listen:  "127.0.0.1:0",
			Auth:    config.Auth{Mode: "bearer", Token: "admin-token"},
		},
	}
}

func testHTTPOptions() httpserver.Options {
	return httpserver.Options{
		ReadHeaderTimeout: time.Second,
		ReadTimeout:       time.Second,
		WriteTimeout:      time.Second,
		IdleTimeout:       time.Second,
		ShutdownTimeout:   time.Second,
		MaxHeaderBytes:    16 << 10,
	}
}

func mustListen(t *testing.T) net.Listener {
	t.Helper()
	listener, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		t.Fatalf("Listen(): %v", err)
	}
	return listener
}

func assertStatus(t *testing.T, method, target string, headers http.Header, want int) {
	t.Helper()
	request, err := http.NewRequest(method, target, nil)
	if err != nil {
		t.Fatalf("NewRequest(): %v", err)
	}
	request.Header = headers.Clone()
	client := &http.Client{Timeout: time.Second}
	response, err := client.Do(request)
	if err != nil {
		t.Fatalf("Do(%s): %v", target, err)
	}
	defer response.Body.Close()
	if response.StatusCode != want {
		body, _ := io.ReadAll(response.Body)
		t.Fatalf("%s status = %d, want %d; body=%s", target, response.StatusCode, want, body)
	}
}

type stubExtractor struct{}

func (stubExtractor) Extract(context.Context, controllerExtraction.Request) (controllerExtraction.Response, error) {
	return controllerExtraction.Response{}, nil
}

type stubReadiness struct{}

func (stubReadiness) Ready(context.Context) error { return nil }

type stubAdminService struct {
	status      admin.Status
	audit       admin.AuditPage
	statusCalls atomic.Int64
	auditCalls  atomic.Int64
}

func (service *stubAdminService) Status(context.Context) (admin.Status, error) {
	service.statusCalls.Add(1)
	return service.status, nil
}

func (service *stubAdminService) ReadAudit(context.Context, admin.AuditQuery) (admin.AuditPage, error) {
	service.auditCalls.Add(1)
	return service.audit, nil
}

func (*stubAdminService) SetUpstreamEnabled(context.Context, admin.SetUpstreamCommand) (admin.MutationResult, error) {
	return admin.MutationResult{}, nil
}

func (*stubAdminService) SwitchRouting(context.Context, admin.SwitchCommand) (admin.MutationResult, error) {
	return admin.MutationResult{}, nil
}

func (*stubAdminService) ReloadConfiguration(context.Context, admin.ReloadCommand) (admin.MutationResult, error) {
	return admin.MutationResult{}, nil
}