extern/netbird
1
0
Fork 0
mirror of https://github.com/netbirdio/netbird.git synced 2024-12-24 15:48:52 +01:00
Code Issues Packages Projects Releases Wiki Activity

test: fix management multiple concurrent peers test

Browse Source
This commit is contained in:
braginini 2021-07-22 15:53:15 +02:00
parent 2172d6f1b9
commit 83ac774264
1 changed files with 17 additions and 14 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

31
management/management_test.go
View File

@ -5,7 +5,6 @@ import (
pb "github.com/golang/protobuf/proto" //nolint
log "github.com/sirupsen/logrus"
"github.com/wiretrustee/wiretrustee/encryption"
"io"
"io/ioutil"
"math/rand"
"net"
@ -34,7 +33,6 @@ var _ = Describe("Management service", func() {
var (
addr string
server *grpc.Server
tmpDir string
dataDir string
client mgmtProto.ManagementServiceClient
serverPubKey wgtypes.Key
@ -67,7 +65,7 @@ var _ = Describe("Management service", func() {
server.Stop()
err := conn.Close()
Expect(err).NotTo(HaveOccurred())
err = os.RemoveAll(tmpDir)
err = os.RemoveAll(dataDir)
Expect(err).NotTo(HaveOccurred())
})
@ -237,29 +235,30 @@ var _ = Describe("Management service", func() {
wg := sync2.WaitGroup{}
wg.Add(initialPeers + initialPeers*additionalPeers)
var clients []mgmtProto.ManagementService_SyncClient
for _, peer := range peers {
messageBytes, err := pb.Marshal(&mgmtProto.SyncRequest{})
Expect(err).NotTo(HaveOccurred())
encryptedBytes, err := encryption.Encrypt(messageBytes, serverPubKey, peer)
Expect(err).NotTo(HaveOccurred())
// open stream
sync, err := client.Sync(context.TODO(), &mgmtProto.EncryptedMessage{
WgPubKey: peer.PublicKey().String(),
Body: encryptedBytes,
})
Expect(err).NotTo(HaveOccurred())
clients = append(clients, sync)
// receive stream
peer := peer
go func() {
// open stream
sync, err := client.Sync(context.TODO(), &mgmtProto.EncryptedMessage{
WgPubKey: peer.PublicKey().String(),
Body: encryptedBytes,
})
Expect(err).NotTo(HaveOccurred())
for {
encryptedResponse := &mgmtProto.EncryptedMessage{}
err = sync.RecvMsg(encryptedResponse)
if err == io.EOF {
if err != nil {
break
} else if err != nil {
Expect(err).NotTo(HaveOccurred())
}
decryptedBytes, err := encryption.Decrypt(encryptedResponse.Body, serverPubKey, peer)
Expect(err).NotTo(HaveOccurred())
@ -268,7 +267,6 @@ var _ = Describe("Management service", func() {
err = pb.Unmarshal(decryptedBytes, resp)
Expect(err).NotTo(HaveOccurred())
wg.Done()
}
}()
}
@ -284,6 +282,11 @@ var _ = Describe("Management service", func() {
wg.Wait()
for _, syncClient := range clients {
err := syncClient.CloseSend()
Expect(err).NotTo(HaveOccurred())
}
})
})
})