From 98c35a9aa74dd60ea438183464a8c91b72e459f5 Mon Sep 17 00:00:00 2001
From: Spike Curtis <spike@coder.com>
Date: Mon, 18 Nov 2024 11:56:03 +0400
Subject: [PATCH] fix: fix listening flake on TestTailnet_ForcesWebSockets

---
 tailnet/conn_test.go | 13 +++++++++----
 1 file changed, 9 insertions(+), 4 deletions(-)

diff --git a/tailnet/conn_test.go b/tailnet/conn_test.go
index c7938afd2721f..068032e642d74 100644
--- a/tailnet/conn_test.go
+++ b/tailnet/conn_test.go
@@ -131,23 +131,28 @@ func TestTailnet(t *testing.T) {
 		stitch(t, w2, w1)
 		stitch(t, w1, w2)
 		require.True(t, w2.AwaitReachable(ctx, w1IP))
-		conn := make(chan struct{}, 1)
+		done := make(chan struct{})
+		listening := make(chan struct{})
 		go func() {
+			defer close(done)
 			listener, err := w1.Listen("tcp", ":35565")
-			assert.NoError(t, err)
+			if !assert.NoError(t, err) {
+				return
+			}
 			defer listener.Close()
+			close(listening)
 			nc, err := listener.Accept()
 			if !assert.NoError(t, err) {
 				return
 			}
 			_ = nc.Close()
-			conn <- struct{}{}
 		}()
 
+		testutil.RequireRecvCtx(ctx, t, listening)
 		nc, err := w2.DialContextTCP(ctx, netip.AddrPortFrom(w1IP, 35565))
 		require.NoError(t, err)
 		_ = nc.Close()
-		<-conn
+		testutil.RequireRecvCtx(ctx, t, done)
 
 		nodes := make(chan *tailnet.Node, 1)
 		w2.SetNodeCallback(func(node *tailnet.Node) {