Skip to content

Commit 66d1916

Browse files
threepointonethreepointone
committed
Refactor readonly-connections tests for type safety
Improved type safety in readonly-connections.test.ts by introducing explicit message interfaces and type guards, replacing 'any' with specific types in test helpers and assertions. Also updated TestReadonlyAgent to ignore unused connection parameter in shouldConnectionBeReadonly. These changes enhance code reliability and maintainability in the test suite.
1 parent b2e1e21 commit 66d1916

File tree

2 files changed

+89
-60
lines changed

2 files changed

+89
-60
lines changed

‎packages/agents/src/tests/readonly-connections.test.ts‎

Lines changed: 88 additions & 59 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,35 @@ declare module "cloudflare:test" {
77
interface ProvidedEnv extends Env {}
88
}
99

10+
// Message types used in tests
11+
interface StateMessage {
12+
type: MessageType.CF_AGENT_STATE;
13+
state: { count?: number };
14+
}
15+
16+
interface StateErrorMessage {
17+
type: MessageType.CF_AGENT_STATE_ERROR;
18+
error: string;
19+
}
20+
21+
interface RpcMessage {
22+
type: MessageType.RPC;
23+
id: string;
24+
success?: boolean;
25+
result?: unknown;
26+
}
27+
28+
type TestMessage = StateMessage | StateErrorMessage | RpcMessage;
29+
30+
function isTestMessage(data: unknown): data is TestMessage {
31+
return (
32+
typeof data === "object" &&
33+
data !== null &&
34+
"type" in data &&
35+
typeof (data as TestMessage).type === "string"
36+
);
37+
}
38+
1039
async function connectWS(path: string) {
1140
const ctx = createExecutionContext();
1241
const req = new Request(`http://example.com${path}`, {
@@ -20,17 +49,17 @@ async function connectWS(path: string) {
2049
return { ws, ctx };
2150
}
2251

23-
function waitForMessage(
52+
function waitForMessage<T extends TestMessage>(
2453
ws: WebSocket,
25-
predicate: (data: unknown) => boolean
26-
): Promise<unknown> {
54+
predicate: (data: TestMessage) => boolean
55+
): Promise<T> {
2756
return new Promise((resolve) => {
2857
const handler = (e: MessageEvent) => {
2958
try {
30-
const data = JSON.parse(e.data as string);
31-
if (predicate(data)) {
59+
const data: unknown = JSON.parse(e.data as string);
60+
if (isTestMessage(data) && predicate(data)) {
3261
ws.removeEventListener("message", handler);
33-
resolve(data);
62+
resolve(data as T);
3463
}
3564
} catch {
3665
// Ignore parse errors
@@ -49,9 +78,9 @@ describe("Readonly Connections", () => {
4978
);
5079

5180
// Wait for initial state message
52-
await waitForMessage(
81+
await waitForMessage<StateMessage>(
5382
ws1,
54-
(data: any) => data.type === MessageType.CF_AGENT_STATE
83+
(data) => data.type === MessageType.CF_AGENT_STATE
5584
);
5685

5786
ws1.close();
@@ -61,9 +90,9 @@ describe("Readonly Connections", () => {
6190
`/agents/test-readonly-agent/${room}?readonly=false`
6291
);
6392

64-
await waitForMessage(
93+
await waitForMessage<StateMessage>(
6594
ws2,
66-
(data: any) => data.type === MessageType.CF_AGENT_STATE
95+
(data) => data.type === MessageType.CF_AGENT_STATE
6796
);
6897

6998
// Test passed - connections were established with different readonly query params
@@ -79,15 +108,15 @@ describe("Readonly Connections", () => {
79108
);
80109

81110
// Wait for initial state
82-
await waitForMessage(
111+
await waitForMessage<StateMessage>(
83112
ws,
84-
(data: any) => data.type === MessageType.CF_AGENT_STATE
113+
(data) => data.type === MessageType.CF_AGENT_STATE
85114
);
86115

87116
// Try to update state from readonly connection
88-
const errorPromise = waitForMessage(
117+
const errorPromise = waitForMessage<StateErrorMessage>(
89118
ws,
90-
(data: any) => data.type === MessageType.CF_AGENT_STATE_ERROR
119+
(data) => data.type === MessageType.CF_AGENT_STATE_ERROR
91120
);
92121

93122
ws.send(
@@ -97,7 +126,7 @@ describe("Readonly Connections", () => {
97126
})
98127
);
99128

100-
const errorMsg = (await errorPromise) as any;
129+
const errorMsg = await errorPromise;
101130
expect(errorMsg.type).toBe(MessageType.CF_AGENT_STATE_ERROR);
102131
expect(errorMsg.error).toBe("Connection is readonly");
103132

@@ -111,10 +140,10 @@ describe("Readonly Connections", () => {
111140
);
112141

113142
// Wait for initial state
114-
const initialState = (await waitForMessage(
143+
const initialState = await waitForMessage<StateMessage>(
115144
ws,
116-
(data: any) => data.type === MessageType.CF_AGENT_STATE
117-
)) as any;
145+
(data) => data.type === MessageType.CF_AGENT_STATE
146+
);
118147

119148
expect(initialState.state).toBeDefined();
120149

@@ -130,16 +159,16 @@ describe("Readonly Connections", () => {
130159
);
131160

132161
// Wait for initial state
133-
await waitForMessage(
162+
await waitForMessage<StateMessage>(
134163
ws,
135-
(data: any) => data.type === MessageType.CF_AGENT_STATE
164+
(data) => data.type === MessageType.CF_AGENT_STATE
136165
);
137166

138167
// Call RPC method from readonly connection
139168
const rpcId = Math.random().toString(36).slice(2);
140-
const rpcPromise = waitForMessage(
169+
const rpcPromise = waitForMessage<RpcMessage>(
141170
ws,
142-
(data: any) => data.type === MessageType.RPC && data.id === rpcId
171+
(data) => data.type === MessageType.RPC && data.id === rpcId
143172
);
144173

145174
ws.send(
@@ -151,7 +180,7 @@ describe("Readonly Connections", () => {
151180
})
152181
);
153182

154-
const rpcMsg = (await rpcPromise) as any;
183+
const rpcMsg = await rpcPromise;
155184
expect(rpcMsg.success).toBe(true);
156185
expect(rpcMsg.result).toBe(1);
157186

@@ -167,9 +196,9 @@ describe("Readonly Connections", () => {
167196
);
168197

169198
// Wait for initial state
170-
await waitForMessage(
199+
await waitForMessage<StateMessage>(
171200
ws,
172-
(data: any) => data.type === MessageType.CF_AGENT_STATE
201+
(data) => data.type === MessageType.CF_AGENT_STATE
173202
);
174203

175204
// Call an RPC method to verify connection works
@@ -183,10 +212,10 @@ describe("Readonly Connections", () => {
183212
})
184213
);
185214

186-
const rpcMsg = (await waitForMessage(
215+
const rpcMsg = await waitForMessage<RpcMessage>(
187216
ws,
188-
(data: any) => data.type === MessageType.RPC && data.id === rpcId
189-
)) as any;
217+
(data) => data.type === MessageType.RPC && data.id === rpcId
218+
);
190219

191220
expect(rpcMsg.success).toBe(true);
192221
expect(rpcMsg.result).toBeDefined();
@@ -203,9 +232,9 @@ describe("Readonly Connections", () => {
203232
);
204233

205234
// Wait for connection
206-
await waitForMessage(
235+
await waitForMessage<StateMessage>(
207236
ws,
208-
(data: any) => data.type === MessageType.CF_AGENT_STATE
237+
(data) => data.type === MessageType.CF_AGENT_STATE
209238
);
210239

211240
// Check that readonly status is in the database
@@ -219,14 +248,14 @@ describe("Readonly Connections", () => {
219248
})
220249
);
221250

222-
const dbResult = (await waitForMessage(
251+
const dbResult = await waitForMessage<RpcMessage>(
223252
ws,
224-
(data: any) => data.type === MessageType.RPC && data.id === checkDbId
225-
)) as any;
253+
(data) => data.type === MessageType.RPC && data.id === checkDbId
254+
);
226255

227256
// Should have at least one entry
228257
expect(Array.isArray(dbResult.result)).toBe(true);
229-
expect(dbResult.result.length).toBeGreaterThan(0);
258+
expect((dbResult.result as unknown[]).length).toBeGreaterThan(0);
230259

231260
ws.close();
232261
});
@@ -239,9 +268,9 @@ describe("Readonly Connections", () => {
239268
`/agents/test-readonly-agent/${room}?readonly=true`
240269
);
241270

242-
await waitForMessage(
271+
await waitForMessage<StateMessage>(
243272
ws1,
244-
(data: any) => data.type === MessageType.CF_AGENT_STATE
273+
(data) => data.type === MessageType.CF_AGENT_STATE
245274
);
246275

247276
// Close connection (simulates hibernation scenario)
@@ -255,15 +284,15 @@ describe("Readonly Connections", () => {
255284
`/agents/test-readonly-agent/${room}?readonly=true`
256285
);
257286

258-
await waitForMessage(
287+
await waitForMessage<StateMessage>(
259288
ws2,
260-
(data: any) => data.type === MessageType.CF_AGENT_STATE
289+
(data) => data.type === MessageType.CF_AGENT_STATE
261290
);
262291

263292
// Try state update - should still be blocked
264-
const errorPromise = waitForMessage(
293+
const errorPromise = waitForMessage<StateErrorMessage>(
265294
ws2,
266-
(data: any) => data.type === MessageType.CF_AGENT_STATE_ERROR
295+
(data) => data.type === MessageType.CF_AGENT_STATE_ERROR
267296
);
268297

269298
ws2.send(
@@ -273,7 +302,7 @@ describe("Readonly Connections", () => {
273302
})
274303
);
275304

276-
const errorMsg = (await errorPromise) as any;
305+
const errorMsg = await errorPromise;
277306
expect(errorMsg.type).toBe(MessageType.CF_AGENT_STATE_ERROR);
278307

279308
ws2.close();
@@ -287,9 +316,9 @@ describe("Readonly Connections", () => {
287316
`/agents/test-readonly-agent/${room}?readonly=true`
288317
);
289318

290-
await waitForMessage(
319+
await waitForMessage<StateMessage>(
291320
ws,
292-
(data: any) => data.type === MessageType.CF_AGENT_STATE
321+
(data) => data.type === MessageType.CF_AGENT_STATE
293322
);
294323

295324
// Verify it's in the database
@@ -303,12 +332,12 @@ describe("Readonly Connections", () => {
303332
})
304333
);
305334

306-
const dbResult1 = (await waitForMessage(
335+
const dbResult1 = await waitForMessage<RpcMessage>(
307336
ws,
308-
(data: any) => data.type === MessageType.RPC && data.id === checkDbId1
309-
)) as any;
337+
(data) => data.type === MessageType.RPC && data.id === checkDbId1
338+
);
310339

311-
expect(dbResult1.result.length).toBeGreaterThan(0);
340+
expect((dbResult1.result as unknown[]).length).toBeGreaterThan(0);
312341

313342
// Close connection
314343
ws.close();
@@ -321,9 +350,9 @@ describe("Readonly Connections", () => {
321350
`/agents/test-readonly-agent/${room}?readonly=false`
322351
);
323352

324-
await waitForMessage(
353+
await waitForMessage<StateMessage>(
325354
ws2,
326-
(data: any) => data.type === MessageType.CF_AGENT_STATE
355+
(data) => data.type === MessageType.CF_AGENT_STATE
327356
);
328357

329358
const checkDbId2 = Math.random().toString(36).slice(2);
@@ -336,10 +365,10 @@ describe("Readonly Connections", () => {
336365
})
337366
);
338367

339-
const dbResult2 = (await waitForMessage(
368+
const dbResult2 = await waitForMessage<RpcMessage>(
340369
ws2,
341-
(data: any) => data.type === MessageType.RPC && data.id === checkDbId2
342-
)) as any;
370+
(data) => data.type === MessageType.RPC && data.id === checkDbId2
371+
);
343372

344373
// Old connection should be cleaned up
345374
expect(dbResult2.result).toEqual([]);
@@ -357,15 +386,15 @@ describe("Readonly Connections", () => {
357386
`/agents/test-readonly-agent/${room}?readonly=true`
358387
);
359388

360-
await waitForMessage(
389+
await waitForMessage<StateMessage>(
361390
ws1,
362-
(data: any) => data.type === MessageType.CF_AGENT_STATE
391+
(data) => data.type === MessageType.CF_AGENT_STATE
363392
);
364393

365394
// ws1 (readonly) should not be able to update state
366-
const errorPromise = waitForMessage(
395+
const errorPromise = waitForMessage<StateErrorMessage>(
367396
ws1,
368-
(data: any) => data.type === MessageType.CF_AGENT_STATE_ERROR
397+
(data) => data.type === MessageType.CF_AGENT_STATE_ERROR
369398
);
370399

371400
ws1.send(
@@ -375,7 +404,7 @@ describe("Readonly Connections", () => {
375404
})
376405
);
377406

378-
const errorMsg = (await errorPromise) as any;
407+
const errorMsg = await errorPromise;
379408
expect(errorMsg.error).toBe("Connection is readonly");
380409

381410
ws1.close();
@@ -385,9 +414,9 @@ describe("Readonly Connections", () => {
385414
`/agents/test-readonly-agent/${room}?readonly=false`
386415
);
387416

388-
await waitForMessage(
417+
await waitForMessage<StateMessage>(
389418
ws2,
390-
(data: any) => data.type === MessageType.CF_AGENT_STATE
419+
(data) => data.type === MessageType.CF_AGENT_STATE
391420
);
392421

393422
ws2.close();

‎packages/agents/src/tests/worker.ts‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -666,7 +666,7 @@ export class TestReadonlyAgent extends Agent<Env, { count: number }> {
666666
}> = [];
667667

668668
shouldConnectionBeReadonly(
669-
connection: Connection,
669+
_connection: Connection,
670670
ctx: import("../index").ConnectionContext
671671
): boolean {
672672
// Check query parameter to determine readonly status

0 commit comments

Comments
 (0)