feat: use Agent v2 API for Service Banner (#11806)

Agent uses the v2 API for the service banner, rather than the v1 HTTP API.

One of several for #10534
This commit is contained in:
Spike Curtis
2024-01-30 07:44:47 +04:00
committed by GitHub
parent 4f5a2f0a9b
commit 13e24f21e4
11 changed files with 274 additions and 133 deletions

View File

@ -5,12 +5,12 @@ import (
"sync/atomic"
"testing"
"github.com/stretchr/testify/require"
"golang.org/x/xerrors"
agentproto "github.com/coder/coder/v2/agent/proto"
"github.com/coder/coder/v2/coderd/appearance"
"github.com/coder/coder/v2/codersdk"
"github.com/stretchr/testify/require"
)
func TestGetServiceBanner(t *testing.T) {

View File

@ -915,23 +915,67 @@ func AwaitWorkspaceBuildJobCompleted(t testing.TB, client *codersdk.Client, buil
// AwaitWorkspaceAgents waits for all resources with agents to be connected. If
// specific agents are provided, it will wait for those agents to be connected
// but will not fail if other agents are not connected.
//
// Deprecated: Use NewWorkspaceAgentWaiter
func AwaitWorkspaceAgents(t testing.TB, client *codersdk.Client, workspaceID uuid.UUID, agentNames ...string) []codersdk.WorkspaceResource {
t.Helper()
return NewWorkspaceAgentWaiter(t, client, workspaceID).AgentNames(agentNames).Wait()
}
agentNamesMap := make(map[string]struct{}, len(agentNames))
for _, name := range agentNames {
// WorkspaceAgentWaiter waits for all resources with agents to be connected. If
// specific agents are provided using AgentNames(), it will wait for those agents
// to be connected but will not fail if other agents are not connected.
type WorkspaceAgentWaiter struct {
t testing.TB
client *codersdk.Client
workspaceID uuid.UUID
agentNames []string
resourcesMatcher func([]codersdk.WorkspaceResource) bool
}
// NewWorkspaceAgentWaiter returns an object that waits for agents to connect when
// you call Wait() on it.
func NewWorkspaceAgentWaiter(t testing.TB, client *codersdk.Client, workspaceID uuid.UUID) WorkspaceAgentWaiter {
return WorkspaceAgentWaiter{
t: t,
client: client,
workspaceID: workspaceID,
}
}
// AgentNames instructs the waiter to wait for the given, named agents to be connected and will
// return even if other agents are not connected.
func (w WorkspaceAgentWaiter) AgentNames(names []string) WorkspaceAgentWaiter {
//nolint: revive // returns modified struct
w.agentNames = names
return w
}
// MatchResources instructs the waiter to wait until the workspace has resources that cause the
// provided matcher function to return true.
func (w WorkspaceAgentWaiter) MatchResources(m func([]codersdk.WorkspaceResource) bool) WorkspaceAgentWaiter {
//nolint: revive // returns modified struct
w.resourcesMatcher = m
return w
}
// Wait waits for the agent(s) to connect and fails the test if they do not within testutil.WaitLong
func (w WorkspaceAgentWaiter) Wait() []codersdk.WorkspaceResource {
w.t.Helper()
agentNamesMap := make(map[string]struct{}, len(w.agentNames))
for _, name := range w.agentNames {
agentNamesMap[name] = struct{}{}
}
ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong)
defer cancel()
t.Logf("waiting for workspace agents (workspace %s)", workspaceID)
w.t.Logf("waiting for workspace agents (workspace %s)", w.workspaceID)
var resources []codersdk.WorkspaceResource
require.Eventually(t, func() bool {
require.Eventually(w.t, func() bool {
var err error
workspace, err := client.Workspace(ctx, workspaceID)
if !assert.NoError(t, err) {
workspace, err := w.client.Workspace(ctx, w.workspaceID)
if !assert.NoError(w.t, err) {
return false
}
if workspace.LatestBuild.Job.CompletedAt == nil {
@ -943,23 +987,25 @@ func AwaitWorkspaceAgents(t testing.TB, client *codersdk.Client, workspaceID uui
for _, resource := range workspace.LatestBuild.Resources {
for _, agent := range resource.Agents {
if len(agentNames) > 0 {
if len(w.agentNames) > 0 {
if _, ok := agentNamesMap[agent.Name]; !ok {
continue
}
}
if agent.Status != codersdk.WorkspaceAgentConnected {
t.Logf("agent %s not connected yet", agent.Name)
w.t.Logf("agent %s not connected yet", agent.Name)
return false
}
}
}
resources = workspace.LatestBuild.Resources
return true
if w.resourcesMatcher == nil {
return true
}
return w.resourcesMatcher(resources)
}, testutil.WaitLong, testutil.IntervalMedium)
t.Logf("got workspace agents (workspace %s)", workspaceID)
w.t.Logf("got workspace agents (workspace %s)", w.workspaceID)
return resources
}

View File

@ -29,6 +29,8 @@ import (
"cdr.dev/slog"
"cdr.dev/slog/sloggers/slogtest"
"github.com/coder/coder/v2/agent"
"github.com/coder/coder/v2/agent/agenttest"
agentproto "github.com/coder/coder/v2/agent/proto"
"github.com/coder/coder/v2/coderd/wsconncache"
"github.com/coder/coder/v2/codersdk"
"github.com/coder/coder/v2/codersdk/agentsdk"
@ -171,13 +173,12 @@ func setupAgent(t *testing.T, manifest agentsdk.Manifest, ptyTimeout time.Durati
_ = coordinator.Close()
})
manifest.AgentID = uuid.New()
aC := &client{
t: t,
agentID: manifest.AgentID,
manifest: manifest,
coordinator: coordinator,
derpMapUpdates: make(chan *tailcfg.DERPMap),
}
aC := newClient(
t,
slogtest.Make(t, nil).Leveled(slog.LevelDebug),
manifest,
coordinator,
)
t.Cleanup(aC.close)
closer := agent.New(agent.Options{
Client: aC,
@ -239,6 +240,45 @@ type client struct {
coordinator tailnet.Coordinator
closeOnce sync.Once
derpMapUpdates chan *tailcfg.DERPMap
server *drpcserver.Server
fakeAgentAPI *agenttest.FakeAgentAPI
}
func newClient(t *testing.T, logger slog.Logger, manifest agentsdk.Manifest, coordinator tailnet.Coordinator) *client {
logger = logger.Named("drpc")
coordPtr := atomic.Pointer[tailnet.Coordinator]{}
coordPtr.Store(&coordinator)
mux := drpcmux.New()
derpMapUpdates := make(chan *tailcfg.DERPMap)
drpcService := &tailnet.DRPCService{
CoordPtr: &coordPtr,
Logger: logger,
DerpMapUpdateFrequency: time.Microsecond,
DerpMapFn: func() *tailcfg.DERPMap { return <-derpMapUpdates },
}
err := proto.DRPCRegisterTailnet(mux, drpcService)
require.NoError(t, err)
fakeAAPI := agenttest.NewFakeAgentAPI(t, logger)
err = agentproto.DRPCRegisterAgent(mux, fakeAAPI)
require.NoError(t, err)
server := drpcserver.NewWithOptions(mux, drpcserver.Options{
Log: func(err error) {
if xerrors.Is(err, io.EOF) {
return
}
logger.Debug(context.Background(), "drpc server error", slog.Error(err))
},
})
return &client{
t: t,
agentID: manifest.AgentID,
manifest: manifest,
coordinator: coordinator,
derpMapUpdates: derpMapUpdates,
server: server,
fakeAgentAPI: fakeAAPI,
}
}
func (c *client) close() {
@ -250,35 +290,12 @@ func (c *client) Manifest(_ context.Context) (agentsdk.Manifest, error) {
}
func (c *client) Listen(_ context.Context) (drpc.Conn, error) {
logger := slogtest.Make(c.t, nil).Leveled(slog.LevelDebug).Named("drpc")
conn, lis := drpcsdk.MemTransportPipe()
c.t.Cleanup(func() {
_ = conn.Close()
_ = lis.Close()
})
coordPtr := atomic.Pointer[tailnet.Coordinator]{}
coordPtr.Store(&c.coordinator)
mux := drpcmux.New()
drpcService := &tailnet.DRPCService{
CoordPtr: &coordPtr,
Logger: logger,
DerpMapUpdateFrequency: time.Microsecond,
DerpMapFn: func() *tailcfg.DERPMap { return <-c.derpMapUpdates },
}
err := proto.DRPCRegisterTailnet(mux, drpcService)
if err != nil {
return nil, xerrors.Errorf("register DRPC service: %w", err)
}
server := drpcserver.NewWithOptions(mux, drpcserver.Options{
Log: func(err error) {
if xerrors.Is(err, io.EOF) ||
xerrors.Is(err, context.Canceled) ||
xerrors.Is(err, context.DeadlineExceeded) {
return
}
logger.Debug(context.Background(), "drpc server error", slog.Error(err))
},
})
serveCtx, cancel := context.WithCancel(context.Background())
c.t.Cleanup(cancel)
auth := tailnet.AgentTunnelAuth{}
@ -289,7 +306,7 @@ func (c *client) Listen(_ context.Context) (drpc.Conn, error) {
}
serveCtx = tailnet.WithStreamID(serveCtx, streamID)
go func() {
server.Serve(serveCtx, lis)
c.server.Serve(serveCtx, lis)
}()
return conn, nil
}
@ -317,7 +334,3 @@ func (*client) PostStartup(_ context.Context, _ agentsdk.PostStartupRequest) err
func (*client) PatchLogs(_ context.Context, _ agentsdk.PatchLogs) error {
return nil
}
func (*client) GetServiceBanner(_ context.Context) (codersdk.ServiceBannerConfig, error) {
return codersdk.ServiceBannerConfig{}, nil
}