github.com/choria-io/go-choria@v0.28.1-0.20240416190746-b3bf9c7d5a45/providers/discovery/broadcast/broadcast_test.go (about)

     1  // Copyright (c) 2021-2022, R.I. Pienaar and the Choria Project contributors
     2  //
     3  // SPDX-License-Identifier: Apache-2.0
     4  
     5  package broadcast
     6  
     7  import (
     8  	"context"
     9  	"fmt"
    10  	"sort"
    11  	"testing"
    12  	"time"
    13  
    14  	"github.com/choria-io/go-choria/build"
    15  	"github.com/choria-io/go-choria/client/client"
    16  	"github.com/choria-io/go-choria/inter"
    17  	imock "github.com/choria-io/go-choria/inter/imocks"
    18  	"github.com/choria-io/go-choria/message"
    19  	"github.com/choria-io/go-choria/protocol"
    20  	v1 "github.com/choria-io/go-choria/protocol/v1"
    21  	"github.com/choria-io/go-choria/providers/security/filesec"
    22  
    23  	"github.com/golang/mock/gomock"
    24  
    25  	"github.com/choria-io/go-choria/config"
    26  
    27  	. "github.com/onsi/ginkgo/v2"
    28  	. "github.com/onsi/gomega"
    29  )
    30  
    31  func TestBroadcast(t *testing.T) {
    32  	RegisterFailHandler(Fail)
    33  	RunSpecs(t, "Providers/Discovery/Broadcast")
    34  }
    35  
    36  var _ = Describe("Broadcast", func() {
    37  	var (
    38  		fw      *imock.MockFramework
    39  		cfg     *config.Config
    40  		mockctl *gomock.Controller
    41  		cl      *MockChoriaClient
    42  		b       *Broadcast
    43  	)
    44  
    45  	BeforeEach(func() {
    46  		mockctl = gomock.NewController(GinkgoT())
    47  		cl = NewMockChoriaClient(mockctl)
    48  		fw, cfg = imock.NewFrameworkForTests(mockctl, GinkgoWriter, imock.WithCallerID())
    49  		cfg.Collectives = []string{"mcollective", "test"}
    50  
    51  		fw.EXPECT().NewMessage(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).DoAndReturn(func(payload []byte, agent string, collective string, msgType string, request inter.Message) (msg inter.Message, err error) {
    52  			return message.NewMessage(payload, agent, collective, msgType, request, fw)
    53  		}).AnyTimes()
    54  
    55  		sec, err := filesec.New(filesec.WithChoriaConfig(&build.Info{}, cfg), filesec.WithLog(fw.Logger("")))
    56  		Expect(err).ToNot(HaveOccurred())
    57  
    58  		fw.EXPECT().NewTransportFromJSON(gomock.Any()).DoAndReturn(func(data []byte) (message protocol.TransportMessage, err error) {
    59  			return v1.NewTransportFromJSON(data)
    60  		}).AnyTimes()
    61  		fw.EXPECT().NewReplyTransportForMessage(gomock.Any(), gomock.Any()).DoAndReturn(func(msg inter.Message, request protocol.Request) (protocol.TransportMessage, error) {
    62  			reply, err := v1.NewReply(request, cfg.Identity)
    63  			Expect(err).ToNot(HaveOccurred())
    64  			reply.SetMessage(msg.Payload())
    65  
    66  			sreply, err := v1.NewSecureReply(reply, sec)
    67  			Expect(err).ToNot(HaveOccurred())
    68  
    69  			transport, err := v1.NewTransportMessage(cfg.Identity)
    70  			Expect(err).ToNot(HaveOccurred())
    71  
    72  			err = transport.SetReplyData(sreply)
    73  			Expect(err).ToNot(HaveOccurred())
    74  
    75  			return transport, nil
    76  		}).AnyTimes()
    77  
    78  		b = New(fw)
    79  	})
    80  
    81  	AfterEach(func() {
    82  		mockctl.Finish()
    83  	})
    84  
    85  	Describe("New", func() {
    86  		It("Should initialize timeout to default", func() {
    87  			Expect(b.timeout).To(Equal(2 * time.Second))
    88  			cfg.DiscoveryTimeout = 100
    89  			b = New(fw)
    90  			Expect(b.timeout).To(Equal(100 * time.Second))
    91  		})
    92  	})
    93  
    94  	Describe("Discover", func() {
    95  		It("Should request and return discovered nodes", func() {
    96  			ctx, cancel := context.WithCancel(context.Background())
    97  			defer cancel()
    98  
    99  			f := protocol.NewFilter()
   100  			f.AddAgentFilter("choria")
   101  
   102  			cl.EXPECT().Request(gomock.Any(), gomock.Any(), gomock.Any()).Return(nil).Do(func(ctx context.Context, msg *message.Message, handler client.Handler) {
   103  				Expect(msg.Collective()).To(Equal("test"))
   104  				Expect(msg.Payload()).To(Equal([]byte("cGluZw==")))
   105  
   106  				req, err := v1.NewRequest(msg.Agent(), msg.SenderID(), msg.CallerID(), msg.TTL(), msg.RequestID(), msg.Collective())
   107  				Expect(err).ToNot(HaveOccurred())
   108  				req.SetMessage(msg.Payload())
   109  
   110  				reply, err := message.NewMessageFromRequest(req, msg.ReplyTo(), fw)
   111  				Expect(err).ToNot(HaveOccurred())
   112  
   113  				t, err := reply.Transport(context.Background())
   114  				Expect(err).ToNot(HaveOccurred())
   115  
   116  				for i := 0; i < 10; i++ {
   117  					t.SetSender(fmt.Sprintf("test.sender.%d", i))
   118  
   119  					j, err := t.JSON()
   120  					Expect(err).ToNot(HaveOccurred())
   121  
   122  					cm := imock.NewMockConnectorMessage(mockctl)
   123  					cm.EXPECT().Data().Return([]byte(j))
   124  
   125  					handler(ctx, cm)
   126  				}
   127  			})
   128  
   129  			nodes, err := b.Discover(ctx, choriaClient(cl), Filter(f), Collective("test"))
   130  			Expect(err).ToNot(HaveOccurred())
   131  			sort.Strings(nodes)
   132  			Expect(nodes).To(Equal([]string{"test.sender.0", "test.sender.1", "test.sender.2", "test.sender.3", "test.sender.4", "test.sender.5", "test.sender.6", "test.sender.7", "test.sender.8", "test.sender.9"}))
   133  		})
   134  	})
   135  })