github.com/blend/go-sdk@v1.20220411.3/webutil/graceful_http_server_test.go (about)

     1  /*
     2  
     3  Copyright (c) 2022 - Present. Blend Labs, Inc. All rights reserved
     4  Use of this source code is governed by a MIT license that can be found in the LICENSE file.
     5  
     6  */
     7  
     8  package webutil
     9  
    10  import (
    11  	"fmt"
    12  	"net"
    13  	"net/http"
    14  	"os"
    15  	"testing"
    16  
    17  	"github.com/blend/go-sdk/assert"
    18  	"github.com/blend/go-sdk/graceful"
    19  )
    20  
    21  var (
    22  	_ graceful.Graceful = (*GracefulHTTPServer)(nil)
    23  )
    24  
    25  func TestGracefulServer(t *testing.T) {
    26  	assert := assert.New(t)
    27  
    28  	listener, err := net.Listen("tcp", "127.0.0.1:")
    29  	assert.Nil(err)
    30  	typedListener, ok := listener.(*net.TCPListener)
    31  	assert.True(ok)
    32  	assert.NotNil(typedListener)
    33  
    34  	server := &http.Server{
    35  		Handler: http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
    36  			w.WriteHeader(http.StatusOK)
    37  			fmt.Fprintf(w, "OK!\n")
    38  		}),
    39  	}
    40  	gs := NewGracefulHTTPServer(server, OptGracefulHTTPServerListener(typedListener))
    41  	stopSignal := make(chan os.Signal)
    42  	didShutdown := make(chan struct{})
    43  
    44  	go func() {
    45  		defer func() { close(didShutdown) }()
    46  		_ = graceful.ShutdownBySignal([]graceful.Graceful{gs}, graceful.OptShutdownSignal(stopSignal))
    47  	}()
    48  	<-gs.NotifyStarted()
    49  
    50  	res, err := http.Get("http://" + typedListener.Addr().String())
    51  	assert.Nil(err)
    52  	assert.Equal(http.StatusOK, res.StatusCode)
    53  
    54  	stopSignal <- os.Interrupt
    55  	<-didShutdown
    56  }