Skip to content

Commit df980e7

Browse files
johnstcnjohnstcn
committed
chore: improve coverage of dbpurge tests for DeleteOldWorkspaceAgentLogs
1 parent fcd2606 commit df980e7

File tree

1 file changed

+149
-99
lines changed

1 file changed

+149
-99
lines changed

‎coderd/database/dbpurge/dbpurge_test.go

Lines changed: 149 additions & 99 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,7 @@ import (
1111
"time"
1212

1313
"github.com/google/uuid"
14+
"github.com/stretchr/testify/assert"
1415
"github.com/stretchr/testify/require"
1516
"go.uber.org/goleak"
1617
"golang.org/x/exp/slices"
@@ -181,139 +182,188 @@ func containsWorkspaceAgentStat(stats []database.GetWorkspaceAgentStatsRow, need
181182

182183
//nolint:paralleltest // It uses LockIDDBPurge.
183184
func TestDeleteOldWorkspaceAgentLogs(t *testing.T) {
184-
db, _ := dbtestutil.NewDB(t)
185+
ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitShort)
186+
defer cancel()
187+
clk := quartz.NewMock(t)
188+
now := dbtime.Now()
189+
threshold := now.Add(-7 * 24 * time.Hour)
190+
beforeThreshold := threshold.Add(-time.Hour)
191+
afterThreshold := threshold.Add(time.Hour)
192+
clk.Set(now).MustWait(ctx)
193+
194+
db, _ := dbtestutil.NewDB(t, dbtestutil.WithDumpOnFailure())
185195
org := dbgen.Organization(t, db, database.Organization{})
186196
user := dbgen.User(t, db, database.User{})
187197
_ = dbgen.OrganizationMember(t, db, database.OrganizationMember{UserID: user.ID, OrganizationID: org.ID})
188198
tv := dbgen.TemplateVersion(t, db, database.TemplateVersion{OrganizationID: org.ID, CreatedBy: user.ID})
189199
tmpl := dbgen.Template(t, db, database.Template{OrganizationID: org.ID, ActiveVersionID: tv.ID, CreatedBy: user.ID})
190200

191201
logger := slogtest.Make(t, &slogtest.Options{IgnoreErrors: true})
192-
now := dbtime.Now()
193202

194-
//nolint:paralleltest // It uses LockIDDBPurge.
195-
t.Run("AgentHasNotConnectedSinceWeek_LogsExpired", func(t *testing.T) {
196-
ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitShort)
197-
defer cancel()
198-
clk := quartz.NewMock(t)
199-
clk.Set(now).MustWait(ctx)
200-
201-
// After dbpurge completes, the ticker is reset. Trap this call.
202-
trapReset := clk.Trap().TickerReset()
203-
defer trapReset.Close()
204-
205-
// given: an agent with logs older than threshold
206-
agent := mustCreateAgentWithLogs(ctx, t, db, user, org, tmpl, tv, now.Add(-8*24*time.Hour), t.Name())
207-
208-
// when dbpurge runs
209-
closer := dbpurge.New(ctx, logger, db, clk)
210-
defer closer.Close()
211-
// Wait for the initial nanosecond tick.
212-
clk.Advance(time.Nanosecond).MustWait(ctx)
213-
214-
trapReset.MustWait(ctx).Release() // Wait for ticker.Reset()
215-
d, w := clk.AdvanceNext()
216-
require.Equal(t, 10*time.Minute, d)
217-
218-
closer.Close() // doTick() has now run.
219-
w.MustWait(ctx)
220-
221-
// then the logs should be gone
222-
agentLogs, err := db.GetWorkspaceAgentLogsAfter(ctx, database.GetWorkspaceAgentLogsAfterParams{
223-
AgentID: agent,
224-
CreatedAfter: 0,
225-
})
226-
require.NoError(t, err)
227-
require.Empty(t, agentLogs, "expected agent logs to be empty")
228-
})
203+
// Given the following:
204+
205+
// Workspace A was built once before the threshold, and never connected.
206+
wsA := dbgen.Workspace(t, db, database.Workspace{OwnerID: user.ID, OrganizationID: org.ID, TemplateID: tmpl.ID})
207+
wbA1 := mustCreateWorkspaceBuild(t, db, org, tv, wsA.ID, beforeThreshold, 1)
208+
agentA1 := mustCreateAgent(t, db, wbA1)
209+
mustCreateAgentLogs(ctx, t, db, agentA1.ID, nil, "agent a1 logs should be deleted")
210+
211+
// Workspace B was built twice before the threshold.
212+
wsB := dbgen.Workspace(t, db, database.Workspace{OwnerID: user.ID, OrganizationID: org.ID, TemplateID: tmpl.ID})
213+
wbB1 := mustCreateWorkspaceBuild(t, db, org, tv, wsB.ID, beforeThreshold, 1)
214+
wbB2 := mustCreateWorkspaceBuild(t, db, org, tv, wsB.ID, beforeThreshold, 2)
215+
agentB1 := mustCreateAgent(t, db, wbB1)
216+
agentB2 := mustCreateAgent(t, db, wbB2)
217+
mustCreateAgentLogs(ctx, t, db, agentB1.ID, &beforeThreshold, "agent b1 logs should be deleted")
218+
mustCreateAgentLogs(ctx, t, db, agentB2.ID, &beforeThreshold, "agent b2 logs should be retained")
219+
220+
// Workspace C was built once before the threshold, and once after.
221+
wsC := dbgen.Workspace(t, db, database.Workspace{OwnerID: user.ID, OrganizationID: org.ID, TemplateID: tmpl.ID})
222+
wbC1 := mustCreateWorkspaceBuild(t, db, org, tv, wsC.ID, beforeThreshold, 1)
223+
wbC2 := mustCreateWorkspaceBuild(t, db, org, tv, wsC.ID, afterThreshold, 2)
224+
agentC1 := mustCreateAgent(t, db, wbC1)
225+
agentC2 := mustCreateAgent(t, db, wbC2)
226+
mustCreateAgentLogs(ctx, t, db, agentC1.ID, &beforeThreshold, "agent c1 logs should be deleted")
227+
mustCreateAgentLogs(ctx, t, db, agentC2.ID, &afterThreshold, "agent c2 logs should be retained")
228+
229+
// Workspace D was built twice after the threshold.
230+
wsD := dbgen.Workspace(t, db, database.Workspace{OwnerID: user.ID, OrganizationID: org.ID, TemplateID: tmpl.ID})
231+
wbD1 := mustCreateWorkspaceBuild(t, db, org, tv, wsD.ID, afterThreshold, 1)
232+
wbD2 := mustCreateWorkspaceBuild(t, db, org, tv, wsD.ID, afterThreshold, 2)
233+
agentD1 := mustCreateAgent(t, db, wbD1)
234+
agentD2 := mustCreateAgent(t, db, wbD2)
235+
mustCreateAgentLogs(ctx, t, db, agentD1.ID, &afterThreshold, "agent d1 logs should be retained")
236+
mustCreateAgentLogs(ctx, t, db, agentD2.ID, &afterThreshold, "agent d2 logs should be retained")
237+
238+
// Workspace E was build once after threshold but never connected.
239+
wsE := dbgen.Workspace(t, db, database.Workspace{OwnerID: user.ID, OrganizationID: org.ID, TemplateID: tmpl.ID})
240+
wbE1 := mustCreateWorkspaceBuild(t, db, org, tv, wsE.ID, beforeThreshold, 1)
241+
agentE1 := mustCreateAgent(t, db, wbE1)
242+
mustCreateAgentLogs(ctx, t, db, agentE1.ID, nil, "agent e1 logs should be retained")
243+
244+
// when dbpurge runs
245+
246+
// After dbpurge completes, the ticker is reset. Trap this call.
247+
trapReset := clk.Trap().TickerReset()
248+
defer trapReset.Close()
229249

230-
//nolint:paralleltest // It uses LockIDDBPurge.
231-
t.Run("AgentConnectedSixDaysAgo_LogsValid", func(t *testing.T) {
232-
ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitShort)
233-
defer cancel()
234-
clk := quartz.NewMock(t)
235-
clk.Set(now).MustWait(ctx)
236-
237-
// After dbpurge completes, the ticker is reset. Trap this call.
238-
trapReset := clk.Trap().TickerReset()
239-
defer trapReset.Close()
240-
241-
// given: an agent with logs newer than threshold
242-
agent := mustCreateAgentWithLogs(ctx, t, db, user, org, tmpl, tv, now.Add(-6*24*time.Hour), t.Name())
243-
244-
// when dbpurge runs
245-
closer := dbpurge.New(ctx, logger, db, clk)
246-
defer closer.Close()
247-
248-
// Wait for the initial nanosecond tick.
249-
clk.Advance(time.Nanosecond).MustWait(ctx)
250-
251-
trapReset.MustWait(ctx).Release() // Wait for ticker.Reset()
252-
d, w := clk.AdvanceNext()
253-
require.Equal(t, 10*time.Minute, d)
254-
255-
closer.Close() // doTick() has now run.
256-
w.MustWait(ctx)
257-
258-
// then the logs should still be there
259-
agentLogs, err := db.GetWorkspaceAgentLogsAfter(ctx, database.GetWorkspaceAgentLogsAfterParams{
260-
AgentID: agent,
261-
})
262-
require.NoError(t, err)
263-
require.NotEmpty(t, agentLogs)
264-
for _, al := range agentLogs {
265-
require.Equal(t, t.Name(), al.Output)
266-
}
267-
})
268-
}
250+
closer := dbpurge.New(ctx, logger, db, clk)
251+
defer closer.Close()
252+
// Wait for the initial nanosecond tick.
253+
clk.Advance(time.Nanosecond).MustWait(ctx)
269254

