diff --git a/pkg/connectionset/statefulness_test.go b/pkg/connectionset/statefulness_test.go
index c10b3ba..f3608f5 100644
--- a/pkg/connectionset/statefulness_test.go
+++ b/pkg/connectionset/statefulness_test.go
@@ -33,15 +33,15 @@ func allButTCP() *ConnectionSet {
 
 type statefulnessTest struct {
 	name     string
-	srcToDst *connectionset.ConnectionSet
-	dstToSrc *connectionset.ConnectionSet
+	srcToDst *ConnectionSet
+	dstToSrc *ConnectionSet
 	// expectedIsStateful represents the expected IsStateful computed value for srcToDst,
 	// which should be either StatefulTrue or StatefulFalse, given the input dstToSrc connection.
 	// the computation applies only to the TCP protocol within those connections.
 	expectedIsStateful int
 	// expectedStatefulConn represents the subset from srcToDst which is not related to the "non-stateful" mark (*) on the srcToDst connection,
 	// the stateless part for TCP is srcToDst.Subtract(statefuleConn)
-	expectedStatefulConn *connectionset.ConnectionSet
+	expectedStatefulConn *ConnectionSet
 }
 
 var testCasesStatefulness = []statefulnessTest{
@@ -54,7 +54,7 @@ var testCasesStatefulness = []statefulnessTest{
 	},
 	{
 		name:     "first_all_cons_second_tcp_with_ports",
-		srcToDst: connectionset.NewConnectionSet(true), // all connections
+		srcToDst: NewConnectionSet(true), // all connections
 		dstToSrc: newTCPConn(80, 80, MinPort, MaxPort), // TCP , src-ports: 80, dst-ports: all
 
 		// there is a subset of the tcp connection which is not stateful