11package conntrack
22
33import (
4- "net"
4+ "net/netip "
55 "testing"
66
7- "github.com/google/go-cmp/cmp"
87 "github.com/stretchr/testify/assert"
98 "github.com/stretchr/testify/require"
109 "github.com/ti-mo/netfilter"
@@ -183,8 +182,8 @@ var corpusExpect = []struct {
183182 exp : Expect {
184183 TupleMaster : Tuple {
185184 IP : IPTuple {
186- SourceAddress : [] byte { 127 , 0 , 0 , 1 } ,
187- DestinationAddress : [] byte { 127 , 0 , 0 , 2 } ,
185+ SourceAddress : netip . MustParseAddr ( " 127.0.0.1" ) ,
186+ DestinationAddress : netip . MustParseAddr ( " 127.0.0.2" ) ,
188187 },
189188 Proto : ProtoTuple {
190189 Protocol : 6 ,
@@ -194,8 +193,8 @@ var corpusExpect = []struct {
194193 },
195194 Tuple : Tuple {
196195 IP : IPTuple {
197- SourceAddress : [] byte { 127 , 0 , 0 , 1 } ,
198- DestinationAddress : [] byte { 127 , 0 , 0 , 2 } ,
196+ SourceAddress : netip . MustParseAddr ( " 127.0.0.1" ) ,
197+ DestinationAddress : netip . MustParseAddr ( " 127.0.0.2" ) ,
199198 },
200199 Proto : ProtoTuple {
201200 Protocol : 6 ,
@@ -204,8 +203,8 @@ var corpusExpect = []struct {
204203 },
205204 Mask : Tuple {
206205 IP : IPTuple {
207- SourceAddress : [] byte { 255 , 255 , 255 , 255 } ,
208- DestinationAddress : [] byte { 255 , 255 , 255 , 255 } ,
206+ SourceAddress : netip . MustParseAddr ( " 255. 255. 255. 255" ) ,
207+ DestinationAddress : netip . MustParseAddr ( " 255. 255. 255. 255" ) ,
209208 },
210209 Proto : ProtoTuple {
211210 Protocol : 6 ,
@@ -263,11 +262,8 @@ func TestExpectUnmarshal(t *testing.T) {
263262 for _ , tt := range corpusExpect {
264263 t .Run (tt .name , func (t * testing.T ) {
265264 var ex Expect
266- assert .NoError (t , ex .unmarshal (mustDecodeAttributes (tt .attrs )))
267-
268- if diff := cmp .Diff (tt .exp , ex ); diff != "" {
269- t .Fatalf ("unexpected unmarshal (-want +got):\n %s" , diff )
270- }
265+ require .NoError (t , ex .unmarshal (mustDecodeAttributes (tt .attrs )))
266+ assert .Equal (t , tt .exp , ex , "unexpected unmarshal" )
271267 })
272268 }
273269
@@ -355,9 +351,7 @@ func TestExpectMarshal(t *testing.T) {
355351 },
356352 }
357353
358- if diff := cmp .Diff (want , exm ); diff != "" {
359- t .Fatalf ("unexpected Expect marshal (-want +got):\n %s" , diff )
360- }
354+ assert .Equal (t , want , exm , "unexpected Expect marshal" )
361355
362356 // Cannot marshal without tuple/mask/master Tuples
363357 _ , err = Expect {}.marshal ()
@@ -424,10 +418,7 @@ func TestExpectNATUnmarshal(t *testing.T) {
424418 }
425419
426420 require .NoError (t , err )
427-
428- if diff := cmp .Diff (tt .enat , enat ); diff != "" {
429- t .Fatalf ("unexpected unmarshal (-want +got):\n %s" , diff )
430- }
421+ assert .Equal (t , tt .enat , enat , "unexpected unmarshal" )
431422 })
432423 }
433424}
@@ -439,8 +430,8 @@ func TestExpectNATMarshal(t *testing.T) {
439430 Direction : true ,
440431 Tuple : Tuple {
441432 IP : IPTuple {
442- SourceAddress : net . ParseIP ("baa:baa::b" ),
443- DestinationAddress : net . ParseIP ("ef00:3f00::ba13" ),
433+ SourceAddress : netip . MustParseAddr ("baa:baa::b" ),
434+ DestinationAddress : netip . MustParseAddr ("ef00:3f00::ba13" ),
444435 },
445436 Proto : ProtoTuple {
446437 Protocol : 13 ,
@@ -458,9 +449,7 @@ func TestExpectNATMarshal(t *testing.T) {
458449
459450 // Only verify first attribute (direction); Tuple marshal has its own tests
460451 want := netfilter.Attribute {Type : uint16 (ctaExpectNATDir ), Data : []byte {0 , 0 , 0 , 1 }}
461- if diff := cmp .Diff (want , enm .Children [0 ]); diff != "" {
462- t .Fatalf ("unexpected ExpectNAT marshal (-want +got):\n %s" , diff )
463- }
452+ assert .Equal (t , want , enm .Children [0 ], "unexpected ExpectNAT marshal" )
464453}
465454
466455func TestExpectTypeString (t * testing.T ) {
0 commit comments