feat: add support for WorkspaceUpdates to WebsocketDialer (#15534)

closes #14730

Adds support for WorkspaceUpdates to the WebsocketDialer. This allows us to dial the new endpoint added in #14847 and connect it up to a `tailnet.Controllers` to connect to all agents over the tailnet.

I refactored the fakeWorkspaceUpdatesProvider to a mock and moved it to `tailnettest` so it could be more easily reused.  The Mock is a little more full-featured.

codersdk/workspacesdk/dialer_test.go

@@ -9,8 +9,10 @@ import (
 	"testing"
 	"time"
 
+	"github.com/google/uuid"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
+	"go.uber.org/mock/gomock"
 	"nhooyr.io/websocket"
 	"tailscale.com/tailcfg"
 
@@ -21,7 +23,7 @@ import (
 	"github.com/coder/coder/v2/codersdk"
 	"github.com/coder/coder/v2/codersdk/workspacesdk"
 	"github.com/coder/coder/v2/tailnet"
-	"github.com/coder/coder/v2/tailnet/proto"
+	tailnetproto "github.com/coder/coder/v2/tailnet/proto"
 	"github.com/coder/coder/v2/tailnet/tailnettest"
 	"github.com/coder/coder/v2/testutil"
 )
@@ -102,6 +104,7 @@ func TestWebsocketDialer_TokenController(t *testing.T) {
 	require.Equal(t, "", gotToken)
 
 	clients = testutil.RequireRecvCtx(ctx, t, clientCh)
+	require.Nil(t, clients.WorkspaceUpdates)
 	clients.Closer.Close()
 
 	err = testutil.RequireRecvCtx(ctx, t, wsErr)
@@ -273,7 +276,7 @@ func TestWebsocketDialer_UplevelVersion(t *testing.T) {
 	logger := slogtest.Make(t, &slogtest.Options{IgnoreErrors: true}).Leveled(slog.LevelDebug)
 
 	svr := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
-		sVer := apiversion.New(proto.CurrentMajor, proto.CurrentMinor-1)
+		sVer := apiversion.New(2, 2)
 
 		// the following matches what Coderd does;
 		// c.f. coderd/workspaceagents.go: workspaceAgentClientCoordinate
@@ -291,7 +294,10 @@ func TestWebsocketDialer_UplevelVersion(t *testing.T) {
 	svrURL, err := url.Parse(svr.URL)
 	require.NoError(t, err)
 
-	uut := workspacesdk.NewWebsocketDialer(logger, svrURL, &websocket.DialOptions{})
+	uut := workspacesdk.NewWebsocketDialer(
+		logger, svrURL, &websocket.DialOptions{},
+		workspacesdk.WithWorkspaceUpdates(&tailnetproto.WorkspaceUpdatesRequest{}),
+	)
 
 	errCh := make(chan error, 1)
 	go func() {
@@ -307,6 +313,84 @@ func TestWebsocketDialer_UplevelVersion(t *testing.T) {
 	require.NotEmpty(t, sdkErr.Helper)
 }
 
+func TestWebsocketDialer_WorkspaceUpdates(t *testing.T) {
+	t.Parallel()
+	ctx := testutil.Context(t, testutil.WaitShort)
+	logger := slogtest.Make(t, &slogtest.Options{
+		IgnoreErrors: true,
+	}).Leveled(slog.LevelDebug)
+
+	fCoord := tailnettest.NewFakeCoordinator()
+	var coord tailnet.Coordinator = fCoord
+	coordPtr := atomic.Pointer[tailnet.Coordinator]{}
+	coordPtr.Store(&coord)
+	ctrl := gomock.NewController(t)
+	mProvider := tailnettest.NewMockWorkspaceUpdatesProvider(ctrl)
+
+	svc, err := tailnet.NewClientService(tailnet.ClientServiceOptions{
+		Logger:                   logger,
+		CoordPtr:                 &coordPtr,
+		DERPMapUpdateFrequency:   time.Hour,
+		DERPMapFn:                func() *tailcfg.DERPMap { return &tailcfg.DERPMap{} },
+		WorkspaceUpdatesProvider: mProvider,
+	})
+	require.NoError(t, err)
+
+	wsErr := make(chan error, 1)
+	svr := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		// need 2.3 for WorkspaceUpdates RPC
+		cVer := r.URL.Query().Get("version")
+		assert.Equal(t, "2.3", cVer)
+
+		sws, err := websocket.Accept(w, r, nil)
+		if !assert.NoError(t, err) {
+			return
+		}
+		wsCtx, nc := codersdk.WebsocketNetConn(ctx, sws, websocket.MessageBinary)
+		// streamID can be empty because we don't call RPCs in this test.
+		wsErr <- svc.ServeConnV2(wsCtx, nc, tailnet.StreamID{})
+	}))
+	defer svr.Close()
+	svrURL, err := url.Parse(svr.URL)
+	require.NoError(t, err)
+
+	userID := uuid.UUID{88}
+
+	mSub := tailnettest.NewMockSubscription(ctrl)
+	updateCh := make(chan *tailnetproto.WorkspaceUpdate, 1)
+	mProvider.EXPECT().Subscribe(gomock.Any(), userID).Times(1).Return(mSub, nil)
+	mSub.EXPECT().Updates().MinTimes(1).Return(updateCh)
+	mSub.EXPECT().Close().Times(1).Return(nil)
+
+	uut := workspacesdk.NewWebsocketDialer(
+		logger, svrURL, &websocket.DialOptions{},
+		workspacesdk.WithWorkspaceUpdates(&tailnetproto.WorkspaceUpdatesRequest{
+			WorkspaceOwnerId: userID[:],
+		}),
+	)
+
+	clients, err := uut.Dial(ctx, nil)
+	require.NoError(t, err)
+	require.NotNil(t, clients.WorkspaceUpdates)
+
+	wsID := uuid.UUID{99}
+	expectedUpdate := &tailnetproto.WorkspaceUpdate{
+		UpsertedWorkspaces: []*tailnetproto.Workspace{
+			{Id: wsID[:]},
+		},
+	}
+	updateCh <- expectedUpdate
+
+	gotUpdate, err := clients.WorkspaceUpdates.Recv()
+	require.NoError(t, err)
+	require.Equal(t, wsID[:], gotUpdate.GetUpsertedWorkspaces()[0].GetId())
+
+	clients.Closer.Close()
+
+	err = testutil.RequireRecvCtx(ctx, t, wsErr)
+	require.NoError(t, err)
+}
+
 type fakeResumeTokenController struct {
 	ctx        context.Context
 	t          testing.TB