Function TestWebsocketDialer_TokenController

codersdk/workspacesdk/dialer_test.go:31–112

(t *testing.T)

)

func TestWebsocketDialer_TokenController(t *testing.T) {
	t.Parallel()
	ctx := testutil.Context(t, testutil.WaitMedium)
	logger := slogtest.Make(t, &slogtest.Options{
		IgnoreErrors: true,
	}).Leveled(slog.LevelDebug)

	fTokenProv := newFakeTokenController(ctx, t)
	fCoord := tailnettest.NewFakeCoordinator()
	var coord tailnet.Coordinator = fCoord
	coordPtr := atomic.Pointer[tailnet.Coordinator]{}
	coordPtr.Store(&coord)

	svc, err := tailnet.NewClientService(tailnet.ClientServiceOptions{
		Logger:                 logger,
		CoordPtr:               &coordPtr,
		DERPMapUpdateFrequency: time.Hour,
		DERPMapFn:              func() *tailcfg.DERPMap { return &tailcfg.DERPMap{} },
	})
	require.NoError(t, err)

	dialTokens := make(chan string, 1)
	wsErr := make(chan error, 1)
	svr := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		select {
		case <-ctx.Done():
			t.Error("timed out sending token")
		case dialTokens <- r.URL.Query().Get("resume_token"):
			// OK
		}

		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)

	uut := workspacesdk.NewWebsocketDialer(logger, svrURL, &websocket.DialOptions{})

	clientCh := make(chan tailnet.ControlProtocolClients, 1)
	go func() {
		clients, err := uut.Dial(ctx, fTokenProv)
		assert.NoError(t, err)
		clientCh <- clients
	}()

	call := testutil.TryReceive(ctx, t, fTokenProv.tokenCalls)
	call <- tokenResponse{"test token", true}
	gotToken := <-dialTokens
	require.Equal(t, "test token", gotToken)

	clients := testutil.TryReceive(ctx, t, clientCh)