chore(coderd/database/dbpurge): improve dbpurge tests

johnstcn · johnstcn · commit 0fdb3a33148f · 2023-12-01T10:15:05.000Z
diff --git a/coderd/database/dbpurge/dbpurge_test.go b/coderd/database/dbpurge/dbpurge_test.go
@@ -2,15 +2,23 @@ package dbpurge_test
 
 import (
 	"context"
+	"encoding/json"
+	"math/rand"
 	"testing"
+	"time"
 
-	"go.uber.org/goleak"
-
+	"github.com/google/uuid"
 	"github.com/stretchr/testify/require"
+	"go.uber.org/goleak"
 
 	"cdr.dev/slog/sloggers/slogtest"
-	"github.com/coder/coder/v2/coderd/database/dbmem"
+	"github.com/coder/coder/v2/coderd/database"
+	"github.com/coder/coder/v2/coderd/database/dbfake"
 	"github.com/coder/coder/v2/coderd/database/dbpurge"
+	"github.com/coder/coder/v2/coderd/database/dbtestutil"
+	"github.com/coder/coder/v2/coderd/database/dbtime"
+	"github.com/coder/coder/v2/cryptorand"
+	"github.com/coder/coder/v2/provisionersdk/proto"
 )
 
 func TestMain(m *testing.M) {
@@ -20,7 +28,111 @@ func TestMain(m *testing.M) {
 // Ensures no goroutines leak.
 func TestPurge(t *testing.T) {
 	t.Parallel()
-	purger := dbpurge.New(context.Background(), slogtest.Make(t, nil), dbmem.New())
+	db, _ := dbtestutil.NewDB(t)
+	opts := seedOpts{
+		NumAgents:        10,
+		NumLogsPerAgent:  10,
+		NumStatsPerAgent: 10,
+	}
+	seed(t, db, opts)
+	purger := dbpurge.New(context.Background(), slogtest.Make(t, nil), db)
 	err := purger.Close()
 	require.NoError(t, err)
 }
+
+type seedOpts struct {
+	NumAgents        int
+	NumLogsPerAgent  int
+	NumStatsPerAgent int
+}
+
+func seed(t *testing.T, db database.Store, opts seedOpts) {
+	t.Helper()
+
+	// Create a number of agents
+	agentIDs := make([]uuid.UUID, opts.NumAgents)
+	agentWSes := make(map[uuid.UUID]dbfake.WorkspaceResponse)
+	for i := 0; i < opts.NumAgents; i++ {
+		agentID := uuid.New()
+		wsr := dbfake.Workspace(t, db).WithAgent(func(agents []*proto.Agent) []*proto.Agent {
+			for _, agt := range agents {
+				agt.Id = agentID.String()
+			}
+			return agents
+		}).Do()
+		agentIDs[i] = agentID
+		agentWSes[agentID] = wsr
+	}
+
+	for _, agentID := range agentIDs {
+		// Create a number of logs for each agent
+		var entries []string
+		for i := 0; i < opts.NumLogsPerAgent; i++ {
+			randStr, err := cryptorand.String(1024)
+			require.NoError(t, err)
+			entries = append(entries, randStr)
+		}
+		_, err := db.InsertWorkspaceAgentLogs(context.Background(), database.InsertWorkspaceAgentLogsParams{
+			AgentID:      agentID,
+			CreatedAt:    time.Now(),
+			Output:       entries,
+			Level:        times(database.LogLevelInfo, opts.NumLogsPerAgent),
+			LogSourceID:  uuid.UUID{},
+			OutputLength: 0,
+		})
+		require.NoError(t, err)
+
+		// Insert a number of stats for each agent
+		err = db.InsertWorkspaceAgentStats(context.Background(), database.InsertWorkspaceAgentStatsParams{
+			ID:                          times(agentID, opts.NumStatsPerAgent),
+			CreatedAt:                   times(dbtime.Now(), opts.NumStatsPerAgent),
+			UserID:                      times(agentWSes[agentID].Workspace.OwnerID, opts.NumStatsPerAgent),
+			WorkspaceID:                 times(agentWSes[agentID].Workspace.ID, opts.NumStatsPerAgent),
+			TemplateID:                  times(agentWSes[agentID].Workspace.TemplateID, opts.NumStatsPerAgent),
+			AgentID:                     times(agentID, opts.NumStatsPerAgent),
+			ConnectionsByProto:          fakeConnectionsByProto(t, opts.NumStatsPerAgent),
+			ConnectionCount:             timesf(rand.Int63, opts.NumStatsPerAgent),
+			RxPackets:                   timesf(rand.Int63, opts.NumStatsPerAgent),
+			RxBytes:                     timesf(rand.Int63, opts.NumStatsPerAgent),
+			TxPackets:                   timesf(rand.Int63, opts.NumStatsPerAgent),
+			TxBytes:                     timesf(rand.Int63, opts.NumStatsPerAgent),
+			SessionCountVSCode:          timesf(rand.Int63, opts.NumStatsPerAgent),
+			SessionCountJetBrains:       timesf(rand.Int63, opts.NumStatsPerAgent),
+			SessionCountReconnectingPTY: timesf(rand.Int63, opts.NumStatsPerAgent),
+			SessionCountSSH:             timesf(rand.Int63, opts.NumStatsPerAgent),
+			ConnectionMedianLatencyMS:   timesf(rand.Float64, opts.NumStatsPerAgent),
+		})
+		require.NoError(t, err)
+	}
+}
+
+func fakeConnectionsByProto(t testing.TB, n int) json.RawMessage {
+	ms := make([]map[string]int64, n)
+	for i := 0; i < n; i++ {
+		m := map[string]int64{
+			"vscode": rand.Int63(),
+			"tty":    rand.Int63(),
+			"ssh":    rand.Int63(),
+		}
+		ms = append(ms, m)
+	}
+	bytes, err := json.Marshal(ms)
+	require.NoError(t, err)
+	return bytes
+}
+
+func times[T any](t T, n int) []T {
+	ts := make([]T, n)
+	for i := 0; i < n; i++ {
+		ts[i] = t
+	}
+	return ts
+}
+
+func timesf[T any](f func() T, n int) []T {
+	ts := make([]T, n)
+	for i := 0; i < n; i++ {
+		ts[i] = f()
+	}
+	return ts
+}
