Merge pull request #2201 from josephschorr/stats-test-time

Remove sleep in stats test unless needed

pkg/datastore/test/stats.go

@@ -2,6 +2,7 @@ package test
 
 import (
 	"context"
+	"strings"
 	"testing"
 	"time"
 
@@ -10,7 +11,7 @@ import (
 	"github.com/authzed/spicedb/internal/testfixtures"
 )
 
-const statsRetryCount = 3
+const statsRetryCount = 10
 
 func StatsTest(t *testing.T, tester DatastoreTester) {
 	ctx := context.Background()
@@ -21,23 +22,27 @@ func StatsTest(t *testing.T, tester DatastoreTester) {
 
 	ds, _ = testfixtures.StandardDatastoreWithData(ds, require)
 
-	// stats use follower reads, need to wait a bit so that the base tables
-	// have a chance to be follower-read
-	time.Sleep(5 * time.Second)
-
 	for retryCount := statsRetryCount; retryCount >= 0; retryCount-- {
 		stats, err := ds.Statistics(ctx)
+
+		// If the error contains code 3D000, the *database* is not yet available, which
+		// can happen due to using follower reads in CockroachDB. We should retry in this case.
+		if err != nil && strings.Contains(err.Error(), "3D000") {
+			time.Sleep(5 * time.Second)
+			continue
+		}
+
 		require.NoError(err)
 
 		require.Len(stats.UniqueID, 36, "unique ID must be a valid UUID")
-		require.Len(stats.ObjectTypeStatistics, 3, "must report object stats")
 
 		if stats.EstimatedRelationshipCount == uint64(0) && retryCount > 0 {
 			// Sleep for a bit to get the stats table to update.
 			time.Sleep(500 * time.Millisecond)
 			continue
 		}
 
+		require.Len(stats.ObjectTypeStatistics, 3, "must report object stats")
 		require.Greater(stats.EstimatedRelationshipCount, uint64(0), "must report some relationships")
 
 		newStats, err := ds.Statistics(ctx)