270-
func mustCreateAgentWithLogs(ctx context.Context, t *testing.T, db database.Store, user database.User, org database.Organization, tmpl database.Template, tv database.TemplateVersion, agentLastConnectedAt time.Time, output string) uuid.UUID {
271-
agent := mustCreateAgent(t, db, user, org, tmpl, tv)
255+
trapReset.MustWait(ctx).Release() // Wait for ticker.Reset()
256+
d, w := clk.AdvanceNext()
257+
require.Equal(t, 10*time.Minute, d)
258+
259+
closer.Close() // doTick() has now run.
260+
w.MustWait(ctx)
261+
262+
// then logs related to the following agents should be deleted:
263+
// Agent A1 never connected and was created before the threshold.
264+
assertNoWorkspaceAgentLogs(ctx, t, db, agentA1.ID)
265+
// Agent B1 is not the latest build and the logs are from before threshold.
266+
assertNoWorkspaceAgentLogs(ctx, t, db, agentB1.ID)
267+
// Agent C1 is not the latest build and the logs are from before threshold.
268+
assertNoWorkspaceAgentLogs(ctx, t, db, agentC1.ID)
269+
270+
// then logs related to the following agents should be retained:
271+
// Agent B2 is the latest build.
272+
assertWorkspaceAgentLogs(ctx, t, db, agentB2.ID, "agent b2 logs should be retained")
273+
// Agent C2 is the latest build.
274+
assertWorkspaceAgentLogs(ctx, t, db, agentC2.ID, "agent c2 logs should be retained")
275+
// Agents D1, D2, and E1 are all after threshold.
276+
assertWorkspaceAgentLogs(ctx, t, db, agentD1.ID, "agent d1 logs should be retained")
277+
assertWorkspaceAgentLogs(ctx, t, db, agentD2.ID, "agent d2 logs should be retained")
278+
assertWorkspaceAgentLogs(ctx, t, db, agentE1.ID, "agent e1 logs should be retained")
279+
}
272280

