fm_be/internal/app/api/server_test.go

package api

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

	"github.com/gofiber/fiber/v2"
)

func TestServerShutdownDrainsInFlightRequestsAndRejectsNewRequests(t *testing.T) {
	app := fiber.New(fiber.Config{
		DisableStartupMessage: true,
	})
	UseDefaultMiddlewares(app, testHTTPConfig())

	slowStarted := make(chan struct{})
	app.Get("/slow", func(c *fiber.Ctx) error {
		select {
		case <-slowStarted:
		default:
			close(slowStarted)
		}

		select {
		case <-time.After(300 * time.Millisecond):
		case <-c.UserContext().Done():
			return c.Status(http.StatusInternalServerError).SendString("user context canceled")
		}

		return c.SendString("slow done")
	})
	app.Get("/fast", func(c *fiber.Ctx) error {
		return c.SendString("fast done")
	})

	ln, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		t.Fatalf("listen: %v", err)
	}

	serveErrCh := make(chan error, 1)
	go func() {
		serveErrCh <- app.Listener(ln)
	}()

	t.Cleanup(func() {
		_ = app.Shutdown()
	})

	baseURL := "http://" + ln.Addr().String()
	fastClient := newTestHTTPClient()
	slowClient := newTestHTTPClient()
	defer fastClient.CloseIdleConnections()
	defer slowClient.CloseIdleConnections()

	if err := waitForHTTP200(fastClient, baseURL+"/fast", time.Second); err != nil {
		t.Fatalf("wait for server ready: %v", err)
	}

	slowRespCh := make(chan httpCallResult, 1)
	go func() {
		resp, err := slowClient.Get(baseURL + "/slow")
		if err != nil {
			slowRespCh <- httpCallResult{err: err}
			return
		}
		defer resp.Body.Close()

		body, readErr := io.ReadAll(resp.Body)
		slowRespCh <- httpCallResult{
			statusCode: resp.StatusCode,
			body:       string(body),
			err:        readErr,
		}
	}()

	select {
	case <-slowStarted:
	case <-time.After(time.Second):
		t.Fatal("slow request did not start")
	}

	srv := NewServer(app, WithShutdownTimeout(2*time.Second))
	shutdownErrCh := make(chan error, 1)
	go func() {
		shutdownErrCh <- srv.Shutdown(context.Background())
	}()

	if err := waitForRequestRejection(fastClient, baseURL+"/fast", time.Second); err == nil {
		t.Fatal("expected new requests to be rejected during shutdown")
	}

	slowResult := <-slowRespCh
	if slowResult.err != nil {
		t.Fatalf("slow request failed: %v", slowResult.err)
	}
	if slowResult.statusCode != http.StatusOK {
		t.Fatalf("expected slow request 200, got %d with body %q", slowResult.statusCode, slowResult.body)
	}
	if slowResult.body != "slow done" {
		t.Fatalf("expected slow request body %q, got %q", "slow done", slowResult.body)
	}

	if err := <-shutdownErrCh; err != nil {
		t.Fatalf("shutdown failed: %v", err)
	}
	if err := <-serveErrCh; err != nil {
		t.Fatalf("listener exited with error: %v", err)
	}
}

type httpCallResult struct {
	statusCode int
	body       string
	err        error
}

func newTestHTTPClient() *http.Client {
	return &http.Client{
		Timeout: time.Second,
		Transport: &http.Transport{
			DisableKeepAlives: true,
		},
	}
}

func waitForHTTP200(client *http.Client, url string, timeout time.Duration) error {
	deadline := time.Now().Add(timeout)
	for time.Now().Before(deadline) {
		resp, err := client.Get(url)
		if err == nil {
			resp.Body.Close()
			if resp.StatusCode == http.StatusOK {
				return nil
			}
		}
		time.Sleep(10 * time.Millisecond)
	}
	return errors.New("server did not become ready")
}

func waitForRequestRejection(client *http.Client, url string, timeout time.Duration) error {
	deadline := time.Now().Add(timeout)
	var lastErr error
	for time.Now().Before(deadline) {
		resp, err := client.Get(url)
		if err != nil {
			return err
		}
		lastErr = errors.New("request still accepted")
		resp.Body.Close()
		time.Sleep(10 * time.Millisecond)
	}
	return lastErr
}