From 0de5a45b4fdc9b68460acdea2edeffe8c50777e0 Mon Sep 17 00:00:00 2001
From: Cian Johnston <cian@coder.com>
Date: Thu, 24 Jul 2025 20:30:15 +0000
Subject: [PATCH] chore(enterprise/coderd/coderdenttest): ensure WorkspaceProxy
 is connected before returning

---
 enterprise/coderd/coderdenttest/proxytest.go | 18 +++++++++++++++++-
 1 file changed, 17 insertions(+), 1 deletion(-)

diff --git a/enterprise/coderd/coderdenttest/proxytest.go b/enterprise/coderd/coderdenttest/proxytest.go
index 5aaaf4a88a725..7fb8082166e8e 100644
--- a/enterprise/coderd/coderdenttest/proxytest.go
+++ b/enterprise/coderd/coderdenttest/proxytest.go
@@ -68,7 +68,6 @@ func NewWorkspaceProxyReplica(t *testing.T, coderdAPI *coderd.API, owner *coders
 	t.Helper()
 
 	ctx, cancelFunc := context.WithCancel(context.Background())
-	t.Cleanup(cancelFunc)
 
 	if options == nil {
 		options = &ProxyOptions{}
@@ -179,6 +178,23 @@ func NewWorkspaceProxyReplica(t *testing.T, coderdAPI *coderd.API, owner *coders
 	handler = wssrv.Handler
 	mutex.Unlock()
 
+	// Wait for the tailnet connection to be established before returning.
+	// Returning the proxy before it's connected may lead to test ending
+	// before the proxy is fully connected, which will result in an error log
+	// that may fail the test.
+	waitCtx, waitCancel := context.WithTimeout(ctx, testutil.WaitShort)
+	defer waitCancel()
+	dialer, err := wssrv.SDKClient.TailnetDialer()
+	require.NoError(t, err, "create test tailnet dialer")
+
+	select {
+	case <-dialer.Connected():
+		t.Logf("Test dialer connected")
+	case <-waitCtx.Done():
+		t.Logf("Waiting for tailnet connection timed out or was canceled: %v", waitCtx.Err())
+	}
+	t.Cleanup(cancelFunc)
+
 	return WorkspaceProxy{
 		Server:    wssrv,
 		ServerURL: serverURL,