105 lines
2.7 KiB
Go
105 lines
2.7 KiB
Go
package server_test
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
"io"
|
|
"net"
|
|
"net/http"
|
|
"sync"
|
|
"testing"
|
|
"time"
|
|
|
|
"forgejo.riotpiao.com/rock/homelab-frontend/internal/server"
|
|
)
|
|
|
|
// TestGracefulShutdown verifies that:
|
|
// - A request in-flight when shutdown starts receives its full, uncorrupted response body
|
|
// - A request arriving after shutdown starts is refused on a new connection
|
|
// - The shutdown completes with exit code 0 (no timeout)
|
|
func TestGracefulShutdown(t *testing.T) {
|
|
// Create a handler that responds slowly
|
|
const responseBody = "slow response body content"
|
|
const sleepDuration = 500 * time.Millisecond
|
|
|
|
handler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
// Simulate a slow LLM response
|
|
time.Sleep(sleepDuration)
|
|
fmt.Fprint(w, responseBody)
|
|
})
|
|
|
|
// Create server with a shutdown timeout longer than the sleep
|
|
srv := server.New("127.0.0.1:0", 5*time.Second, handler)
|
|
|
|
// Start server in a goroutine
|
|
var listenErr error
|
|
var wg sync.WaitGroup
|
|
wg.Add(1)
|
|
go func() {
|
|
defer wg.Done()
|
|
listenErr = srv.ListenAndServe()
|
|
// http.ErrServerClosed is expected after shutdown
|
|
if listenErr != nil && listenErr != http.ErrServerClosed {
|
|
t.Logf("unexpected listen error: %v", listenErr)
|
|
}
|
|
}()
|
|
|
|
// Give server time to start listening
|
|
time.Sleep(100 * time.Millisecond)
|
|
|
|
// Issue a slow request in a goroutine
|
|
var responseBody_got string
|
|
var requestErr error
|
|
var requestWg sync.WaitGroup
|
|
requestWg.Add(1)
|
|
go func() {
|
|
defer requestWg.Done()
|
|
resp, err := http.Get(fmt.Sprintf("http://%s/", srv.Addr()))
|
|
if err != nil {
|
|
requestErr = err
|
|
return
|
|
}
|
|
defer resp.Body.Close()
|
|
|
|
body, err := io.ReadAll(resp.Body)
|
|
if err != nil {
|
|
requestErr = err
|
|
return
|
|
}
|
|
responseBody_got = string(body)
|
|
}()
|
|
|
|
// Give the request time to reach the handler
|
|
time.Sleep(100 * time.Millisecond)
|
|
|
|
// Now initiate shutdown while request is in-flight
|
|
shutdownErr := srv.Shutdown(context.Background())
|
|
|
|
// Wait for the in-flight request to complete
|
|
requestWg.Wait()
|
|
|
|
// Verify the in-flight request completed successfully
|
|
if requestErr != nil {
|
|
t.Fatalf("in-flight request failed: %v", requestErr)
|
|
}
|
|
if responseBody_got != responseBody {
|
|
t.Fatalf("in-flight request got wrong body: %q (expected %q)", responseBody_got, responseBody)
|
|
}
|
|
|
|
// Verify shutdown succeeded (no timeout)
|
|
if shutdownErr != nil {
|
|
t.Fatalf("shutdown failed: %v", shutdownErr)
|
|
}
|
|
|
|
// Verify in-flight requests were allowed to complete
|
|
wg.Wait()
|
|
|
|
// Now verify that a new request is refused after shutdown
|
|
_, err := net.Dial("tcp", srv.Addr())
|
|
if err == nil {
|
|
// Connection succeeded when it should have failed
|
|
t.Fatalf("new connection accepted after shutdown (should have been refused)")
|
|
}
|
|
// If we get here, the connection was properly refused, which is what we want
|
|
}
|