chore(natpmp): fix determinism for test Test_Client_ExternalAddress

qdm12 · Dec 14, 2024 · ddd9f4d · ddd9f4d
1 parent 7e58b4b
commit ddd9f4d
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 4 deletions.
diff --git a/internal/natpmp/externaladdress_test.go b/internal/natpmp/externaladdress_test.go
@@ -2,6 +2,7 @@ package natpmp
 
 import (
 	"context"
+	"net"
 	"net/netip"
 	"testing"
 	"time"
@@ -23,14 +24,15 @@ func Test_Client_ExternalAddress(t *testing.T) {
 		durationSinceStartOfEpoch time.Duration
 		externalIPv4Address       netip.Addr
 		err                       error
-		errMessage                string
+		errMessageRegex           string
 	}{
 		"failure": {
 			ctx:                 canceledCtx,
 			gateway:             netip.AddrFrom4([4]byte{127, 0, 0, 1}),
 			initialConnDuration: initialConnectionDuration,
-			err:                 context.Canceled,
-			errMessage:          "executing remote procedure call: reading from udp connection: context canceled",
+			err:                 net.ErrClosed,
+			errMessageRegex: "executing remote procedure call: setting connection deadline: " +
+				"set udp 127.0.0.1:[1-9][0-9]{1,4}: use of closed network connection",
 		},
 		"success": {
 			ctx:                 context.Background(),
@@ -60,7 +62,7 @@ func Test_Client_ExternalAddress(t *testing.T) {
 			durationSinceStartOfEpoch, externalIPv4Address, err := client.ExternalAddress(testCase.ctx, testCase.gateway)
 			assert.ErrorIs(t, err, testCase.err)
 			if testCase.err != nil {
-				assert.EqualError(t, err, testCase.errMessage)
+				assert.Regexp(t, testCase.errMessageRegex, err.Error())
 			}
 			assert.Equal(t, testCase.durationSinceStartOfEpoch, durationSinceStartOfEpoch)
 			assert.Equal(t, testCase.externalIPv4Address, externalIPv4Address)

diff --git a/internal/natpmp/rpc.go b/internal/natpmp/rpc.go
@@ -45,8 +45,10 @@ func (c *Client) rpc(ctx context.Context, gateway netip.Addr,
 		cancel()
 		<-endGoroutineDone
 	}()
+	ctxListeningReady := make(chan struct{})
 	go func() {
 		defer close(endGoroutineDone)
+		close(ctxListeningReady)
 		// Context is canceled either by the parent context or
 		// when this function returns.
 		<-ctx.Done()
@@ -60,6 +62,7 @@ func (c *Client) rpc(ctx context.Context, gateway netip.Addr,
 		}
 		err = fmt.Errorf("%w; closing connection: %w", err, closeErr)
 	}()
+	<-ctxListeningReady // really to make unit testing reliable
 
 	const maxResponseSize = 16
 	response = make([]byte, maxResponseSize)