From acd9f0441afc12e9de71403d1f448f2167c92dbb Mon Sep 17 00:00:00 2001
From: Erik Brakkee <erik@brakkee.org>
Date: Thu, 22 Aug 2024 17:13:34 +0200
Subject: [PATCH] some refactoring inpraparation for multiple agents and
 clients.

---
 pkg/server/admin/admin_test.go | 102 ++++++++++++++++++---------------
 1 file changed, 57 insertions(+), 45 deletions(-)

diff --git a/pkg/server/admin/admin_test.go b/pkg/server/admin/admin_test.go
index 3062ada..a00af65 100644
--- a/pkg/server/admin/admin_test.go
+++ b/pkg/server/admin/admin_test.go
@@ -17,18 +17,7 @@ import (
 	"time"
 )
 
-// test cases
-//
-// Agent only: verify state, verify agentregistration message
-// - Connect single agent
-// - Connect agent, connect second agent with duplicate id
-//    -> new id taken out
-// - Connect more than 100 agents with the same id
-//   -> 101th agents gets error
-//
-// Client connected to agent: Verify clientConnection and agentCOnnection, verify state, verify clientInfo message.
-// - Connect agent + connect client with mmtching id
-// - Connect agent + connect client with wrong id
+// test case
 //
 // Overall:
 // - Connect agent, connect 2 clients
@@ -171,34 +160,29 @@ func (s *AdminTestSuite) Test_connectClient() {
 	res := testsupport.RunAndWait(
 		&s.Suite,
 		func() any {
-			// server
-			clientConn, err := s.admin.AddClient(models.RendezVousId(publicId),
-				iowrappers.NewSimpleReadWriteAddrCloser(serverToClientRW, testsupport.DummyRemoteAddr("remoteaddr")))
-			s.Nil(err)
-			// Connection to agent over yamux
-			serverToAgentYamux := clientConn.agentConnection
-			// test by sending a message to the agent.
-			testsupport.AssertWriteData(&s.Suite, data, serverToAgentYamux)
-			return clientConn
+			return s.connectClient(publicId, serverToClientRW, data)
 		},
 		func() any {
-			// agent
-			listener := comms.NewAgentListener(agentRes.commChannel.Session)
-			//.Connection from server over yamux
-			agentToServerYamux, err := listener.Accept()
-			s.Nil(err)
-			// Test by receiving a message from the server
-			testsupport.AssertReadData(&s.Suite, data, agentToServerYamux)
-			return agentToServerYamux
+			return s.clientConnection(agentRes, data)
 		})
 
-	// Now we need to verify bi-directional communication between client and agent through the wserv
-
+	// bidirectional communciataion check
 	clientConn := res[0].(*clientConnection)
-	go clientConn.Synchronize()
-
 	agentToServerYamux := res[1].(io.ReadWriter)
+	go clientConn.Synchronize()
+	s.bidirectionalConnectionCheck(clientToServerRW, agentToServerYamux)
 
+	// verify state
+	state := s.admin.CreateNotifification()
+	s.Equal(1, len(state.Agents))
+	s.Equal(1, len(state.Clients))
+	s.Equal(clientConn.Info, state.Clients[clientConn.Info.Guid])
+
+	// will close the connections and as a result also th synchronize goroutine.
+	s.cancelFunc()
+}
+
+func (s *AdminTestSuite) bidirectionalConnectionCheck(clientToServerRW io.ReadWriteCloser, agentToServerYamux io.ReadWriter) {
 	data1 := "mytestdata"
 	data2 := "mytestdata-2"
 	testsupport.RunAndWait(
@@ -213,9 +197,6 @@ func (s *AdminTestSuite) Test_connectClient() {
 			testsupport.AssertWriteData(&s.Suite, data2, clientToServerRW)
 			return nil
 		})
-
-	// will close the connections and as a result also th synchronize goroutine.
-	s.cancelFunc()
 }
 
 func (s *AdminTestSuite) Test_connectClientUnknownRendezVousId() {
@@ -229,8 +210,27 @@ func (s *AdminTestSuite) Test_connectClientUnknownRendezVousId() {
 	_, err := s.admin.AddClient(models.RendezVousId(publicId+"sothatisunknown"),
 		iowrappers.NewSimpleReadWriteAddrCloser(serverToClientRW, testsupport.DummyRemoteAddr("remoteaddr")))
 	s.NotNil(err)
+
+	// verify state
+	state := s.admin.CreateNotifification()
+	s.Equal(1, len(state.Agents))
+	s.Equal(0, len(state.Clients))
 }
 
+// Registering an agent on the server
+
+func (s *AdminTestSuite) addAgent(publicId string, assignedPublicId string, serverRW io.ReadWriteCloser) (*agentConnection, error) {
+	agentConn, err := s.admin.AddAgent(
+		s.hostKey, models.RendezVousId(publicId), comms.EnvironmentInfo{},
+		serverRW)
+	if err != nil {
+		return nil, err
+	}
+	s.Equal(assignedPublicId, string(agentConn.Info.PublicId))
+	return agentConn, nil
+}
+
+// Agent activities registring on the server
 func (s *AdminTestSuite) agentRegistration(agentRW io.ReadWriteCloser) AgentRegisterResult {
 	// verify registration message received
 	agentRegistration, err := comms.ReceiveRegistrationMessage(agentRW)
@@ -245,13 +245,25 @@ func (s *AdminTestSuite) agentRegistration(agentRW io.ReadWriteCloser) AgentRegi
 	return AgentRegisterResult{registration: agentRegistration, commChannel: commChannel, err: nil}
 }
 
-func (s *AdminTestSuite) addAgent(publicId string, assignedPublicId string, serverRW io.ReadWriteCloser) (*agentConnection, error) {
-	agentConn, err := s.admin.AddAgent(
-		s.hostKey, models.RendezVousId(publicId), comms.EnvironmentInfo{},
-		serverRW)
-	if err != nil {
-		return nil, err
-	}
-	s.Equal(assignedPublicId, string(agentConn.Info.PublicId))
-	return agentConn, nil
+func (s *AdminTestSuite) connectClient(publicId string, serverToClientRW io.ReadWriteCloser, data string) any {
+	// server
+	clientConn, err := s.admin.AddClient(models.RendezVousId(publicId),
+		iowrappers.NewSimpleReadWriteAddrCloser(serverToClientRW, testsupport.DummyRemoteAddr("remoteaddr")))
+	s.Nil(err)
+	// Connection to agent over yamux
+	serverToAgentYamux := clientConn.agentConnection
+	// test by sending a message to the agent.
+	testsupport.AssertWriteData(&s.Suite, data, serverToAgentYamux)
+	return clientConn
+}
+
+func (s *AdminTestSuite) clientConnection(agentRes AgentRegisterResult, data string) any {
+	// agent
+	listener := comms.NewAgentListener(agentRes.commChannel.Session)
+	//.Connection from server over yamux
+	agentToServerYamux, err := listener.Accept()
+	s.Nil(err)
+	// Test by receiving a message from the server
+	testsupport.AssertReadData(&s.Suite, data, agentToServerYamux)
+	return agentToServerYamux
 }