github.com/blend/go-sdk@v1.20220411.3/graceful/shutdown_by_signal_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 graceful
     9  
    10  import (
    11  	"fmt"
    12  	"os"
    13  	"testing"
    14  
    15  	"github.com/blend/go-sdk/assert"
    16  )
    17  
    18  func newHosted() *hosted {
    19  	return &hosted{
    20  		started: make(chan struct{}),
    21  		stopped: make(chan struct{}),
    22  	}
    23  }
    24  
    25  type hosted struct {
    26  	state   int32
    27  	started chan struct{}
    28  	stopped chan struct{}
    29  }
    30  
    31  func (h *hosted) Start() error {
    32  	h.state = 1
    33  	h.stopped = make(chan struct{})
    34  	close(h.started)
    35  	<-h.stopped
    36  	return nil
    37  }
    38  
    39  func (h *hosted) Stop() error {
    40  	if h.state != 1 {
    41  		return fmt.Errorf("cannot stop")
    42  	}
    43  	h.state = 0
    44  	h.started = make(chan struct{})
    45  	close(h.stopped)
    46  	return nil
    47  }
    48  
    49  func (h *hosted) NotifyStarted() <-chan struct{} {
    50  	return h.started
    51  }
    52  
    53  func (h *hosted) NotifyStopped() <-chan struct{} {
    54  	return h.stopped
    55  }
    56  
    57  func TestShutdownBySignal(t *testing.T) {
    58  	assert := assert.New(t)
    59  
    60  	hosted := newHosted()
    61  
    62  	terminateSignal := make(chan os.Signal)
    63  	var err error
    64  	done := make(chan struct{})
    65  	go func() {
    66  		err = ShutdownBySignal([]Graceful{hosted}, OptShutdownSignal(terminateSignal))
    67  		close(done)
    68  	}()
    69  	<-hosted.NotifyStarted()
    70  
    71  	close(terminateSignal)
    72  	<-done
    73  	assert.Nil(err)
    74  }
    75  
    76  func TestShutdownBySignalMany(t *testing.T) {
    77  	assert := assert.New(t)
    78  
    79  	workers := []Graceful{
    80  		newHosted(),
    81  		newHosted(),
    82  		newHosted(),
    83  		newHosted(),
    84  		newHosted(),
    85  	}
    86  
    87  	terminateSignal := make(chan os.Signal)
    88  	var err error
    89  	done := make(chan struct{})
    90  
    91  	go func() {
    92  		err = ShutdownBySignal(workers, OptShutdownSignal(terminateSignal))
    93  		close(done)
    94  	}()
    95  
    96  	// wait for the workers to start
    97  	for _, h := range workers {
    98  		<-h.(*hosted).started
    99  	}
   100  	for _, h := range workers {
   101  		assert.Equal(1, h.(*hosted).state)
   102  	}
   103  
   104  	close(terminateSignal)
   105  	<-done
   106  	assert.Nil(err)
   107  	for _, h := range workers {
   108  		assert.Equal(0, h.(*hosted).state)
   109  	}
   110  }