From b01f7514fcd7a33934b113c381418b37a7dcf15f Mon Sep 17 00:00:00 2001
From: JoukoVirtanen <jouko@stackrox.com>
Date: Wed, 13 Nov 2024 19:41:33 -0800
Subject: [PATCH] Not using assertMismatch

---
 .../pkg/mock_sensor/expect_conn.go            | 26 +++++++------------
 integration-tests/pkg/mock_sensor/server.go   |  7 +++++
 2 files changed, 17 insertions(+), 16 deletions(-)

diff --git a/integration-tests/pkg/mock_sensor/expect_conn.go b/integration-tests/pkg/mock_sensor/expect_conn.go
index 95fd111344..dd2be10ae9 100644
--- a/integration-tests/pkg/mock_sensor/expect_conn.go
+++ b/integration-tests/pkg/mock_sensor/expect_conn.go
@@ -96,26 +96,14 @@ func (s *MockSensor) checkIfConnectionsMatchExpected(t *testing.T, connections [
 	return false
 }
 
-// getConnectionsAndCompare gets the connections for a container, sorts them, and compares with a set of expected
-// connections. If asssertMismatch is true and the set of observed connections does not match the set of expected connections an assert is
-// triggered. If assertMismatch is not set then just return if the observed and expected connections match.
-func (s *MockSensor) getConnectionsAndCompare(t *testing.T, containerID string, assertMismatch bool, expected ...types.NetworkInfo) bool {
-	connections := s.Connections(containerID)
-	types.SortConnections(connections)
-	success := s.checkIfConnectionsMatchExpected(t, connections, expected)
-	if assertMismatch && !success {
-		return assert.ElementsMatch(t, expected, connections, "networking connections do not match")
-	}
-	return success
-}
-
 // ExpectSameElementsConnections compares a list of expected connections to the observed connections. This comparison is done at the beginning, when a new
 // connection arrives, and after a timeout period. The number of connections must match and the expected and observed connections must match, but the order
 // does not matter.
 func (s *MockSensor) ExpectSameElementsConnections(t *testing.T, containerID string, timeout time.Duration, expected ...types.NetworkInfo) bool {
 	types.SortConnections(expected)
 
-	success := s.getConnectionsAndCompare(t, containerID, false, expected...)
+	connections := s.SortedConnections(containerID)
+	success := s.checkIfConnectionsMatchExpected(t, connections, expected)
 	if success {
 		return true
 	}
@@ -125,12 +113,18 @@ func (s *MockSensor) ExpectSameElementsConnections(t *testing.T, containerID str
 	for {
 		select {
 		case <-timer:
-			return s.getConnectionsAndCompare(t, containerID, true, expected...)
+			connections := s.SortedConnections(containerID)
+			success := s.checkIfConnectionsMatchExpected(t, connections, expected)
+			if !success {
+				return assert.ElementsMatch(t, expected, connections, "networking connections do not match")
+			}
+			return success
 		case conn := <-s.LiveConnections():
 			if conn.GetContainerId() != containerID {
 				continue
 			}
-			success := s.getConnectionsAndCompare(t, containerID, false, expected...)
+			connections := s.SortedConnections(containerID)
+			success := s.checkIfConnectionsMatchExpected(t, connections, expected)
 			if success {
 				return true
 			}
diff --git a/integration-tests/pkg/mock_sensor/server.go b/integration-tests/pkg/mock_sensor/server.go
index f78e538db5..fc7febafc7 100644
--- a/integration-tests/pkg/mock_sensor/server.go
+++ b/integration-tests/pkg/mock_sensor/server.go
@@ -160,6 +160,13 @@ func (m *MockSensor) Connections(containerID string) []types.NetworkInfo {
 	return make([]types.NetworkInfo, 0)
 }
 
+func (m *MockSensor) SortedConnections(containerID string) []types.NetworkInfo {
+	connections := m.Connections(containerID)
+	types.SortConnections(connections)
+
+	return connections
+}
+
 // HasConnection returns whether a given connection has been seen for a given
 // container ID
 func (m *MockSensor) HasConnection(containerID string, conn types.NetworkInfo) bool {