Function TestWorkspaceAgent_UpdatedDERP

coderd/workspaceagents_test.go:2762–2856

TestWorkspaceAgent_UpdatedDERP runs a real coderd server, with a real agent and a real client, and updates the DERP map live to ensure connections still work.

(t *testing.T)

// and a real client, and updates the DERP map live to ensure connections still
// work.
func TestWorkspaceAgent_UpdatedDERP(t *testing.T) {
	t.Parallel()

	logger := testutil.Logger(t)

	dv := coderdtest.DeploymentValues(t)
	err := dv.DERP.Config.BlockDirect.Set("true")
	require.NoError(t, err)

	client, closer, api := coderdtest.NewWithAPI(t, &coderdtest.Options{
		DeploymentValues: dv,
	})
	defer closer.Close()
	user := coderdtest.CreateFirstUser(t, client)

	// Change the DERP mapper to our custom one.
	var currentDerpMap atomic.Pointer[tailcfg.DERPMap]
	originalDerpMap, _ := tailnettest.RunDERPAndSTUN(t)
	currentDerpMap.Store(originalDerpMap)
	derpMapFn := func(_ *tailcfg.DERPMap) *tailcfg.DERPMap {
		return currentDerpMap.Load().Clone()
	}
	api.DERPMapper.Store(&derpMapFn)

	// Start workspace a workspace agent.
	r := dbfake.WorkspaceBuild(t, api.Database, database.WorkspaceTable{
		OrganizationID: user.OrganizationID,
		OwnerID:        user.UserID,
	}).WithAgent().Do()

	agentCloser := agenttest.New(t, client.URL, r.AgentToken)
	resources := coderdtest.AwaitWorkspaceAgents(t, client, r.Workspace.ID)
	agentID := resources[0].Agents[0].ID

	// Connect from a client.
	conn1, err := func() (workspacesdk.AgentConn, error) {
		ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong)
		defer cancel() // Connection should remain open even if the dial context is canceled.

		return workspacesdk.New(client).
			DialAgent(ctx, agentID, &workspacesdk.DialAgentOptions{
				Logger: logger.Named("client1"),
			})
	}()
	require.NoError(t, err)
	defer conn1.Close()

	ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong)
	defer cancel()

	ok := conn1.AwaitReachable(ctx)
	require.True(t, ok)

	// Change the DERP map and change the region ID.
	newDerpMap, _ := tailnettest.RunDERPAndSTUN(t)
	require.NotNil(t, newDerpMap)
	newDerpMap.Regions[2] = newDerpMap.Regions[1]
	delete(newDerpMap.Regions, 1)