Skip to content

Commit 36bfcbd

Browse files
deansheatherdeansheather
committed
chore: add agentapi tests
1 parent 3071863 commit 36bfcbd

12 files changed

+1729
-92
lines changed

‎coderd/agentapi/api.go

Lines changed: 16 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -25,7 +25,6 @@ import (
2525
"github.com/coder/coder/v2/coderd/schedule"
2626
"github.com/coder/coder/v2/coderd/tracing"
2727
"github.com/coder/coder/v2/codersdk/agentsdk"
28-
"github.com/coder/coder/v2/tailnet"
2928
)
3029

3130
const AgentAPIVersionDRPC = "2.0"
@@ -58,21 +57,18 @@ type Options struct {
5857
Database database.Store
5958
Pubsub pubsub.Pubsub
6059
DerpMapFn func() *tailcfg.DERPMap
61-
TailnetCoordinator *atomic.Pointer[tailnet.Coordinator]
6260
TemplateScheduleStore *atomic.Pointer[schedule.TemplateScheduleStore]
6361
StatsBatcher *batchstats.Batcher
6462
PublishWorkspaceUpdateFn func(ctx context.Context, workspaceID uuid.UUID)
6563
PublishWorkspaceAgentLogsUpdateFn func(ctx context.Context, workspaceAgentID uuid.UUID, msg agentsdk.LogsNotifyMessage)
6664

67-
AccessURL *url.URL
68-
AppHostname string
69-
AgentInactiveDisconnectTimeout time.Duration
70-
AgentFallbackTroubleshootingURL string
71-
AgentStatsRefreshInterval time.Duration
72-
DisableDirectConnections bool
73-
DerpForceWebSockets bool
74-
DerpMapUpdateFrequency time.Duration
75-
ExternalAuthConfigs []*externalauth.Config
65+
AccessURL *url.URL
66+
AppHostname string
67+
AgentStatsRefreshInterval time.Duration
68+
DisableDirectConnections bool
69+
DerpForceWebSockets bool
70+
DerpMapUpdateFrequency time.Duration
71+
ExternalAuthConfigs []*externalauth.Config
7672

7773
// Optional:
7874
// WorkspaceID avoids a future lookup to find the workspace ID by setting
@@ -89,17 +85,15 @@ func New(opts Options) *API {
8985
}
9086

9187
api.ManifestAPI = &ManifestAPI{
92-
AccessURL: opts.AccessURL,
93-
AppHostname: opts.AppHostname,
94-
AgentInactiveDisconnectTimeout: opts.AgentInactiveDisconnectTimeout,
95-
AgentFallbackTroubleshootingURL: opts.AgentFallbackTroubleshootingURL,
96-
ExternalAuthConfigs: opts.ExternalAuthConfigs,
97-
DisableDirectConnections: opts.DisableDirectConnections,
98-
DerpForceWebSockets: opts.DerpForceWebSockets,
99-
AgentFn: api.agent,
100-
Database: opts.Database,
101-
DerpMapFn: opts.DerpMapFn,
102-
TailnetCoordinator: opts.TailnetCoordinator,
88+
AccessURL: opts.AccessURL,
89+
AppHostname: opts.AppHostname,
90+
ExternalAuthConfigs: opts.ExternalAuthConfigs,
91+
DisableDirectConnections: opts.DisableDirectConnections,
92+
DerpForceWebSockets: opts.DerpForceWebSockets,
93+
AgentFn: api.agent,
94+
WorkspaceIDFn: api.workspaceID,
95+
Database: opts.Database,
96+
DerpMapFn: opts.DerpMapFn,
10397
}
10498

10599
api.ServiceBannerAPI = &ServiceBannerAPI{

‎coderd/agentapi/apps.go

Lines changed: 5 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -90,9 +90,11 @@ func (a *AppsAPI) BatchUpdateAppHealths(ctx context.Context, req *agentproto.Bat
9090
}
9191
}
9292

93-
err = a.PublishWorkspaceUpdateFn(ctx, &workspaceAgent)
94-
if err != nil {
95-
return nil, xerrors.Errorf("publish workspace update: %w", err)
93+
if a.PublishWorkspaceUpdateFn != nil && len(newApps) > 0 {
94+
err = a.PublishWorkspaceUpdateFn(ctx, &workspaceAgent)
95+
if err != nil {
96+
return nil, xerrors.Errorf("publish workspace update: %w", err)
97+
}
9698
}
9799
return &agentproto.BatchUpdateAppHealthResponse{}, nil
98100
}

‎coderd/agentapi/apps_test.go

Lines changed: 253 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,253 @@
1+
package agentapi_test
2+
3+
import (
4+
"context"
5+
"sync/atomic"
6+
"testing"
7+
8+
"github.com/golang/mock/gomock"
9+
"github.com/google/uuid"
10+
"github.com/stretchr/testify/require"
11+
12+
"cdr.dev/slog/sloggers/slogtest"
13+
14+
agentproto "github.com/coder/coder/v2/agent/proto"
15+
"github.com/coder/coder/v2/coderd/agentapi"
16+
"github.com/coder/coder/v2/coderd/database"
17+
"github.com/coder/coder/v2/coderd/database/dbmock"
18+
)
19+
20+
func TestBatchUpdateAppHealths(t *testing.T) {
21+
t.Parallel()
22+
23+
var (
24+
agent = database.WorkspaceAgent{
25+
ID: uuid.New(),
26+
}
27+
app1 = database.WorkspaceApp{
28+
ID: uuid.New(),
29+
AgentID: agent.ID,
30+
Slug: "code-server-1",
31+
DisplayName: "code-server 1",
32+
HealthcheckUrl: "http://localhost:3000",
33+
Health: database.WorkspaceAppHealthInitializing,
34+
}
35+
app2 = database.WorkspaceApp{
36+
ID: uuid.New(),
37+
AgentID: agent.ID,
38+
Slug: "code-server-2",
39+
DisplayName: "code-server 2",
40+
HealthcheckUrl: "http://localhost:3001",
41+
Health: database.WorkspaceAppHealthHealthy,
42+
}
43+
)
44+
45+
t.Run("OK", func(t *testing.T) {
46+
t.Parallel()
47+
48+
dbM := dbmock.NewMockStore(gomock.NewController(t))
49+
dbM.EXPECT().GetWorkspaceAppsByAgentID(gomock.Any(), agent.ID).Return([]database.WorkspaceApp{app1, app2}, nil)
50+
dbM.EXPECT().UpdateWorkspaceAppHealthByID(gomock.Any(), database.UpdateWorkspaceAppHealthByIDParams{
51+
ID: app1.ID,
52+
Health: database.WorkspaceAppHealthHealthy,
53+
}).Return(nil)
54+
dbM.EXPECT().UpdateWorkspaceAppHealthByID(gomock.Any(), database.UpdateWorkspaceAppHealthByIDParams{
55+
ID: app2.ID,
56+
Health: database.WorkspaceAppHealthUnhealthy,
57+
}).Return(nil)
58+
59+
var publishCalled int64
60+
api := &agentapi.AppsAPI{
61+
AgentFn: func(context.Context) (database.WorkspaceAgent, error) {
62+
return agent, nil
63+
},
64+
Database: dbM,
65+
Log: slogtest.Make(t, nil),
66+
PublishWorkspaceUpdateFn: func(ctx context.Context, wa *database.WorkspaceAgent) error {
67+
atomic.AddInt64(&publishCalled, 1)
68+
return nil
69+
},
70+
}
71+
72+
// Set both to healthy, only one should be updated in the DB.
73+
resp, err := api.BatchUpdateAppHealths(context.Background(), &agentproto.BatchUpdateAppHealthRequest{
74+
Updates: []*agentproto.BatchUpdateAppHealthRequest_HealthUpdate{
75+
{
76+
Id: app1.ID[:],
77+
Health: agentproto.AppHealth_HEALTHY,
78+
},
79+
{
80+
Id: app2.ID[:],
81+
Health: agentproto.AppHealth_UNHEALTHY,
82+
},
83+
},
84+
})
85+
require.NoError(t, err)
86+
require.Equal(t, &agentproto.BatchUpdateAppHealthResponse{}, resp)
87+
88+
require.EqualValues(t, 1, atomic.LoadInt64(&publishCalled))
89+
})
90+
91+
t.Run("Unchanged", func(t *testing.T) {
92+
t.Parallel()
93+
94+
dbM := dbmock.NewMockStore(gomock.NewController(t))
95+
dbM.EXPECT().GetWorkspaceAppsByAgentID(gomock.Any(), agent.ID).Return([]database.WorkspaceApp{app1, app2}, nil)
96+
97+
var publishCalled int64
98+
api := &agentapi.AppsAPI{
99+
AgentFn: func(context.Context) (database.WorkspaceAgent, error) {
100+
return agent, nil
101+
},
102+
Database: dbM,
103+
Log: slogtest.Make(t, nil),
104+
PublishWorkspaceUpdateFn: func(ctx context.Context, wa *database.WorkspaceAgent) error {
105+
atomic.AddInt64(&publishCalled, 1)
106+
return nil
107+
},
108+
}
109+
110+
// Set both to their current status, neither should be updated in the
111+
// DB.
112+
resp, err := api.BatchUpdateAppHealths(context.Background(), &agentproto.BatchUpdateAppHealthRequest{
113+
Updates: []*agentproto.BatchUpdateAppHealthRequest_HealthUpdate{
114+
{
115+
Id: app1.ID[:],
116+
Health: agentproto.AppHealth_INITIALIZING,
117+
},
118+
{
119+
Id: app2.ID[:],
120+
Health: agentproto.AppHealth_HEALTHY,
121+
},
122+
},
123+
})
124+
require.NoError(t, err)
125+
require.Equal(t, &agentproto.BatchUpdateAppHealthResponse{}, resp)
126+
127+
require.EqualValues(t, 0, atomic.LoadInt64(&publishCalled))
128+
})
129+
130+
t.Run("Empty", func(t *testing.T) {
131+
t.Parallel()
132+
133+
// No DB queries are made if there are no updates to process.
134+
dbM := dbmock.NewMockStore(gomock.NewController(t))
135+
136+
var publishCalled int64
137+
api := &agentapi.AppsAPI{
138+
AgentFn: func(context.Context) (database.WorkspaceAgent, error) {
139+
return agent, nil
140+
},
141+
Database: dbM,
142+
Log: slogtest.Make(t, nil),
143+
PublishWorkspaceUpdateFn: func(ctx context.Context, wa *database.WorkspaceAgent) error {
144+
atomic.AddInt64(&publishCalled, 1)
145+
return nil
146+
},
147+
}
148+
149+
// Do nothing.
150+
resp, err := api.BatchUpdateAppHealths(context.Background(), &agentproto.BatchUpdateAppHealthRequest{
151+
Updates: []*agentproto.BatchUpdateAppHealthRequest_HealthUpdate{},
152+
})
153+
require.NoError(t, err)
154+
require.Equal(t, &agentproto.BatchUpdateAppHealthResponse{}, resp)
155+
156+
require.EqualValues(t, 0, atomic.LoadInt64(&publishCalled))
157+
})
158+
159+
t.Run("AppNoHealthcheck", func(t *testing.T) {
160+
t.Parallel()
161+
162+
app3 := database.WorkspaceApp{
163+
ID: uuid.New(),
164+
AgentID: agent.ID,
165+
Slug: "code-server-3",
166+
DisplayName: "code-server 3",
167+
}
168+
169+
dbM := dbmock.NewMockStore(gomock.NewController(t))
170+
dbM.EXPECT().GetWorkspaceAppsByAgentID(gomock.Any(), agent.ID).Return([]database.WorkspaceApp{app3}, nil)
171+
172+
api := &agentapi.AppsAPI{
173+
AgentFn: func(context.Context) (database.WorkspaceAgent, error) {
174+
return agent, nil
175+
},
176+
Database: dbM,
177+
Log: slogtest.Make(t, nil),
178+
PublishWorkspaceUpdateFn: nil,
179+
}
180+
181+
// Set app3 to healthy, should error.
182+
resp, err := api.BatchUpdateAppHealths(context.Background(), &agentproto.BatchUpdateAppHealthRequest{
183+
Updates: []*agentproto.BatchUpdateAppHealthRequest_HealthUpdate{
184+
{
185+
Id: app3.ID[:],
186+
Health: agentproto.AppHealth_HEALTHY,
187+
},
188+
},
189+
})
190+
require.Error(t, err)
191+
require.ErrorContains(t, err, "does not have healthchecks enabled")
192+
require.Nil(t, resp)
193+
})
194+
195+
t.Run("UnknownApp", func(t *testing.T) {
196+
t.Parallel()
197+
198+
dbM := dbmock.NewMockStore(gomock.NewController(t))
199+
dbM.EXPECT().GetWorkspaceAppsByAgentID(gomock.Any(), agent.ID).Return([]database.WorkspaceApp{app1, app2}, nil)
200+
201+
api := &agentapi.AppsAPI{
202+
AgentFn: func(context.Context) (database.WorkspaceAgent, error) {
203+
return agent, nil
204+
},
205+
Database: dbM,
206+
Log: slogtest.Make(t, nil),
207+
PublishWorkspaceUpdateFn: nil,
208+
}
209+
210+
// Set an unknown app to healthy, should error.
211+
id := uuid.New()
212+
resp, err := api.BatchUpdateAppHealths(context.Background(), &agentproto.BatchUpdateAppHealthRequest{
213+
Updates: []*agentproto.BatchUpdateAppHealthRequest_HealthUpdate{
214+
{
215+
Id: id[:],
216+
Health: agentproto.AppHealth_HEALTHY,
217+
},
218+
},
219+
})
220+
require.Error(t, err)
221+
require.ErrorContains(t, err, "not found")
222+
require.Nil(t, resp)
223+
})
224+
225+
t.Run("InvalidHealth", func(t *testing.T) {
226+
t.Parallel()
227+
228+
dbM := dbmock.NewMockStore(gomock.NewController(t))
229+
dbM.EXPECT().GetWorkspaceAppsByAgentID(gomock.Any(), agent.ID).Return([]database.WorkspaceApp{app1, app2}, nil)
230+
231+
api := &agentapi.AppsAPI{
232+
AgentFn: func(context.Context) (database.WorkspaceAgent, error) {
233+
return agent, nil
234+
},
235+
Database: dbM,
236+
Log: slogtest.Make(t, nil),
237+
PublishWorkspaceUpdateFn: nil,
238+
}
239+
240+
// Set an unknown app to healthy, should error.
241+
resp, err := api.BatchUpdateAppHealths(context.Background(), &agentproto.BatchUpdateAppHealthRequest{
242+
Updates: []*agentproto.BatchUpdateAppHealthRequest_HealthUpdate{
243+
{
244+
Id: app1.ID[:],
245+
Health: -999,
246+
},
247+
},
248+
})
249+
require.Error(t, err)
250+
require.ErrorContains(t, err, "unknown health status")
251+
require.Nil(t, resp)
252+
})
253+
}

‎coderd/agentapi/lifecycle.go

Lines changed: 22 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@ package agentapi
33
import (
44
"context"
55
"database/sql"
6+
"time"
67

78
"github.com/google/uuid"
89
"golang.org/x/mod/semver"
@@ -21,6 +22,15 @@ type LifecycleAPI struct {
2122
Database database.Store
2223
Log slog.Logger
2324
PublishWorkspaceUpdateFn func(context.Context, *database.WorkspaceAgent) error
25+
26+
TimeNowFn func() time.Time // defaults to dbtime.Now()
27+
}
28+
29+
func (a *LifecycleAPI) now() time.Time {
30+
if a.TimeNowFn != nil {
31+
return a.TimeNowFn()
32+
}
33+
return dbtime.Now()
2434
}
2535

2636
func (a *LifecycleAPI) UpdateLifecycle(ctx context.Context, req *agentproto.UpdateLifecycleRequest) (*agentproto.Lifecycle, error) {
@@ -68,7 +78,7 @@ func (a *LifecycleAPI) UpdateLifecycle(ctx context.Context, req *agentproto.Upda
6878

6979
changedAt := req.Lifecycle.ChangedAt.AsTime()
7080
if changedAt.IsZero() {
71-
changedAt = dbtime.Now()
81+
changedAt = a.now()
7282
req.Lifecycle.ChangedAt = timestamppb.New(changedAt)
7383
}
7484
dbChangedAt := sql.NullTime{Time: changedAt, Valid: true}
@@ -78,8 +88,13 @@ func (a *LifecycleAPI) UpdateLifecycle(ctx context.Context, req *agentproto.Upda
7888
switch lifecycleState {
7989
case database.WorkspaceAgentLifecycleStateStarting:
8090
startedAt = dbChangedAt
81-
readyAt.Valid = false // This agent is re-starting, so it's not ready yet.
91+
// This agent is (re)starting, so it's not ready yet.
92+
readyAt.Time = time.Time{}
93+
readyAt.Valid = false
8294
case database.WorkspaceAgentLifecycleStateReady, database.WorkspaceAgentLifecycleStateStartError:
95+
if !startedAt.Valid {
96+
startedAt = dbChangedAt
97+
}
8398
readyAt = dbChangedAt
8499
}
85100

@@ -97,9 +112,11 @@ func (a *LifecycleAPI) UpdateLifecycle(ctx context.Context, req *agentproto.Upda
97112
return nil, xerrors.Errorf("update workspace agent lifecycle state: %w", err)
98113
}
99114

100-
err = a.PublishWorkspaceUpdateFn(ctx, &workspaceAgent)
101-
if err != nil {
102-
return nil, xerrors.Errorf("publish workspace update: %w", err)
115+
if a.PublishWorkspaceUpdateFn != nil {
116+
err = a.PublishWorkspaceUpdateFn(ctx, &workspaceAgent)
117+
if err != nil {
118+
return nil, xerrors.Errorf("publish workspace update: %w", err)
119+
}
103120
}
104121

105122
return req.Lifecycle, nil

0 commit comments

Comments
 (0)