From 07a5454e88e8c0f1c7d967df576aa9a9109fb609 Mon Sep 17 00:00:00 2001
From: Erik Brakkee <erik@brakkee.org>
Date: Thu, 22 Aug 2024 10:38:45 +0200
Subject: [PATCH] prepare test for more agents and clietns.

---
 pkg/server/admin/admin_test.go | 27 +++++++++++----------------
 1 file changed, 11 insertions(+), 16 deletions(-)

diff --git a/pkg/server/admin/admin_test.go b/pkg/server/admin/admin_test.go
index 5392e54..b235e99 100644
--- a/pkg/server/admin/admin_test.go
+++ b/pkg/server/admin/admin_test.go
@@ -37,13 +37,15 @@ type AdminTestSuite struct {
 	cancelFunc  context.CancelFunc
 	pprofServer *http.Server
 
-	agentReadWriter  io.ReadWriteCloser
-	serverReadWriter io.ReadWriteCloser
-
 	admin   *Admin
 	hostKey []byte
 }
 
+func (s *AdminTestSuite) createPipe() (io.ReadWriteCloser, io.ReadWriteCloser) {
+	bitpipe := testsupport.NewInmemoryConnection(s.ctx, "inmemory", 10)
+	return bitpipe.Front(), bitpipe.Back()
+}
+
 func (s *AdminTestSuite) SetupSuite() {
 	s.pprofServer = testsupport.StartPprof("")
 }
@@ -57,16 +59,6 @@ func (s *AdminTestSuite) SetupTest() {
 	s.ctx = ctx
 	s.cancelFunc = cancelFunc
 
-	// Could have also used net.Pipe but net.Pipe uses synchronous communication
-	// by default and the bitpipe implementation can become asynchronous when
-	// a channels ize > 0 is passed in. Also the test utility respects the context
-	// so also deals with cancellation much better than net.Pipe.
-	bitpipe := testsupport.NewInmemoryConnection(s.ctx, "inmemory", 10)
-	agentReadWriter := bitpipe.Front()
-	serverReadWriter := bitpipe.Back()
-	s.agentReadWriter = agentReadWriter
-	s.serverReadWriter = serverReadWriter
-
 	s.admin = NewAdmin()
 	s.hostKey = make([]byte, 100)
 	rand.Read(s.hostKey)
@@ -81,10 +73,13 @@ func TestAdminTestSuite(t *testing.T) {
 
 func (s *AdminTestSuite) Test_AgentRegisters() {
 	publicId := "abc"
+	agentRW, serverRW := s.createPipe()
 	testsupport.RunAndWait(
 		&s.Suite,
 		func() any {
-			agentConn, err := s.admin.AddAgent(s.hostKey, models.RendezVousId(publicId), comms.EnvironmentInfo{}, s.serverReadWriter)
+			agentConn, err := s.admin.AddAgent(
+				s.hostKey, models.RendezVousId(publicId), comms.EnvironmentInfo{},
+				serverRW)
 			s.Nil(err)
 			s.Equal(publicId, string(agentConn.Info.PublicId))
 			state := s.admin.CreateNotifification()
@@ -95,12 +90,12 @@ func (s *AdminTestSuite) Test_AgentRegisters() {
 		},
 		func() any {
 			// verify registration message received
-			agentRegistration, err := comms.ReceiveRegistrationMessage(s.agentReadWriter)
+			agentRegistration, err := comms.ReceiveRegistrationMessage(agentRW)
 			s.Nil(err)
 			s.True(agentRegistration.Ok)
 			s.Equal(s.hostKey, agentRegistration.HostPrivateKey)
 
-			commChannel, err := comms.NewCommChannel(comms.Agent, s.agentReadWriter)
+			commChannel, err := comms.NewCommChannel(comms.Agent, agentRW)
 			s.Nil(err)
 			s.NotNil(commChannel)
 			return nil