func TestWatch(t *testing.T) {
storeMock, err := libkvmock.New([]string{"127.0.0.1:1234"}, nil)
assert.NotNil(t, storeMock)
assert.NoError(t, err)
d := &Discovery{backend: store.CONSUL}
d.Initialize("127.0.0.1:1234/path", 0, 0)
d.store = storeMock
s := d.store.(*libkvmock.Mock)
mockCh := make(chan []*store.KVPair)
// The first watch will fail.
s.On("WatchTree", "path/"+discoveryPath, mock.Anything).Return(mockCh, errors.New("test error")).Once()
// The second one will succeed.
s.On("WatchTree", "path/"+discoveryPath, mock.Anything).Return(mockCh, nil).Once()
expected := discovery.Entries{
&discovery.Entry{Host: "1.1.1.1", Port: "1111"},
&discovery.Entry{Host: "2.2.2.2", Port: "2222"},
}
kvs := []*store.KVPair{
{Key: path.Join("path", discoveryPath, "1.1.1.1"), Value: []byte("1.1.1.1:1111")},
{Key: path.Join("path", discoveryPath, "2.2.2.2"), Value: []byte("2.2.2.2:2222")},
}
stopCh := make(chan struct{})
ch, errCh := d.Watch(stopCh)
// It should fire an error since the first WatchTree call failed.
assert.EqualError(t, <-errCh, "test error")
// We have to drain the error channel otherwise Watch will get stuck.
go func() {
for range errCh {
}
}()
// Push the entries into the store channel and make sure discovery emits.
mockCh <- kvs
assert.Equal(t, <-ch, expected)
// Add a new entry.
expected = append(expected, &discovery.Entry{Host: "3.3.3.3", Port: "3333"})
kvs = append(kvs, &store.KVPair{Key: path.Join("path", discoveryPath, "3.3.3.3"), Value: []byte("3.3.3.3:3333")})
mockCh <- kvs
assert.Equal(t, <-ch, expected)
// Make sure that if an error occurs it retries.
// This third call to WatchTree will be checked later by AssertExpectations.
s.On("WatchTree", "path/"+discoveryPath, mock.Anything).Return(mockCh, nil)
close(mockCh)
// Give it enough time to call WatchTree.
time.Sleep(3)
// Stop and make sure it closes all channels.
close(stopCh)
assert.Nil(t, <-ch)
assert.Nil(t, <-errCh)
s.AssertExpectations(t)
}
func TestInitialize(t *testing.T) {
storeMock, err := libkvmock.New([]string{"127.0.0.1"}, nil)
assert.NotNil(t, storeMock)
assert.NoError(t, err)
d := &Discovery{backend: store.CONSUL}
d.Initialize("127.0.0.1", 0, 0, nil)
d.store = storeMock
s := d.store.(*libkvmock.Mock)
assert.Len(t, s.Endpoints, 1)
assert.Equal(t, s.Endpoints[0], "127.0.0.1")
assert.Equal(t, d.path, discoveryPath)
storeMock, err = libkvmock.New([]string{"127.0.0.1:1234"}, nil)
assert.NotNil(t, storeMock)
assert.NoError(t, err)
d = &Discovery{backend: store.CONSUL}
d.Initialize("127.0.0.1:1234/path", 0, 0, nil)
d.store = storeMock
s = d.store.(*libkvmock.Mock)
assert.Len(t, s.Endpoints, 1)
assert.Equal(t, s.Endpoints[0], "127.0.0.1:1234")
assert.Equal(t, d.path, "path/"+discoveryPath)
storeMock, err = libkvmock.New([]string{"127.0.0.1:1234", "127.0.0.2:1234", "127.0.0.3:1234"}, nil)
assert.NotNil(t, storeMock)
assert.NoError(t, err)
d = &Discovery{backend: store.CONSUL}
d.Initialize("127.0.0.1:1234,127.0.0.2:1234,127.0.0.3:1234/path", 0, 0, nil)
d.store = storeMock
s = d.store.(*libkvmock.Mock)
if assert.Len(t, s.Endpoints, 3) {
assert.Equal(t, s.Endpoints[0], "127.0.0.1:1234")
assert.Equal(t, s.Endpoints[1], "127.0.0.2:1234")
assert.Equal(t, s.Endpoints[2], "127.0.0.3:1234")
}
assert.Equal(t, d.path, "path/"+discoveryPath)
}
func (ds *DiscoverySuite) TestInitialize(c *check.C) {
storeMock, err := libkvmock.New([]string{"127.0.0.1"}, nil)
c.Assert(storeMock, check.NotNil)
c.Assert(err, check.IsNil)
d := &Discovery{backend: store.CONSUL}
d.Initialize("127.0.0.1", 0, 0)
d.store = storeMock
s := d.store.(*libkvmock.Mock)
c.Assert(s.Endpoints, check.HasLen, 1)
c.Assert(s.Endpoints[0], check.Equals, "127.0.0.1")
c.Assert(d.path, check.Equals, discoveryPath)
storeMock, err = libkvmock.New([]string{"127.0.0.1:1234"}, nil)
c.Assert(storeMock, check.NotNil)
c.Assert(err, check.IsNil)
d = &Discovery{backend: store.CONSUL}
d.Initialize("127.0.0.1:1234/path", 0, 0)
d.store = storeMock
s = d.store.(*libkvmock.Mock)
c.Assert(s.Endpoints, check.HasLen, 1)
c.Assert(s.Endpoints[0], check.Equals, "127.0.0.1:1234")
c.Assert(d.path, check.Equals, "path/"+discoveryPath)
storeMock, err = libkvmock.New([]string{"127.0.0.1:1234", "127.0.0.2:1234", "127.0.0.3:1234"}, nil)
c.Assert(storeMock, check.NotNil)
c.Assert(err, check.IsNil)
d = &Discovery{backend: store.CONSUL}
d.Initialize("127.0.0.1:1234,127.0.0.2:1234,127.0.0.3:1234/path", 0, 0)
d.store = storeMock
s = d.store.(*libkvmock.Mock)
c.Assert(s.Endpoints, check.HasLen, 3)
c.Assert(s.Endpoints[0], check.Equals, "127.0.0.1:1234")
c.Assert(s.Endpoints[1], check.Equals, "127.0.0.2:1234")
c.Assert(s.Endpoints[2], check.Equals, "127.0.0.3:1234")
c.Assert(d.path, check.Equals, "path/"+discoveryPath)
}
func TestCandidate(t *testing.T) {
kv, err := libkvmock.New([]string{}, nil)
assert.NoError(t, err)
assert.NotNil(t, kv)
mockStore := kv.(*libkvmock.Mock)
mockLock := &libkvmock.Lock{}
mockStore.On("NewLock", "test_key", mock.Anything).Return(mockLock, nil)
// Lock and unlock always succeeds.
lostCh := make(chan struct{})
var mockLostCh <-chan struct{} = lostCh
mockLock.On("Lock", mock.Anything).Return(mockLostCh, nil)
mockLock.On("Unlock").Return(nil)
candidate := NewCandidate(kv, "test_key", "test_node", 0)
electedCh, _, err := candidate.RunForElection()
assert.Nil(t, err)
// Should issue a false upon start, no matter what.
assert.False(t, <-electedCh)
// Since the lock always succeeeds, we should get elected.
assert.True(t, <-electedCh)
assert.True(t, candidate.IsLeader())
// Signaling a lost lock should get us de-elected...
close(lostCh)
assert.False(t, <-electedCh)
// And we should attempt to get re-elected again.
assert.True(t, <-electedCh)
// When we resign, unlock will get called, we'll be notified of the
// de-election and we'll try to get the lock again.
go candidate.Resign()
assert.False(t, <-electedCh)
assert.True(t, <-electedCh)
candidate.Stop()
// Ensure that the chan closes after some time
for {
select {
case _, open := <-electedCh:
if !open {
mockStore.AssertExpectations(t)
return
}
case <-time.After(1 * time.Second):
t.Fatalf("electedCh not closed correctly")
}
}
}
func TestFollower(t *testing.T) {
kv, err := libkvmock.New([]string{}, nil)
assert.NoError(t, err)
assert.NotNil(t, kv)
mockStore := kv.(*libkvmock.Mock)
kvCh := make(chan *store.KVPair)
var mockKVCh <-chan *store.KVPair = kvCh
mockStore.On("Watch", "test_key", mock.Anything).Return(mockKVCh, nil)
follower := NewFollower(kv, "test_key")
leaderCh, errCh, err := follower.FollowElection()
assert.Nil(t, err)
// Simulate leader updates
go func() {
kvCh <- &store.KVPair{Key: "test_key", Value: []byte("leader1")}
kvCh <- &store.KVPair{Key: "test_key", Value: []byte("leader1")}
kvCh <- &store.KVPair{Key: "test_key", Value: []byte("leader2")}
kvCh <- &store.KVPair{Key: "test_key", Value: []byte("leader1")}
}()
// We shouldn't see duplicate events.
assert.Equal(t, <-leaderCh, "leader1")
assert.Equal(t, <-leaderCh, "leader2")
assert.Equal(t, <-leaderCh, "leader1")
assert.Equal(t, follower.Leader(), "leader1")
// Once stopped, iteration over the leader channel should stop.
follower.Stop()
close(kvCh)
// Assert that we receive an error from the error chan to deal with the failover
err, open := <-errCh
assert.True(t, open)
assert.NotNil(t, err)
// Ensure that the chan is closed
_, open = <-leaderCh
assert.False(t, open)
mockStore.AssertExpectations(t)
}
func TestFollower(t *testing.T) {
kv, err := libkvmock.New([]string{}, nil)
assert.NoError(t, err)
assert.NotNil(t, kv)
mockStore := kv.(*libkvmock.Mock)
kvCh := make(chan *store.KVPair)
var mockKVCh <-chan *store.KVPair = kvCh
mockStore.On("Watch", "test_key", mock.Anything).Return(mockKVCh, nil)
follower := NewFollower(kv, "test_key")
follower.FollowElection()
leaderCh := follower.LeaderCh()
// Simulate leader updates
go func() {
kvCh <- &store.KVPair{Key: "test_key", Value: []byte("leader1")}
kvCh <- &store.KVPair{Key: "test_key", Value: []byte("leader1")}
kvCh <- &store.KVPair{Key: "test_key", Value: []byte("leader2")}
kvCh <- &store.KVPair{Key: "test_key", Value: []byte("leader1")}
}()
// We shouldn't see duplicate events.
assert.Equal(t, <-leaderCh, "leader1")
assert.Equal(t, <-leaderCh, "leader2")
assert.Equal(t, <-leaderCh, "leader1")
assert.Equal(t, follower.Leader(), "leader1")
// Once stopped, iteration over the leader channel should stop.
follower.Stop()
close(kvCh)
assert.Equal(t, "", <-leaderCh)
mockStore.AssertExpectations(t)
}