From 6e9716292301ea599a1b278abd674b8a59ee1110 Mon Sep 17 00:00:00 2001
From: Erik Brakkee <erik@brakkee.org>
Date: Thu, 22 Aug 2024 17:01:04 +0200
Subject: [PATCH] using AssertReadData/AssertWriteData now where possible.

---
 pkg/comms/agentserver_test.go              | 24 ++++------------------
 pkg/testsupport/inmemoryconnection_test.go | 10 ++-------
 2 files changed, 6 insertions(+), 28 deletions(-)

diff --git a/pkg/comms/agentserver_test.go b/pkg/comms/agentserver_test.go
index 45a77d9..a08f0e2 100644
--- a/pkg/comms/agentserver_test.go
+++ b/pkg/comms/agentserver_test.go
@@ -130,32 +130,16 @@ func (s *AgentServerTestSuite) Test_ConnectThroughYamux() {
 		func() any {
 			conn, err := agentCommChannel.Session.OpenStream()
 			s.Nil(err)
-			n, err := conn.Write([]byte(dataAgentToServer))
-			s.Nil(err)
-			s.Equal(len(dataAgentToServer), n)
-
-			buf := make([]byte, len(dataServerToAgent))
-			n, err = conn.Read(buf)
-			s.Nil(err)
-			s.Equal(len(dataServerToAgent), n)
-			s.Equal([]byte(dataServerToAgent), buf[:n])
-
+			testsupport.AssertWriteData(&s.Suite, dataAgentToServer, conn)
+			testsupport.AssertReadData(&s.Suite, dataServerToAgent, conn)
 			return nil
 		},
 		func() any {
 			conn, err := serverCommChannel.Session.Accept()
 
 			s.Nil(err)
-			buf := make([]byte, len(dataAgentToServer))
-			n, err := conn.Read(buf)
-			s.Nil(err)
-			s.Equal(len(dataAgentToServer), n)
-			s.Equal([]byte(dataAgentToServer), buf[:n])
-
-			n, err = conn.Write([]byte(dataServerToAgent))
-			s.Nil(err)
-			s.Equal(len(dataServerToAgent), n)
-
+			testsupport.AssertReadData(&s.Suite, dataAgentToServer, conn)
+			testsupport.AssertWriteData(&s.Suite, dataServerToAgent, conn)
 			return nil
 		})
 }
diff --git a/pkg/testsupport/inmemoryconnection_test.go b/pkg/testsupport/inmemoryconnection_test.go
index 40c3991..198e971 100644
--- a/pkg/testsupport/inmemoryconnection_test.go
+++ b/pkg/testsupport/inmemoryconnection_test.go
@@ -49,14 +49,8 @@ func (s *InMemoryTestSuite) TearDownTest() {
 
 func (s *InMemoryTestSuite) SendReceive(a io.ReadWriteCloser, b io.ReadWriteCloser) {
 	data := "hello"
-	n, err := a.Write([]byte(data))
-	s.Nil(err)
-	s.Equal(len(data), n)
-
-	buf := make([]byte, len(data)*2)
-	n, err = b.Read(buf)
-	s.Nil(err)
-	s.Equal([]byte(data), buf[:n])
+	AssertWriteData(&s.Suite, data, a)
+	AssertReadData(&s.Suite, data, b)
 }
 
 func (s *InMemoryTestSuite) Test_SendReceive1() {