@@ -59,10 +59,145 @@ func (fakeAgentProvider) Close() error {
59
59
return nil
60
60
}
61
61
62
+ type channelCloser struct {
63
+ closeFn func ()
64
+ }
65
+
66
+ func (c * channelCloser )Close ()error {
67
+ c .closeFn ()
68
+ return nil
69
+ }
70
+
62
71
func TestWatchAgentContainers (t * testing.T ) {
63
72
t .Parallel ()
64
73
65
- t .Run ("WebSocketClosesProperly" ,func (t * testing.T ) {
74
+ t .Run ("CoderdWebSocketCanHandleClientClosing" ,func (t * testing.T ) {
75
+ t .Parallel ()
76
+
77
+ // This test ensures that the agent containers `/watch` websocket can gracefully
78
+ // handle the client websocket closing. This test was created in
79
+ // response to this issue: https://github.com/coder/coder/issues/19449
80
+
81
+ var (
82
+ ctx = testutil .Context (t ,testutil .WaitLong )
83
+ logger = slogtest .Make (t ,& slogtest.Options {IgnoreErrors :true }).Leveled (slog .LevelDebug ).Named ("coderd" )
84
+
85
+ mCtrl = gomock .NewController (t )
86
+ mDB = dbmock .NewMockStore (mCtrl )
87
+ mCoordinator = tailnettest .NewMockCoordinator (mCtrl )
88
+ mAgentConn = agentconnmock .NewMockAgentConn (mCtrl )
89
+
90
+ fAgentProvider = fakeAgentProvider {
91
+ agentConn :func (ctx context.Context ,agentID uuid.UUID ) (_ workspacesdk.AgentConn ,release func (),_ error ) {
92
+ return mAgentConn ,func () {},nil
93
+ },
94
+ }
95
+
96
+ workspaceID = uuid .New ()
97
+ agentID = uuid .New ()
98
+ resourceID = uuid .New ()
99
+ jobID = uuid .New ()
100
+ buildID = uuid .New ()
101
+
102
+ containersCh = make (chan codersdk.WorkspaceAgentListContainersResponse )
103
+
104
+ r = chi .NewMux ()
105
+
106
+ api = API {
107
+ ctx :ctx ,
108
+ Options :& Options {
109
+ AgentInactiveDisconnectTimeout :testutil .WaitShort ,
110
+ Database :mDB ,
111
+ Logger :logger ,
112
+ DeploymentValues :& codersdk.DeploymentValues {},
113
+ TailnetCoordinator :tailnettest .NewFakeCoordinator (),
114
+ },
115
+ }
116
+ )
117
+
118
+ var tailnetCoordinator tailnet.Coordinator = mCoordinator
119
+ api .TailnetCoordinator .Store (& tailnetCoordinator )
120
+ api .agentProvider = fAgentProvider
121
+
122
+ // Setup: Allow `ExtractWorkspaceAgentParams` to complete.
123
+ mDB .EXPECT ().GetWorkspaceAgentByID (gomock .Any (),agentID ).Return (database.WorkspaceAgent {
124
+ ID :agentID ,
125
+ ResourceID :resourceID ,
126
+ LifecycleState :database .WorkspaceAgentLifecycleStateReady ,
127
+ FirstConnectedAt : sql.NullTime {Valid :true ,Time :dbtime .Now ()},
128
+ LastConnectedAt : sql.NullTime {Valid :true ,Time :dbtime .Now ()},
129
+ },nil )
130
+ mDB .EXPECT ().GetWorkspaceResourceByID (gomock .Any (),resourceID ).Return (database.WorkspaceResource {
131
+ ID :resourceID ,
132
+ JobID :jobID ,
133
+ },nil )
134
+ mDB .EXPECT ().GetProvisionerJobByID (gomock .Any (),jobID ).Return (database.ProvisionerJob {
135
+ ID :jobID ,
136
+ Type :database .ProvisionerJobTypeWorkspaceBuild ,
137
+ },nil )
138
+ mDB .EXPECT ().GetWorkspaceBuildByJobID (gomock .Any (),jobID ).Return (database.WorkspaceBuild {
139
+ WorkspaceID :workspaceID ,
140
+ ID :buildID ,
141
+ },nil )
142
+
143
+ // And: Allow `db2dsk.WorkspaceAgent` to complete.
144
+ mCoordinator .EXPECT ().Node (gomock .Any ()).Return (nil )
145
+
146
+ // And: Allow `WatchContainers` to be called, returing our `containersCh` channel.
147
+ mAgentConn .EXPECT ().WatchContainers (gomock .Any (),gomock .Any ()).
148
+ DoAndReturn (func (_ context.Context ,_ slog.Logger ) (<- chan codersdk.WorkspaceAgentListContainersResponse , io.Closer ,error ) {
149
+ return containersCh ,& channelCloser {closeFn :func () {
150
+ close (containersCh )
151
+ }},nil
152
+ })
153
+
154
+ // And: We mount the HTTP Handler
155
+ r .With (httpmw .ExtractWorkspaceAgentParam (mDB )).
156
+ Get ("/workspaceagents/{workspaceagent}/containers/watch" ,api .watchWorkspaceAgentContainers )
157
+
158
+ // Given: We create the HTTP server
159
+ srv := httptest .NewServer (r )
160
+ defer srv .Close ()
161
+
162
+ // And: Dial the WebSocket
163
+ wsURL := strings .Replace (srv .URL ,"http://" ,"ws://" ,1 )
164
+ conn ,resp ,err := websocket .Dial (ctx ,fmt .Sprintf ("%s/workspaceagents/%s/containers/watch" ,wsURL ,agentID ),nil )
165
+ require .NoError (t ,err )
166
+ if resp .Body != nil {
167
+ defer resp .Body .Close ()
168
+ }
169
+
170
+ // And: Create a streaming decoder
171
+ decoder := wsjson .NewDecoder [codersdk.WorkspaceAgentListContainersResponse ](conn ,websocket .MessageText ,logger )
172
+ defer decoder .Close ()
173
+ decodeCh := decoder .Chan ()
174
+
175
+ // And: We can successfully send through the channel.
176
+ testutil .RequireSend (ctx ,t ,containersCh , codersdk.WorkspaceAgentListContainersResponse {
177
+ Containers : []codersdk.WorkspaceAgentContainer {{
178
+ ID :"test-container-id" ,
179
+ }},
180
+ })
181
+
182
+ // And: Receive the data.
183
+ containerResp := testutil .RequireReceive (ctx ,t ,decodeCh )
184
+ require .Len (t ,containerResp .Containers ,1 )
185
+ require .Equal (t ,"test-container-id" ,containerResp .Containers [0 ].ID )
186
+
187
+ // When: We close the WebSocket
188
+ conn .Close (websocket .StatusNormalClosure ,"test closing connection" )
189
+
190
+ // Then: We expect `containersCh` to be closed.
191
+ select {
192
+ case <- ctx .Done ():
193
+ t .Fail ()
194
+
195
+ case _ ,ok := <- containersCh :
196
+ require .False (t ,ok ,"channel is expected to be closed" )
197
+ }
198
+ })
199
+
200
+ t .Run ("CoderdWebSocketCanHandleAgentClosing" ,func (t * testing.T ) {
66
201
t .Parallel ()
67
202
68
203
// This test ensures that the agent containers `/watch` websocket can gracefully