Fix flaky natpmp tests

qdm12 · Jun 30, 2023 · 93c515d · 93c515d
1 parent e7d4246
commit 93c515d
Show file tree

Hide file tree

Showing 3 changed files with 21 additions and 11 deletions.
diff --git a/internal/natpmp/helpers_test.go b/internal/natpmp/helpers_test.go
@@ -75,7 +75,10 @@ func launchUDPServer(t *testing.T, exchanges []udpExchange) (
 
 			if exchange.close {
 				err = conn.Close()
-				assert.NoError(t, err)
+				if !errors.Is(err, net.ErrClosed) {
+					// connection might be already closed by client production code
+					assert.NoError(t, err)
+				}
 				return
 			}
 

diff --git a/internal/natpmp/portmapping_test.go b/internal/natpmp/portmapping_test.go
@@ -136,9 +136,13 @@ func Test_Client_AddPortMapping(t *testing.T) {
 			assert.Equal(t, testCase.assignedInternalPort, assignedInternalPort)
 			assert.Equal(t, testCase.assignedExternalPort, assignedExternalPort)
 			assert.Equal(t, testCase.assignedLifetime, assignedLifetime)
-			assert.ErrorIs(t, err, testCase.err)
-			if testCase.err != nil {
-				assert.EqualError(t, err, testCase.errMessage)
+			if testCase.errMessage != "" {
+				if testCase.err != nil {
+					assert.ErrorIs(t, err, testCase.err)
+				}
+				assert.Regexp(t, "^"+testCase.errMessage+"$", err.Error())
+			} else {
+				assert.NoError(t, err)
 			}
 		})
 	}

diff --git a/internal/natpmp/rpc_test.go b/internal/natpmp/rpc_test.go
@@ -7,7 +7,6 @@ import (
 	"time"
 
 	"github.com/stretchr/testify/assert"
-	"github.com/stretchr/testify/require"
 )
 
 func Test_Client_rpc(t *testing.T) {
@@ -51,9 +50,11 @@ func Test_Client_rpc(t *testing.T) {
 			gateway:      netip.AddrFrom4([4]byte{127, 0, 0, 1}),
 			request:      []byte{0, 1},
 			initialRetry: time.Millisecond,
-			exchanges:    []udpExchange{{close: true}},
-			err:          ErrConnectionTimeout,
-			errMessage:   "connection timeout: after 1ms",
+			exchanges: []udpExchange{
+				{request: []byte{0, 1}, close: true},
+			},
+			err:        ErrConnectionTimeout,
+			errMessage: "connection timeout: after 1ms",
 		},
 		"response_too_small": {
 			ctx:          context.Background(),
@@ -152,10 +153,12 @@ func Test_Client_rpc(t *testing.T) {
 				testCase.request, testCase.responseSize)
 
 			if testCase.errMessage != "" {
-				require.Error(t, err)
-				assert.Regexp(t, testCase.errMessage, err.Error())
+				if testCase.err != nil {
+					assert.ErrorIs(t, err, testCase.err)
+				}
+				assert.Regexp(t, "^"+testCase.errMessage+"$", err.Error())
 			} else {
-				assert.ErrorIs(t, err, testCase.err)
+				assert.NoError(t, err)
 			}
 			assert.Equal(t, testCase.expectedResponse, response)
 		})