fix(coderd/prometheusmetrics): close batcher to force flush before asserting agent stats (#9465)

johnstcn · web-flow · commit 22f31e721cf5 · 2023-08-31T11:40:57.000+01:00
diff --git a/coderd/prometheusmetrics/prometheusmetrics_test.go b/coderd/prometheusmetrics/prometheusmetrics_test.go
@@ -379,13 +379,15 @@ func TestAgentStats(t *testing.T) {
 	t.Cleanup(cancelFunc)
 
 	db, pubsub := dbtestutil.NewDB(t)
-	log := slogtest.Make(t, nil)
+	log := slogtest.Make(t, nil).Leveled(slog.LevelDebug)
 
 	batcher, closeBatcher, err := batchstats.New(ctx,
+		// We had previously set the batch size to 1 here, but that caused
+		// intermittent test flakes due to a race between the batcher completing
+		// its flush and the test asserting that the metrics were collected.
+		// Instead, we close the batcher after all stats have been posted, which
+		// forces a flush.
 		batchstats.WithStore(db),
-		// We want our stats, and we want them NOW.
-		batchstats.WithBatchSize(1),
-		batchstats.WithInterval(time.Hour),
 		batchstats.WithLogger(log),
 	)
 	require.NoError(t, err, "create stats batcher failed")
@@ -435,6 +437,11 @@ func TestAgentStats(t *testing.T) {
 		require.NoError(t, err)
 	}
 
+	// Ensure that all stats are flushed to the database
+	// before we query them. We do not expect any more stats
+	// to be posted after this.
+	closeBatcher()
+
 	// when
 	//
 	// Set initialCreateAfter to some time in the past, so that AgentStats would include all above PostStats,