273-
err := db.UpdateWorkspaceAgentConnectionByID(ctx, database.UpdateWorkspaceAgentConnectionByIDParams{
274-
ID: agent.ID,
275-
LastConnectedAt: sql.NullTime{Time: agentLastConnectedAt, Valid: true},
276-
})
277-
require.NoError(t, err)
278-
_, err = db.InsertWorkspaceAgentLogs(ctx, database.InsertWorkspaceAgentLogsParams{
279-
AgentID: agent.ID,
280-
CreatedAt: agentLastConnectedAt,
281-
Output: []string{output},
282-
Level: []database.LogLevel{database.LogLevelDebug},
281+
func assertNoWorkspaceAgentLogs(ctx context.Context, t *testing.T, db database.Store, agentID uuid.UUID) {
282+
t.Helper()
283+
agentLogs, err := db.GetWorkspaceAgentLogsAfter(ctx, database.GetWorkspaceAgentLogsAfterParams{
284+
AgentID: agentID,
285+
CreatedAfter: 0,
283286
})
284287
require.NoError(t, err)
285-
// Make sure that agent logs have been collected.
288+
assert.Empty(t, agentLogs)
289+
}
290+
291+
func assertWorkspaceAgentLogs(ctx context.Context, t *testing.T, db database.Store, agentID uuid.UUID, msg string) {
292+
t.Helper()
286293
agentLogs, err := db.GetWorkspaceAgentLogsAfter(ctx, database.GetWorkspaceAgentLogsAfterParams{
287-
AgentID: agent.ID,
294+
AgentID: agentID,
295+
CreatedAfter: 0,
288296
})
289297
require.NoError(t, err)
290-
require.NotZero(t, agentLogs, "agent logs must be present")
291-
return agent.ID
298+
assert.NotEmpty(t, agentLogs)
299+
for _, al := range agentLogs {
300+
assert.Equal(t, msg, al.Output)
301+
}
292302
}
293303

294-
func mustCreateAgent(t *testing.T, db database.Store, user database.User, org database.Organization, tmpl database.Template, tv database.TemplateVersion) database.WorkspaceAgent {
295-
workspace := dbgen.Workspace(t, db, database.Workspace{OwnerID: user.ID, OrganizationID: org.ID, TemplateID: tmpl.ID})
304+
func mustCreateWorkspaceBuild(t *testing.T, db database.Store, org database.Organization, tv database.TemplateVersion, wsID uuid.UUID, createdAt time.Time, n int32) database.WorkspaceBuild {
305+
t.Helper()
296306
job := dbgen.ProvisionerJob(t, db, nil, database.ProvisionerJob{
307+
CreatedAt: createdAt,
297308
OrganizationID: org.ID,
298309
Type: database.ProvisionerJobTypeWorkspaceBuild,
299310
Provisioner: database.ProvisionerTypeEcho,
300311
StorageMethod: database.ProvisionerStorageMethodFile,
301312
})
302-
_ = dbgen.WorkspaceBuild(t, db, database.WorkspaceBuild{
303-
WorkspaceID: workspace.ID,
313+
wb := dbgen.WorkspaceBuild(t, db, database.WorkspaceBuild{
314+
CreatedAt: createdAt,
315+
WorkspaceID: wsID,
304316
JobID: job.ID,
305317
TemplateVersionID: tv.ID,
306318
Transition: database.WorkspaceTransitionStart,
307319
Reason: database.BuildReasonInitiator,
320+
BuildNumber: n,
308321
})
322+
require.Equal(t, createdAt.UTC(), wb.CreatedAt.UTC())
323+
return wb
324+
}
325+
326+
func mustCreateAgent(t *testing.T, db database.Store, wb database.WorkspaceBuild) database.WorkspaceAgent {
327+
t.Helper()
309328
resource := dbgen.WorkspaceResource(t, db, database.WorkspaceResource{
310-
JobID: job.ID,
329+
JobID: wb.JobID,
311330
Transition: database.WorkspaceTransitionStart,
331+
CreatedAt: wb.CreatedAt,
312332
})
313333

314-
return dbgen.WorkspaceAgent(t, db, database.WorkspaceAgent{
315-
ResourceID: resource.ID,
334+
wa := dbgen.WorkspaceAgent(t, db, database.WorkspaceAgent{
335+
ResourceID: resource.ID,
336+
CreatedAt: wb.CreatedAt,
337+
FirstConnectedAt: sql.NullTime{},
338+
DisconnectedAt: sql.NullTime{},
339+
LastConnectedAt: sql.NullTime{},
316340
})
341+
require.Equal(t, wb.CreatedAt.UTC(), wa.CreatedAt.UTC())
342+
return wa
343+
}
344+
345+
func mustCreateAgentLogs(ctx context.Context, t *testing.T, db database.Store, agentID uuid.UUID, agentLastConnectedAt *time.Time, output string) uuid.UUID {
346+
t.Helper()
347+
if agentLastConnectedAt != nil {
348+
require.NoError(t, db.UpdateWorkspaceAgentConnectionByID(ctx, database.UpdateWorkspaceAgentConnectionByIDParams{
349+
ID: agentID,
350+
LastConnectedAt: sql.NullTime{Time: *agentLastConnectedAt, Valid: true},
351+
}))
352+
}
353+
_, err := db.InsertWorkspaceAgentLogs(ctx, database.InsertWorkspaceAgentLogsParams{
354+
AgentID: agentID,
355+
// CreatedAt: agentLastConnectedAt,
356+
Output: []string{output},
357+
Level: []database.LogLevel{database.LogLevelDebug},
358+
})
359+
require.NoError(t, err)
360+
// Make sure that agent logs have been collected.
361+
agentLogs, err := db.GetWorkspaceAgentLogsAfter(ctx, database.GetWorkspaceAgentLogsAfterParams{
362+
AgentID: agentID,
363+
})
364+
require.NoError(t, err)
365+
require.NotEmpty(t, agentLogs, "agent logs must be present")
366+
return agentID
317367
}
318368

319369
//nolint:paralleltest // It uses LockIDDBPurge.

0 commit comments

Comments
 (0)