6
6
"time"
7
7
8
8
"github.com/google/uuid"
9
- "github.com/stretchr/testify/assert"
10
9
"github.com/stretchr/testify/require"
11
10
"golang.org/x/xerrors"
12
11
@@ -44,13 +43,8 @@ func TestPubsubWorkspaceClaimPublisher(t *testing.T) {
44
43
err = publisher .PublishWorkspaceClaim (claim )
45
44
require .NoError (t , err )
46
45
47
- // Verify the message was published
48
- select {
49
- case gotUserID := <- userIDCh :
50
- require .Equal (t , userID , gotUserID )
51
- case <- time .After (testutil .WaitShort ):
52
- t .Fatal ("timeout waiting for claim" )
53
- }
46
+ gotUserID := testutil .TryReceive (testutil .Context (t , testutil .WaitShort ), t , userIDCh )
47
+ require .Equal (t , userID , gotUserID )
54
48
})
55
49
56
50
t .Run ("fail to publish claim" , func (t * testing.T ) {
@@ -66,7 +60,7 @@ func TestPubsubWorkspaceClaimPublisher(t *testing.T) {
66
60
}
67
61
68
62
err := publisher .PublishWorkspaceClaim (claim )
69
- require .Error (t , err )
63
+ require .ErrorContains (t , err , "failed to trigger prebuilt workspace reinitialization" )
70
64
})
71
65
}
72
66
@@ -88,7 +82,7 @@ func TestPubsubWorkspaceClaimListener(t *testing.T) {
88
82
// Channel should be closed immediately due to context cancellation
89
83
select {
90
84
case _ , ok := <- claims :
91
- assert .False (t , ok )
85
+ require .False (t , ok )
92
86
case <- time .After (testutil .WaitShort ):
93
87
t .Fatal ("timeout waiting for closed channel" )
94
88
}
@@ -106,7 +100,7 @@ func TestPubsubWorkspaceClaimListener(t *testing.T) {
106
100
cancelFunc ()
107
101
select {
108
102
case _ , ok := <- claims :
109
- assert .False (t , ok )
103
+ require .False (t , ok )
110
104
case <- time .After (testutil .WaitShort ):
111
105
t .Fatal ("timeout waiting for closed channel" )
112
106
}
@@ -132,9 +126,9 @@ func TestPubsubWorkspaceClaimListener(t *testing.T) {
132
126
// Verify we receive the claim
133
127
select {
134
128
case claim := <- claims :
135
- assert .Equal (t , userID , claim .UserID )
136
- assert .Equal (t , workspaceID , claim .WorkspaceID )
137
- assert .Equal (t , agentsdk .ReinitializeReasonPrebuildClaimed , claim .Reason )
129
+ require .Equal (t , userID , claim .UserID )
130
+ require .Equal (t , workspaceID , claim .WorkspaceID )
131
+ require .Equal (t , agentsdk .ReinitializeReasonPrebuildClaimed , claim .Reason )
138
132
case <- time .After (time .Second ):
139
133
t .Fatal ("timeout waiting for claim" )
140
134
}
@@ -173,8 +167,7 @@ func TestPubsubWorkspaceClaimListener(t *testing.T) {
173
167
listener := prebuilds .NewPubsubWorkspaceClaimListener (ps , slogtest .Make (t , nil ))
174
168
175
169
_ , _ , err := listener .ListenForWorkspaceClaims (context .Background (), uuid .New ())
176
- require .Error (t , err )
177
- assert .Contains (t , err .Error (), "failed to subscribe to prebuild claimed channel" )
170
+ require .ErrorContains (t , err , "failed to subscribe to prebuild claimed channel" )
178
171
})
179
172
}
180
173
0 commit comments