github.com/anth0d/nomad@v0.0.0-20221214183521-ae3a0a2cad06/client/allocrunner/tasklifecycle/gate_test.go (about)

     1  package tasklifecycle
     2  
     3  import (
     4  	"testing"
     5  	"time"
     6  
     7  	"github.com/hashicorp/nomad/ci"
     8  	"github.com/hashicorp/nomad/helper"
     9  )
    10  
    11  func TestGate(t *testing.T) {
    12  	ci.Parallel(t)
    13  
    14  	testCases := []struct {
    15  		name string
    16  		test func(*testing.T, *Gate)
    17  	}{
    18  		{
    19  			name: "starts blocked",
    20  			test: func(t *testing.T, g *Gate) {
    21  				requireChannelBlocking(t, g.WaitCh(), "wait")
    22  			},
    23  		},
    24  		{
    25  			name: "block",
    26  			test: func(t *testing.T, g *Gate) {
    27  				g.Close()
    28  				requireChannelBlocking(t, g.WaitCh(), "wait")
    29  			},
    30  		},
    31  		{
    32  			name: "allow",
    33  			test: func(t *testing.T, g *Gate) {
    34  				g.Open()
    35  				requireChannelPassing(t, g.WaitCh(), "wait")
    36  			},
    37  		},
    38  		{
    39  			name: "block twice",
    40  			test: func(t *testing.T, g *Gate) {
    41  				g.Close()
    42  				g.Close()
    43  				requireChannelBlocking(t, g.WaitCh(), "wait")
    44  			},
    45  		},
    46  		{
    47  			name: "allow twice",
    48  			test: func(t *testing.T, g *Gate) {
    49  				g.Open()
    50  				g.Open()
    51  				requireChannelPassing(t, g.WaitCh(), "wait")
    52  			},
    53  		},
    54  		{
    55  			name: "allow block allow",
    56  			test: func(t *testing.T, g *Gate) {
    57  				g.Open()
    58  				requireChannelPassing(t, g.WaitCh(), "first allow")
    59  				g.Close()
    60  				requireChannelBlocking(t, g.WaitCh(), "block")
    61  				g.Open()
    62  				requireChannelPassing(t, g.WaitCh(), "second allow")
    63  			},
    64  		},
    65  	}
    66  
    67  	for _, tc := range testCases {
    68  		t.Run(tc.name, func(t *testing.T) {
    69  			shutdownCh := make(chan struct{})
    70  			defer close(shutdownCh)
    71  
    72  			g := NewGate(shutdownCh)
    73  			tc.test(t, g)
    74  		})
    75  	}
    76  }
    77  
    78  // TestGate_shutdown tests a gate with a closed shutdown channel.
    79  func TestGate_shutdown(t *testing.T) {
    80  	ci.Parallel(t)
    81  
    82  	// Create a Gate with a closed shutdownCh.
    83  	shutdownCh := make(chan struct{})
    84  	close(shutdownCh)
    85  
    86  	g := NewGate(shutdownCh)
    87  
    88  	// Test that Open() and Close() doesn't block forever.
    89  	openCh := make(chan struct{})
    90  	closeCh := make(chan struct{})
    91  
    92  	go func() {
    93  		g.Open()
    94  		close(openCh)
    95  	}()
    96  	go func() {
    97  		g.Close()
    98  		close(closeCh)
    99  	}()
   100  
   101  	timer, stop := helper.NewSafeTimer(time.Second)
   102  	defer stop()
   103  
   104  	select {
   105  	case <-openCh:
   106  	case <-timer.C:
   107  		t.Fatalf("timeout waiting for gate operations")
   108  	}
   109  
   110  	select {
   111  	case <-closeCh:
   112  	case <-timer.C:
   113  		t.Fatalf("timeout waiting for gate operations")
   114  	}
   115  
   116  	// A Gate with a shutdownCh should be closed.
   117  	requireChannelBlocking(t, g.WaitCh(), "gate should be closed")
   118  }