@@ -2,15 +2,23 @@ package dbpurge_test
2
2
3
3
import (
4
4
"context"
5
+ "encoding/json"
6
+ "math/rand"
5
7
"testing"
8
+ "time"
6
9
7
- "go.uber.org/goleak"
8
-
10
+ "github.com/google/uuid"
9
11
"github.com/stretchr/testify/require"
12
+ "go.uber.org/goleak"
10
13
11
14
"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"
13
17
"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"
14
22
)
15
23
16
24
func TestMain (m * testing.M ) {
@@ -20,7 +28,111 @@ func TestMain(m *testing.M) {
20
28
// Ensures no goroutines leak.
21
29
func TestPurge (t * testing.T ) {
22
30
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 )
24
39
err := purger .Close ()
25
40
require .NoError (t , err )
26
41
}
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