Skip to content

Commit 0248254

Browse files
committed
chore(coderd/database/dbpurge): improve dbpurge tests
1 parent 0789500 commit 0248254

File tree

1 file changed

+116
-4
lines changed

1 file changed

+116
-4
lines changed

coderd/database/dbpurge/dbpurge_test.go

Lines changed: 116 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -2,15 +2,23 @@ package dbpurge_test
22

33
import (
44
"context"
5+
"encoding/json"
6+
"math/rand"
57
"testing"
8+
"time"
69

7-
"go.uber.org/goleak"
8-
10+
"github.com/google/uuid"
911
"github.com/stretchr/testify/require"
12+
"go.uber.org/goleak"
1013

1114
"cdr.dev/slog/sloggers/slogtest"
12-
"github.com/coder/coder/v2/coderd/database/dbmem"
15+
"github.com/coder/coder/v2/coderd/database"
16+
"github.com/coder/coder/v2/coderd/database/dbfake"
1317
"github.com/coder/coder/v2/coderd/database/dbpurge"
18+
"github.com/coder/coder/v2/coderd/database/dbtestutil"
19+
"github.com/coder/coder/v2/coderd/database/dbtime"
20+
"github.com/coder/coder/v2/cryptorand"
21+
"github.com/coder/coder/v2/provisionersdk/proto"
1422
)
1523

1624
func TestMain(m *testing.M) {
@@ -20,7 +28,111 @@ func TestMain(m *testing.M) {
2028
// Ensures no goroutines leak.
2129
func TestPurge(t *testing.T) {
2230
t.Parallel()
23-
purger := dbpurge.New(context.Background(), slogtest.Make(t, nil), dbmem.New())
31+
db, _ := dbtestutil.NewDB(t)
32+
opts := seedOpts{
33+
NumAgents: 10,
34+
NumLogsPerAgent: 10,
35+
NumStatsPerAgent: 10,
36+
}
37+
seed(t, db, opts)
38+
purger := dbpurge.New(context.Background(), slogtest.Make(t, nil), db)
2439
err := purger.Close()
2540
require.NoError(t, err)
2641
}
42+
43+
type seedOpts struct {
44+
NumAgents int
45+
NumLogsPerAgent int
46+
NumStatsPerAgent int
47+
}
48+
49+
func seed(t *testing.T, db database.Store, opts seedOpts) {
50+
t.Helper()
51+
52+
// Create a number of agents
53+
agentIDs := make([]uuid.UUID, opts.NumAgents)
54+
agentWSes := make(map[uuid.UUID]dbfake.WorkspaceResponse)
55+
for i := 0; i < opts.NumAgents; i++ {
56+
agentID := uuid.New()
57+
wsr := dbfake.Workspace(t, db).WithAgent(func(agents []*proto.Agent) []*proto.Agent {
58+
for _, agt := range agents {
59+
agt.Id = agentID.String()
60+
}
61+
return agents
62+
}).Do()
63+
agentIDs[i] = agentID
64+
agentWSes[agentID] = wsr
65+
}
66+
67+
for _, agentID := range agentIDs {
68+
// Create a number of logs for each agent
69+
var entries []string
70+
for i := 0; i < opts.NumLogsPerAgent; i++ {
71+
randStr, err := cryptorand.String(1024)
72+
require.NoError(t, err)
73+
entries = append(entries, randStr)
74+
}
75+
_, err := db.InsertWorkspaceAgentLogs(context.Background(), database.InsertWorkspaceAgentLogsParams{
76+
AgentID: agentID,
77+
CreatedAt: time.Now(),
78+
Output: entries,
79+
Level: times(database.LogLevelInfo, opts.NumLogsPerAgent),
80+
LogSourceID: uuid.UUID{},
81+
OutputLength: 0,
82+
})
83+
require.NoError(t, err)
84+
85+
// Insert a number of stats for each agent
86+
err = db.InsertWorkspaceAgentStats(context.Background(), database.InsertWorkspaceAgentStatsParams{
87+
ID: times(agentID, opts.NumStatsPerAgent),
88+
CreatedAt: times(dbtime.Now(), opts.NumStatsPerAgent),
89+
UserID: times(agentWSes[agentID].Workspace.OwnerID, opts.NumStatsPerAgent),
90+
WorkspaceID: times(agentWSes[agentID].Workspace.ID, opts.NumStatsPerAgent),
91+
TemplateID: times(agentWSes[agentID].Workspace.TemplateID, opts.NumStatsPerAgent),
92+
AgentID: times(agentID, opts.NumStatsPerAgent),
93+
ConnectionsByProto: fakeConnectionsByProto(t, opts.NumStatsPerAgent),
94+
ConnectionCount: timesf(rand.Int63, opts.NumStatsPerAgent),
95+
RxPackets: timesf(rand.Int63, opts.NumStatsPerAgent),
96+
RxBytes: timesf(rand.Int63, opts.NumStatsPerAgent),
97+
TxPackets: timesf(rand.Int63, opts.NumStatsPerAgent),
98+
TxBytes: timesf(rand.Int63, opts.NumStatsPerAgent),
99+
SessionCountVSCode: timesf(rand.Int63, opts.NumStatsPerAgent),
100+
SessionCountJetBrains: timesf(rand.Int63, opts.NumStatsPerAgent),
101+
SessionCountReconnectingPTY: timesf(rand.Int63, opts.NumStatsPerAgent),
102+
SessionCountSSH: timesf(rand.Int63, opts.NumStatsPerAgent),
103+
ConnectionMedianLatencyMS: timesf(rand.Float64, opts.NumStatsPerAgent),
104+
})
105+
require.NoError(t, err)
106+
}
107+
}
108+
109+
func fakeConnectionsByProto(t testing.TB, n int) json.RawMessage {
110+
ms := make([]map[string]int64, n)
111+
for i := 0; i < n; i++ {
112+
m := map[string]int64{
113+
"vscode": rand.Int63(),
114+
"tty": rand.Int63(),
115+
"ssh": rand.Int63(),
116+
}
117+
ms = append(ms, m)
118+
}
119+
bytes, err := json.Marshal(ms)
120+
require.NoError(t, err)
121+
return bytes
122+
}
123+
124+
func times[T any](t T, n int) []T {
125+
ts := make([]T, n)
126+
for i := 0; i < n; i++ {
127+
ts[i] = t
128+
}
129+
return ts
130+
}
131+
132+
func timesf[T any](f func() T, n int) []T {
133+
ts := make([]T, n)
134+
for i := 0; i < n; i++ {
135+
ts[i] = f()
136+
}
137+
return ts
138+
}

0 commit comments

Comments
 (0)