Function TestWorkspaceAgentRPCRole

coderd/workspaceagentsrpc_test.go:173–252

(t *testing.T)

}

func TestWorkspaceAgentRPCRole(t *testing.T) {
	t.Parallel()

	t.Run("AgentRoleMonitorsConnection", func(t *testing.T) {
		t.Parallel()
		ctx := testutil.Context(t, testutil.WaitLong)
		client, db := coderdtest.NewWithDatabase(t, nil)
		user := coderdtest.CreateFirstUser(t, client)
		r := dbfake.WorkspaceBuild(t, db, database.WorkspaceTable{
			OrganizationID: user.OrganizationID,
			OwnerID:        user.UserID,
		}).WithAgent().Do()

		// Connect with role=agent using ConnectRPCWithRole. This is
		// how the real workspace agent connects.
		ac := agentsdk.New(client.URL, agentsdk.WithFixedToken(r.AgentToken))
		conn, err := ac.ConnectRPCWithRole(ctx, "agent")
		require.NoError(t, err)
		defer func() {
			_ = conn.Close()
		}()

		// The connection monitor updates the database asynchronously,
		// so we need to wait for first_connected_at to be set.
		var agent database.WorkspaceAgent
		require.Eventually(t, func() bool {
			agent, err = db.GetWorkspaceAgentByID(dbauthz.AsSystemRestricted(ctx), r.Agents[0].ID)
			if err != nil {
				return false
			}
			return agent.FirstConnectedAt.Valid
		}, testutil.WaitShort, testutil.IntervalFast)
		assert.True(t, agent.LastConnectedAt.Valid,
			"last_connected_at should be set for agent role")
	})

	t.Run("NonAgentRoleSkipsMonitoring", func(t *testing.T) {
		t.Parallel()
		ctx := testutil.Context(t, testutil.WaitLong)
		client, db := coderdtest.NewWithDatabase(t, nil)
		user := coderdtest.CreateFirstUser(t, client)
		r := dbfake.WorkspaceBuild(t, db, database.WorkspaceTable{
			OrganizationID: user.OrganizationID,
			OwnerID:        user.UserID,
		}).WithAgent().Do()

		// Connect with a non-agent role using ConnectRPCWithRole.
		// This is how coder-logstream-kube should connect.
		ac := agentsdk.New(client.URL, agentsdk.WithFixedToken(r.AgentToken))
		conn, err := ac.ConnectRPCWithRole(ctx, "logstream-kube")
		require.NoError(t, err)

		// Send a log to confirm the RPC connection is functional.
		agentAPI := agentproto.NewDRPCAgentClient(conn)
		_, err = agentAPI.BatchCreateLogs(ctx, &agentproto.BatchCreateLogsRequest{
			LogSourceId: []byte{0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0},
		})
		// We don't care about the log source error, just that the