chore(coderd/workspaceusage): add integration-style test with multipl…

…e instances
coder · johnstcn · Mar 20, 2024 · Mar 18, 2024 · Mar 18, 2024 · Mar 18, 2024
commit 591e1abffe1f47651ff6fcd1a204575ddacd4498
diff --git a/coderd/workspaceusage/tracker_test.go b/coderd/workspaceusage/tracker_test.go
@@ -14,10 +14,15 @@ import (
 
 	"cdr.dev/slog"
 	"cdr.dev/slog/sloggers/slogtest"
+	"github.com/coder/coder/v2/coderd/coderdtest"
 	"github.com/coder/coder/v2/coderd/database"
+	"github.com/coder/coder/v2/coderd/database/dbfake"
 	"github.com/coder/coder/v2/coderd/database/dbmock"
+	"github.com/coder/coder/v2/coderd/database/dbtestutil"
 	"github.com/coder/coder/v2/coderd/database/dbtime"
 	"github.com/coder/coder/v2/coderd/workspaceusage"
+	"github.com/coder/coder/v2/codersdk"
+	"github.com/coder/coder/v2/testutil"
 )
 
 func TestTracker(t *testing.T) {
@@ -105,6 +110,114 @@ func TestTracker(t *testing.T) {
 	require.Panics(t, wut.Loop)
 }
 
+// This test performs a more 'integration-style' test with multiple instances.
+func TestTracker_MultipleInstances(t *testing.T) {
+	t.Parallel()
+	if !dbtestutil.WillUsePostgres() {
+		t.Skip("this test only makes sense with postgres")
+	}
+
+	// Given we have two coderd instances connected to the same database
+	var (
+		ctx      = testutil.Context(t, testutil.WaitLong)
+		db, ps   = dbtestutil.NewDB(t)
+		wuTickA  = make(chan time.Time)
+		wuFlushA = make(chan int, 1)
+		wutA     = workspaceusage.New(db, workspaceusage.WithFlushChannel(wuFlushA), workspaceusage.WithTickChannel(wuTickA))
+		wuTickB  = make(chan time.Time)
+		wuFlushB = make(chan int, 1)
+		wutB     = workspaceusage.New(db, workspaceusage.WithFlushChannel(wuFlushB), workspaceusage.WithTickChannel(wuTickB))
+		clientA  = coderdtest.New(t, &coderdtest.Options{
+			WorkspaceUsageTracker: wutA,
+			Database:              db,
+			Pubsub:                ps,
+		})
+		clientB = coderdtest.New(t, &coderdtest.Options{
+			WorkspaceUsageTracker: wutB,
+			Database:              db,
+			Pubsub:                ps,
+		})
+		owner = coderdtest.CreateFirstUser(t, clientA)
+		now   = dbtime.Now()
+	)
+
+	clientB.SetSessionToken(clientA.SessionToken())
+
+	// Create a number of workspaces
+	numWorkspaces := 10
+	w := make([]dbfake.WorkspaceResponse, numWorkspaces)
+	for i := 0; i < numWorkspaces; i++ {
+		wr := dbfake.WorkspaceBuild(t, db, database.Workspace{
+			OwnerID:        owner.UserID,
+			OrganizationID: owner.OrganizationID,
+			LastUsedAt:     now,
+		}).WithAgent().Do()
+		w[i] = wr
+	}
+
+	// Use client A to update LastUsedAt of the first three
+	require.NoError(t, clientA.PostWorkspaceUsage(ctx, w[0].Workspace.ID))
+	require.NoError(t, clientA.PostWorkspaceUsage(ctx, w[1].Workspace.ID))
+	require.NoError(t, clientA.PostWorkspaceUsage(ctx, w[2].Workspace.ID))
+	// Use client B to update LastUsedAt of the next three
+	require.NoError(t, clientB.PostWorkspaceUsage(ctx, w[3].Workspace.ID))
+	require.NoError(t, clientB.PostWorkspaceUsage(ctx, w[4].Workspace.ID))
+	require.NoError(t, clientB.PostWorkspaceUsage(ctx, w[5].Workspace.ID))
+	// The next two will have updated from both instances
+	require.NoError(t, clientA.PostWorkspaceUsage(ctx, w[6].Workspace.ID))
+	require.NoError(t, clientB.PostWorkspaceUsage(ctx, w[6].Workspace.ID))
+	require.NoError(t, clientA.PostWorkspaceUsage(ctx, w[7].Workspace.ID))
+	require.NoError(t, clientB.PostWorkspaceUsage(ctx, w[7].Workspace.ID))
+	// The last two will not report any usage.
+
+	// Tick both with different times and wait for both flushes to complete
+	nowA := now.Add(time.Minute)
+	nowB := now.Add(2 * time.Minute)
+	var wg sync.WaitGroup
+	var flushedA, flushedB int
+	wg.Add(1)
+	go func() {
+		defer wg.Done()
+		wuTickA <- nowA
+		flushedA = <-wuFlushA
+	}()
+	wg.Add(1)
+	go func() {
+		defer wg.Done()
+		wuTickB <- nowB
+		flushedB = <-wuFlushB
+	}()
+	wg.Wait()
+
+	// We expect 5 flushed IDs each
+	require.Equal(t, 5, flushedA)
+	require.Equal(t, 5, flushedB)
+
+	// Fetch updated workspaces
+	updated := make([]codersdk.Workspace, numWorkspaces)
+	for i := 0; i < numWorkspaces; i++ {
+		ws, err := clientA.Workspace(ctx, w[i].Workspace.ID)
+		require.NoError(t, err)
+		updated[i] = ws
+	}
+	// We expect the first three to have the timestamp of flushA
+	require.Equal(t, nowA.UTC(), updated[0].LastUsedAt.UTC())
+	require.Equal(t, nowA.UTC(), updated[1].LastUsedAt.UTC())
+	require.Equal(t, nowA.UTC(), updated[2].LastUsedAt.UTC())
+	// We expect the next three to have the timestamp of flushB
+	require.Equal(t, nowB.UTC(), updated[3].LastUsedAt.UTC())
+	require.Equal(t, nowB.UTC(), updated[4].LastUsedAt.UTC())
+	require.Equal(t, nowB.UTC(), updated[5].LastUsedAt.UTC())
+	// The next two should have the timestamp of flushB as it is newer than flushA
+	require.Equal(t, nowB.UTC(), updated[6].LastUsedAt.UTC())
+	require.Equal(t, nowB.UTC(), updated[7].LastUsedAt.UTC())
+	// And the last two should be untouched
+	require.Equal(t, w[8].Workspace.LastUsedAt.UTC(), updated[8].LastUsedAt.UTC())
+	require.Equal(t, w[8].Workspace.LastUsedAt.UTC(), updated[8].LastUsedAt.UTC())
+	require.Equal(t, w[9].Workspace.LastUsedAt.UTC(), updated[9].LastUsedAt.UTC())
+	require.Equal(t, w[9].Workspace.LastUsedAt.UTC(), updated[9].LastUsedAt.UTC())
+}
+
 func TestMain(m *testing.M) {
 	goleak.VerifyTestMain(m)
 }