package ipam import ( "fmt" "math/rand" "sync" "testing" "time" "github.com/weaveworks/mesh" "github.com/rajch/weave/net/address" "github.com/rajch/weave/testing/gossip" "github.com/stretchr/testify/require" ) type mockMessage struct { dst mesh.PeerName msgType byte buf []byte } func (m *mockMessage) String() string { return fmt.Sprintf("-> %s [%x]", m.dst, m.buf) } type mockGossipComms struct { sync.RWMutex *testing.T name string messages []mockMessage } func (m *mockGossipComms) String() string { m.RLock() defer m.RUnlock() return fmt.Sprintf("[mockGossipComms %s]", m.name) } // Note: this style of verification, using equalByteBuffer, requires // that the contents of messages are never re-ordered. Which, for instance, // requires they are not based off iterating through a map. func (m *mockGossipComms) GossipBroadcast(update mesh.GossipData) { m.Lock() defer m.Unlock() buf := []byte{} if len(m.messages) == 0 { require.FailNow(m, fmt.Sprintf("%s: Gossip broadcast message unexpected: \n%x", m.name, buf)) } else if msg := m.messages[0]; msg.dst != mesh.UnknownPeerName { require.FailNow(m, fmt.Sprintf("%s: Expected Gossip message to %s but got broadcast", m.name, msg.dst)) } else if msg.buf != nil && !equalByteBuffer(msg.buf, buf) { require.FailNow(m, fmt.Sprintf("%s: Gossip message not sent as expected: \nwant: %x\ngot : %x", m.name, msg.buf, buf)) } else { // Swallow this message m.messages = m.messages[1:] } } func (m *mockGossipComms) GossipNeighbourSubset(update mesh.GossipData) { m.GossipBroadcast(update) } func equalByteBuffer(a, b []byte) bool { if len(a) != len(b) { return false } for i := 0; i < len(a); i++ { if a[i] != b[i] { return false } } return true } func (m *mockGossipComms) GossipUnicast(dstPeerName mesh.PeerName, buf []byte) error { m.Lock() defer m.Unlock() if len(m.messages) == 0 { require.FailNow(m, fmt.Sprintf("%s: Gossip message to %s unexpected: \n%s", m.name, dstPeerName, buf)) } else if msg := m.messages[0]; msg.dst == mesh.UnknownPeerName { require.FailNow(m, fmt.Sprintf("%s: Expected Gossip broadcast message but got dest %s", m.name, dstPeerName)) } else if msg.dst != dstPeerName { require.FailNow(m, fmt.Sprintf("%s: Expected Gossip message to %s but got dest %s", m.name, msg.dst, dstPeerName)) } else if buf[0] != msg.msgType { require.FailNow(m, fmt.Sprintf("%s: Expected Gossip message of type %d but got type %d", m.name, msg.msgType, buf[0])) } else if msg.buf != nil && !equalByteBuffer(msg.buf, buf[1:]) { require.FailNow(m, fmt.Sprintf("%s: Gossip message not sent as expected: \nwant: %x\ngot : %x", m.name, msg.buf, buf[1:])) } else { // Swallow this message m.messages = m.messages[1:] } return nil } func ExpectBroadcastMessage(alloc *Allocator, buf []byte) { m := alloc.gossip.(*mockGossipComms) m.Lock() m.messages = append(m.messages, mockMessage{mesh.UnknownPeerName, 0, buf}) m.Unlock() } func CheckAllExpectedMessagesSent(allocs ...*Allocator) { for _, alloc := range allocs { m := alloc.gossip.(*mockGossipComms) m.RLock() if len(m.messages) > 0 { require.FailNow(m, fmt.Sprintf("%s: Gossip message(s) not sent as expected: \n%v", m.name, m.messages)) } m.RUnlock() } } type mockDB struct{} func (d *mockDB) Load(_ string, _ interface{}) (bool, error) { return false, nil } func (d *mockDB) Save(_ string, _ interface{}) error { return nil } func makeAllocator(name string, cidrStr string, quorum uint, preClaims ...PreClaim) (*Allocator, address.CIDR) { peername, err := mesh.PeerNameFromString(name) if err != nil { panic(err) } cidr, err := address.ParseCIDR(cidrStr) if err != nil { panic(err) } return NewAllocator(Config{ OurName: peername, OurUID: mesh.PeerUID(rand.Int63()), OurNickname: "nick-" + name, Universe: cidr, IsObserver: quorum == 0, PreClaims: preClaims, Quorum: func() uint { return quorum }, Db: new(mockDB), IsKnownPeer: func(mesh.PeerName) bool { return true }, }), cidr } func makeAllocatorWithMockGossip(t *testing.T, name string, universeCIDR string, quorum uint) (*Allocator, address.CIDR) { alloc, subnet := makeAllocator(name, universeCIDR, quorum) gossip := &mockGossipComms{T: t, name: name} alloc.SetInterfaces(gossip) alloc.Start() return alloc, subnet } func (alloc *Allocator) claimRingForTesting(allocs ...*Allocator) { peers := []mesh.PeerName{alloc.ourName} for _, alloc2 := range allocs { peers = append(peers, alloc2.ourName) } alloc.ring.ClaimForPeers(normalizeConsensus(peers)) alloc.space.AddRanges(alloc.ring.OwnedRanges()) } func (alloc *Allocator) NumFreeAddresses(r address.Range) address.Count { resultChan := make(chan address.Count) alloc.actionChan <- func() { resultChan <- alloc.space.NumFreeAddressesInRange(r) } return <-resultChan } func (alloc *Allocator) OwnedRanges() (result []address.Range) { resultChan := make(chan []address.Range) alloc.actionChan <- func() { resultChan <- alloc.ring.OwnedRanges() } return <-resultChan } // Check whether or not something was sent on a channel func AssertSent(t *testing.T, ch <-chan bool) { timeout := time.After(10 * time.Second) select { case <-ch: // This case is ok case <-timeout: require.FailNow(t, "Nothing sent on channel") } } func AssertNothingSent(t *testing.T, ch <-chan bool) { select { case val := <-ch: require.FailNow(t, fmt.Sprintf("Unexpected value on channel: %v", val)) default: // no message received } } func makeNetworkOfAllocators(size int, cidr string, preClaims ...[]PreClaim) ([]*Allocator, *gossip.TestRouter, address.CIDR) { gossipRouter := gossip.NewTestRouter(0.0) allocs := make([]*Allocator, size) var subnet address.CIDR for i := 0; i < size; i++ { var alloc *Allocator preClaim := []PreClaim{} if i < len(preClaims) { preClaim = preClaims[i] } alloc, subnet = makeAllocator(fmt.Sprintf("%02d:00:00:02:00:00", i), cidr, uint(size/2+1), preClaim...) alloc.SetInterfaces(gossipRouter.Connect(alloc.ourName, alloc)) alloc.Start() allocs[i] = alloc } allocs[size-1].gossip.GossipBroadcast(allocs[size-1].Gossip()) gossipRouter.Flush() return allocs, gossipRouter, subnet } func stopNetworkOfAllocators(allocs []*Allocator, gossipRouter *gossip.TestRouter) { // NB: We must stop the router first since ipam gossip makes // synchronous calls into the allocator. gossipRouter.Stop() for _, alloc := range allocs { alloc.Stop() } }