Skip to content

Commit 591e1ab

Browse files
johnstcnjohnstcn
committed
chore(coderd/workspaceusage): add integration-style test with multiple instances
1 parent 8e40efd commit 591e1ab

File tree

1 file changed

+113
-0
lines changed

1 file changed

+113
-0
lines changed

‎coderd/workspaceusage/tracker_test.go

+113
Original file line numberDiff line numberDiff line change
@@ -14,10 +14,15 @@ import (
1414

1515
"cdr.dev/slog"
1616
"cdr.dev/slog/sloggers/slogtest"
17+
"github.com/coder/coder/v2/coderd/coderdtest"
1718
"github.com/coder/coder/v2/coderd/database"
19+
"github.com/coder/coder/v2/coderd/database/dbfake"
1820
"github.com/coder/coder/v2/coderd/database/dbmock"
21+
"github.com/coder/coder/v2/coderd/database/dbtestutil"
1922
"github.com/coder/coder/v2/coderd/database/dbtime"
2023
"github.com/coder/coder/v2/coderd/workspaceusage"
24+
"github.com/coder/coder/v2/codersdk"
25+
"github.com/coder/coder/v2/testutil"
2126
)
2227

2328
func TestTracker(t *testing.T) {
@@ -105,6 +110,114 @@ func TestTracker(t *testing.T) {
105110
require.Panics(t, wut.Loop)
106111
}
107112

113+
// This test performs a more 'integration-style' test with multiple instances.
114+
func TestTracker_MultipleInstances(t *testing.T) {
115+
t.Parallel()
116+
if !dbtestutil.WillUsePostgres() {
117+
t.Skip("this test only makes sense with postgres")
118+
}
119+
120+
// Given we have two coderd instances connected to the same database
121+
var (
122+
ctx = testutil.Context(t, testutil.WaitLong)
123+
db, ps = dbtestutil.NewDB(t)
124+
wuTickA = make(chan time.Time)
125+
wuFlushA = make(chan int, 1)
126+
wutA = workspaceusage.New(db, workspaceusage.WithFlushChannel(wuFlushA), workspaceusage.WithTickChannel(wuTickA))
127+
wuTickB = make(chan time.Time)
128+
wuFlushB = make(chan int, 1)
129+
wutB = workspaceusage.New(db, workspaceusage.WithFlushChannel(wuFlushB), workspaceusage.WithTickChannel(wuTickB))
130+
clientA = coderdtest.New(t, &coderdtest.Options{
131+
WorkspaceUsageTracker: wutA,
132+
Database: db,
133+
Pubsub: ps,
134+
})
135+
clientB = coderdtest.New(t, &coderdtest.Options{
136+
WorkspaceUsageTracker: wutB,
137+
Database: db,
138+
Pubsub: ps,
139+
})
140+
owner = coderdtest.CreateFirstUser(t, clientA)
141+
now = dbtime.Now()
142+
)
143+
144+
clientB.SetSessionToken(clientA.SessionToken())
145+
146+
// Create a number of workspaces
147+
numWorkspaces := 10
148+
w := make([]dbfake.WorkspaceResponse, numWorkspaces)
149+
for i := 0; i < numWorkspaces; i++ {
150+
wr := dbfake.WorkspaceBuild(t, db, database.Workspace{
151+
OwnerID: owner.UserID,
152+
OrganizationID: owner.OrganizationID,
153+
LastUsedAt: now,
154+
}).WithAgent().Do()
155+
w[i] = wr
156+
}
157+
158+
// Use client A to update LastUsedAt of the first three
159+
require.NoError(t, clientA.PostWorkspaceUsage(ctx, w[0].Workspace.ID))
160+
require.NoError(t, clientA.PostWorkspaceUsage(ctx, w[1].Workspace.ID))
161+
require.NoError(t, clientA.PostWorkspaceUsage(ctx, w[2].Workspace.ID))
162+
// Use client B to update LastUsedAt of the next three
163+
require.NoError(t, clientB.PostWorkspaceUsage(ctx, w[3].Workspace.ID))
164+
require.NoError(t, clientB.PostWorkspaceUsage(ctx, w[4].Workspace.ID))
165+
require.NoError(t, clientB.PostWorkspaceUsage(ctx, w[5].Workspace.ID))
166+
// The next two will have updated from both instances
167+
require.NoError(t, clientA.PostWorkspaceUsage(ctx, w[6].Workspace.ID))
168+
require.NoError(t, clientB.PostWorkspaceUsage(ctx, w[6].Workspace.ID))
169+
require.NoError(t, clientA.PostWorkspaceUsage(ctx, w[7].Workspace.ID))
170+
require.NoError(t, clientB.PostWorkspaceUsage(ctx, w[7].Workspace.ID))
171+
// The last two will not report any usage.
172+
173+
// Tick both with different times and wait for both flushes to complete
174+
nowA := now.Add(time.Minute)
175+
nowB := now.Add(2 * time.Minute)
176+
var wg sync.WaitGroup
177+
var flushedA, flushedB int
178+
wg.Add(1)
179+
go func() {
180+
defer wg.Done()
181+
wuTickA <- nowA
182+
flushedA = <-wuFlushA
183+
}()
184+
wg.Add(1)
185+
go func() {
186+
defer wg.Done()
187+
wuTickB <- nowB
188+
flushedB = <-wuFlushB
189+
}()
190+
wg.Wait()
191+
192+
// We expect 5 flushed IDs each
193+
require.Equal(t, 5, flushedA)
194+
require.Equal(t, 5, flushedB)
195+
196+
// Fetch updated workspaces
197+
updated := make([]codersdk.Workspace, numWorkspaces)
198+
for i := 0; i < numWorkspaces; i++ {
199+
ws, err := clientA.Workspace(ctx, w[i].Workspace.ID)
200+
require.NoError(t, err)
201+
updated[i] = ws
202+
}
203+
// We expect the first three to have the timestamp of flushA
204+
require.Equal(t, nowA.UTC(), updated[0].LastUsedAt.UTC())
205+
require.Equal(t, nowA.UTC(), updated[1].LastUsedAt.UTC())
206+
require.Equal(t, nowA.UTC(), updated[2].LastUsedAt.UTC())
207+
// We expect the next three to have the timestamp of flushB
208+
require.Equal(t, nowB.UTC(), updated[3].LastUsedAt.UTC())
209+
require.Equal(t, nowB.UTC(), updated[4].LastUsedAt.UTC())
210+
require.Equal(t, nowB.UTC(), updated[5].LastUsedAt.UTC())
211+
// The next two should have the timestamp of flushB as it is newer than flushA
212+
require.Equal(t, nowB.UTC(), updated[6].LastUsedAt.UTC())
213+
require.Equal(t, nowB.UTC(), updated[7].LastUsedAt.UTC())
214+
// And the last two should be untouched
215+
require.Equal(t, w[8].Workspace.LastUsedAt.UTC(), updated[8].LastUsedAt.UTC())
216+
require.Equal(t, w[8].Workspace.LastUsedAt.UTC(), updated[8].LastUsedAt.UTC())
217+
require.Equal(t, w[9].Workspace.LastUsedAt.UTC(), updated[9].LastUsedAt.UTC())
218+
require.Equal(t, w[9].Workspace.LastUsedAt.UTC(), updated[9].LastUsedAt.UTC())
219+
}
220+
108221
func TestMain(m *testing.M) {
109222
goleak.VerifyTestMain(m)
110223
}

0 commit comments

Comments
 (0)