github.com/mailgun/holster/v4@v4.20.0/etcdutil/session_test.go (about)

     1  package etcdutil_test
     2  
     3  import (
     4  	"fmt"
     5  	"os"
     6  	"testing"
     7  	"time"
     8  
     9  	"github.com/Shopify/toxiproxy"
    10  	"github.com/mailgun/holster/v4/etcdutil"
    11  	"github.com/sirupsen/logrus"
    12  	"github.com/stretchr/testify/assert"
    13  	"github.com/stretchr/testify/require"
    14  	etcd "go.etcd.io/etcd/client/v3"
    15  )
    16  
    17  var proxy *toxiproxy.Proxy
    18  var client *etcd.Client
    19  
    20  func TestMain(m *testing.M) {
    21  	proxy = toxiproxy.NewProxy()
    22  	proxy.Name = "etcd"
    23  	proxy.Upstream = "localhost:22379"
    24  	proxy.Listen = "0.0.0.0:2379"
    25  
    26  	if err := proxy.Start(); err != nil {
    27  		fmt.Printf("failed to start toxiproxy: %s\n", err)
    28  		os.Exit(1)
    29  	}
    30  
    31  	var err error
    32  	client, err = etcdutil.NewClient(nil)
    33  	if err != nil {
    34  		fmt.Printf("failed to connect to etcd\n")
    35  		os.Exit(1)
    36  	}
    37  
    38  	code := m.Run()
    39  	proxy.Stop()
    40  	client.Close()
    41  	os.Exit(code)
    42  }
    43  
    44  func TestNewSession(t *testing.T) {
    45  	leaseChan := make(chan etcd.LeaseID, 1)
    46  	getLease := func() etcd.LeaseID {
    47  		select {
    48  		case id := <-leaseChan:
    49  			return id
    50  		case <-time.After(time.Second * 5):
    51  			require.FailNow(t, "Timeout waiting for lease id")
    52  		}
    53  		return 0
    54  	}
    55  
    56  	session, err := etcdutil.NewSession(client, etcdutil.SessionConfig{
    57  		Observer: func(leaseId etcd.LeaseID, err error) {
    58  			if err != nil {
    59  				t.Fatal(err)
    60  			}
    61  			leaseChan <- leaseId
    62  		},
    63  	})
    64  	require.Nil(t, err)
    65  	defer session.Close()
    66  
    67  	assert.NotEqual(t, etcdutil.NoLease, getLease())
    68  
    69  	session.Close()
    70  
    71  	assert.Equal(t, etcdutil.NoLease, getLease())
    72  }
    73  
    74  func TestConnectivityLost(t *testing.T) {
    75  	leaseChan := make(chan etcd.LeaseID, 5)
    76  
    77  	getLease := func() etcd.LeaseID {
    78  		select {
    79  		case id := <-leaseChan:
    80  			return id
    81  		case <-time.After(time.Second * 5):
    82  			require.FailNow(t, "Timeout waiting for lease id")
    83  		}
    84  		return 0
    85  	}
    86  
    87  	logrus.SetLevel(logrus.DebugLevel)
    88  
    89  	session, err := etcdutil.NewSession(client, etcdutil.SessionConfig{
    90  		Observer: func(leaseId etcd.LeaseID, err error) {
    91  			if err != nil {
    92  				t.Fatal(err)
    93  			}
    94  			leaseChan <- leaseId
    95  		},
    96  		TTL: 1,
    97  	})
    98  	require.Nil(t, err)
    99  	defer session.Close()
   100  
   101  	// Assert we get a valid lease id
   102  	assert.NotEqual(t, etcdutil.NoLease, getLease())
   103  
   104  	// Interrupt the connection
   105  	proxy.Stop()
   106  
   107  	// Wait for session to realize the connection is gone
   108  	assert.Equal(t, etcdutil.NoLease, getLease())
   109  
   110  	// Restore the connection
   111  	require.Nil(t, proxy.Start())
   112  
   113  	// We should get a new lease
   114  	assert.NotEqual(t, etcdutil.NoLease, getLease())
   115  
   116  	session.Close()
   117  
   118  	// Should get a final NoLease after close
   119  	assert.Equal(t, etcdutil.NoLease, getLease())
   120  }