github.com/lingyao2333/mo-zero@v1.4.1/core/discov/publisher_test.go (about)

     1  package discov
     2  
     3  import (
     4  	"errors"
     5  	"sync"
     6  	"testing"
     7  	"time"
     8  
     9  	"github.com/golang/mock/gomock"
    10  	"github.com/lingyao2333/mo-zero/core/discov/internal"
    11  	"github.com/lingyao2333/mo-zero/core/lang"
    12  	"github.com/lingyao2333/mo-zero/core/logx"
    13  	"github.com/lingyao2333/mo-zero/core/stringx"
    14  	"github.com/stretchr/testify/assert"
    15  	clientv3 "go.etcd.io/etcd/client/v3"
    16  )
    17  
    18  func init() {
    19  	logx.Disable()
    20  }
    21  
    22  func TestPublisher_register(t *testing.T) {
    23  	ctrl := gomock.NewController(t)
    24  	defer ctrl.Finish()
    25  	const id = 1
    26  	cli := internal.NewMockEtcdClient(ctrl)
    27  	restore := setMockClient(cli)
    28  	defer restore()
    29  	cli.EXPECT().Ctx().AnyTimes()
    30  	cli.EXPECT().Grant(gomock.Any(), timeToLive).Return(&clientv3.LeaseGrantResponse{
    31  		ID: id,
    32  	}, nil)
    33  	cli.EXPECT().Put(gomock.Any(), makeEtcdKey("thekey", id), "thevalue", gomock.Any())
    34  	pub := NewPublisher(nil, "thekey", "thevalue",
    35  		WithPubEtcdAccount(stringx.Rand(), "bar"))
    36  	_, err := pub.register(cli)
    37  	assert.Nil(t, err)
    38  }
    39  
    40  func TestPublisher_registerWithId(t *testing.T) {
    41  	ctrl := gomock.NewController(t)
    42  	defer ctrl.Finish()
    43  	const id = 2
    44  	cli := internal.NewMockEtcdClient(ctrl)
    45  	restore := setMockClient(cli)
    46  	defer restore()
    47  	cli.EXPECT().Ctx().AnyTimes()
    48  	cli.EXPECT().Grant(gomock.Any(), timeToLive).Return(&clientv3.LeaseGrantResponse{
    49  		ID: 1,
    50  	}, nil)
    51  	cli.EXPECT().Put(gomock.Any(), makeEtcdKey("thekey", id), "thevalue", gomock.Any())
    52  	pub := NewPublisher(nil, "thekey", "thevalue", WithId(id))
    53  	_, err := pub.register(cli)
    54  	assert.Nil(t, err)
    55  }
    56  
    57  func TestPublisher_registerError(t *testing.T) {
    58  	ctrl := gomock.NewController(t)
    59  	defer ctrl.Finish()
    60  	cli := internal.NewMockEtcdClient(ctrl)
    61  	restore := setMockClient(cli)
    62  	defer restore()
    63  	cli.EXPECT().Ctx().AnyTimes()
    64  	cli.EXPECT().Grant(gomock.Any(), timeToLive).Return(nil, errors.New("error"))
    65  	pub := NewPublisher(nil, "thekey", "thevalue")
    66  	val, err := pub.register(cli)
    67  	assert.NotNil(t, err)
    68  	assert.Equal(t, clientv3.NoLease, val)
    69  }
    70  
    71  func TestPublisher_revoke(t *testing.T) {
    72  	ctrl := gomock.NewController(t)
    73  	defer ctrl.Finish()
    74  	const id clientv3.LeaseID = 1
    75  	cli := internal.NewMockEtcdClient(ctrl)
    76  	restore := setMockClient(cli)
    77  	defer restore()
    78  	cli.EXPECT().Ctx().AnyTimes()
    79  	cli.EXPECT().Revoke(gomock.Any(), id)
    80  	pub := NewPublisher(nil, "thekey", "thevalue")
    81  	pub.lease = id
    82  	pub.revoke(cli)
    83  }
    84  
    85  func TestPublisher_revokeError(t *testing.T) {
    86  	ctrl := gomock.NewController(t)
    87  	defer ctrl.Finish()
    88  	const id clientv3.LeaseID = 1
    89  	cli := internal.NewMockEtcdClient(ctrl)
    90  	restore := setMockClient(cli)
    91  	defer restore()
    92  	cli.EXPECT().Ctx().AnyTimes()
    93  	cli.EXPECT().Revoke(gomock.Any(), id).Return(nil, errors.New("error"))
    94  	pub := NewPublisher(nil, "thekey", "thevalue")
    95  	pub.lease = id
    96  	pub.revoke(cli)
    97  }
    98  
    99  func TestPublisher_keepAliveAsyncError(t *testing.T) {
   100  	ctrl := gomock.NewController(t)
   101  	defer ctrl.Finish()
   102  	const id clientv3.LeaseID = 1
   103  	cli := internal.NewMockEtcdClient(ctrl)
   104  	restore := setMockClient(cli)
   105  	defer restore()
   106  	cli.EXPECT().Ctx().AnyTimes()
   107  	cli.EXPECT().KeepAlive(gomock.Any(), id).Return(nil, errors.New("error"))
   108  	pub := NewPublisher(nil, "thekey", "thevalue")
   109  	pub.lease = id
   110  	assert.NotNil(t, pub.keepAliveAsync(cli))
   111  }
   112  
   113  func TestPublisher_keepAliveAsyncQuit(t *testing.T) {
   114  	ctrl := gomock.NewController(t)
   115  	defer ctrl.Finish()
   116  	const id clientv3.LeaseID = 1
   117  	cli := internal.NewMockEtcdClient(ctrl)
   118  	cli.EXPECT().ActiveConnection()
   119  	cli.EXPECT().Close()
   120  	defer cli.Close()
   121  	cli.ActiveConnection()
   122  	restore := setMockClient(cli)
   123  	defer restore()
   124  	cli.EXPECT().Ctx().AnyTimes()
   125  	cli.EXPECT().KeepAlive(gomock.Any(), id)
   126  	var wg sync.WaitGroup
   127  	wg.Add(1)
   128  	cli.EXPECT().Revoke(gomock.Any(), id).Do(func(_, _ interface{}) {
   129  		wg.Done()
   130  	})
   131  	pub := NewPublisher(nil, "thekey", "thevalue")
   132  	pub.lease = id
   133  	pub.Stop()
   134  	assert.Nil(t, pub.keepAliveAsync(cli))
   135  	wg.Wait()
   136  }
   137  
   138  func TestPublisher_keepAliveAsyncPause(t *testing.T) {
   139  	ctrl := gomock.NewController(t)
   140  	defer ctrl.Finish()
   141  	const id clientv3.LeaseID = 1
   142  	cli := internal.NewMockEtcdClient(ctrl)
   143  	restore := setMockClient(cli)
   144  	defer restore()
   145  	cli.EXPECT().Ctx().AnyTimes()
   146  	cli.EXPECT().KeepAlive(gomock.Any(), id)
   147  	pub := NewPublisher(nil, "thekey", "thevalue")
   148  	var wg sync.WaitGroup
   149  	wg.Add(1)
   150  	cli.EXPECT().Revoke(gomock.Any(), id).Do(func(_, _ interface{}) {
   151  		pub.Stop()
   152  		wg.Done()
   153  	})
   154  	pub.lease = id
   155  	assert.Nil(t, pub.keepAliveAsync(cli))
   156  	pub.Pause()
   157  	wg.Wait()
   158  }
   159  
   160  func TestPublisher_Resume(t *testing.T) {
   161  	publisher := new(Publisher)
   162  	publisher.resumeChan = make(chan lang.PlaceholderType)
   163  	go func() {
   164  		publisher.Resume()
   165  	}()
   166  	go func() {
   167  		time.Sleep(time.Minute)
   168  		t.Fail()
   169  	}()
   170  	<-publisher.resumeChan
   171  }