traefik/pkg/healthcheck/mock_test.go

package healthcheck

import (
	"context"
	"net"
	"net/http"
	"net/http/httptest"
	"net/url"
	"sync"
	"testing"
	"time"

	gokitmetrics "github.com/go-kit/kit/metrics"
	"github.com/stretchr/testify/assert"
	"github.com/traefik/traefik/v2/pkg/config/dynamic"
	"github.com/traefik/traefik/v2/pkg/testhelpers"
	"google.golang.org/grpc"
	healthpb "google.golang.org/grpc/health/grpc_health_v1"
)

type StartTestServer interface {
	Start(t *testing.T, done func()) (*url.URL, time.Duration)
}

type Status interface {
	~int | ~int32
}

type HealthSequence[T Status] struct {
	sequenceMu sync.Mutex
	sequence   []T
}

func (s *HealthSequence[T]) Pop() T {
	s.sequenceMu.Lock()
	defer s.sequenceMu.Unlock()

	stat := s.sequence[0]

	s.sequence = s.sequence[1:]

	return stat
}

func (s *HealthSequence[T]) IsEmpty() bool {
	s.sequenceMu.Lock()
	defer s.sequenceMu.Unlock()

	return len(s.sequence) == 0
}

type GRPCServer struct {
	status HealthSequence[healthpb.HealthCheckResponse_ServingStatus]
	done   func()
}

func newGRPCServer(healthSequence ...healthpb.HealthCheckResponse_ServingStatus) *GRPCServer {
	gRPCService := &GRPCServer{
		status: HealthSequence[healthpb.HealthCheckResponse_ServingStatus]{
			sequence: healthSequence,
		},
	}

	return gRPCService
}

func (s *GRPCServer) Check(_ context.Context, _ *healthpb.HealthCheckRequest) (*healthpb.HealthCheckResponse, error) {
	if s.status.IsEmpty() {
		s.done()
		return &healthpb.HealthCheckResponse{
			Status: healthpb.HealthCheckResponse_SERVICE_UNKNOWN,
		}, nil
	}
	stat := s.status.Pop()

	return &healthpb.HealthCheckResponse{
		Status: stat,
	}, nil
}

func (s *GRPCServer) Watch(_ *healthpb.HealthCheckRequest, server healthpb.Health_WatchServer) error {
	if s.status.IsEmpty() {
		s.done()
		return server.Send(&healthpb.HealthCheckResponse{
			Status: healthpb.HealthCheckResponse_SERVICE_UNKNOWN,
		})
	}
	stat := s.status.Pop()

	return server.Send(&healthpb.HealthCheckResponse{
		Status: stat,
	})
}

func (s *GRPCServer) Start(t *testing.T, done func()) (*url.URL, time.Duration) {
	t.Helper()

	listener, err := net.Listen("tcp4", "127.0.0.1:0")
	assert.NoError(t, err)
	t.Cleanup(func() { _ = listener.Close() })

	server := grpc.NewServer()
	t.Cleanup(server.Stop)

	s.done = done

	healthpb.RegisterHealthServer(server, s)

	go func() {
		err := server.Serve(listener)
		assert.NoError(t, err)
	}()

	// Make test timeout dependent on number of expected requests, health check interval, and a safety margin.
	return testhelpers.MustParseURL("http://" + listener.Addr().String()), time.Duration(len(s.status.sequence)*int(dynamic.DefaultHealthCheckInterval) + 500)
}

type HTTPServer struct {
	status HealthSequence[int]
	done   func()
}

func newHTTPServer(healthSequence ...int) *HTTPServer {
	handler := &HTTPServer{
		status: HealthSequence[int]{
			sequence: healthSequence,
		},
	}

	return handler
}

// ServeHTTP returns HTTP response codes following a status sequences.
// It calls the given 'done' function once all request health indicators have been depleted.
func (s *HTTPServer) ServeHTTP(w http.ResponseWriter, _ *http.Request) {
	if s.status.IsEmpty() {
		s.done()
		return
	}

	stat := s.status.Pop()

	w.WriteHeader(stat)
}

func (s *HTTPServer) Start(t *testing.T, done func()) (*url.URL, time.Duration) {
	t.Helper()

	s.done = done

	ts := httptest.NewServer(s)
	t.Cleanup(ts.Close)

	// Make test timeout dependent on number of expected requests, health check interval, and a safety margin.
	return testhelpers.MustParseURL(ts.URL), time.Duration(len(s.status.sequence)*int(dynamic.DefaultHealthCheckInterval) + 500)
}

type testLoadBalancer struct {
	// RWMutex needed due to parallel test execution: Both the system-under-test
	// and the test assertions reference the counters.
	*sync.RWMutex
	numRemovedServers  int
	numUpsertedServers int
}

func (lb *testLoadBalancer) SetStatus(ctx context.Context, childName string, up bool) {
	if up {
		lb.numUpsertedServers++
	} else {
		lb.numRemovedServers++
	}
}

type MetricsMock struct {
	Gauge gokitmetrics.Gauge
}

func (m *MetricsMock) ServiceServerUpGauge() gokitmetrics.Gauge {
	return m.Gauge
}