1import * as ApiGatewayV2 from "@/AWS/ApiGatewayV2";
2import * as Lambda from "@/AWS/Lambda";
3import * as Context from "effect/Context";
4import * as Duration from "effect/Duration";
5import * as Effect from "effect/Effect";
6import * as Layer from "effect/Layer";
7import * as HttpServerResponse from "effect/http/HttpServerResponse";
8import path from "pathe";
9
10const main = path.resolve(import.meta.dirname, "ws-handler.ts");
11
12export class WebSocketTestFunction extends Lambda.Function<Lambda.Function>()(
13 "WebSocketTestFunction",
14) {}
15
16export class WsApiAndStage extends Context.Service<
17 WsApiAndStage,
18 {
19 api: ApiGatewayV2.Api;
20 stage: ApiGatewayV2.ApiGatewayV2Stage;
21 }
22>()("WsApiAndStage") {}
23
24export const WsApiAndStageLive = Layer.effect(
25 WsApiAndStage,
26 Effect.gen(function* () {
27 const api = yield* ApiGatewayV2.Api("WsTestApi", {
28 protocolType: "WEBSOCKET",
29 routeSelectionExpression: "$request.body.action",
30 });
31 const stage = yield* ApiGatewayV2.Stage("WsTestStage", {
32 api,
33 stageName: "test",
34 autoDeploy: true,
35 });
36 return { api, stage };
37 }),
38);
39
40/**
41 * Fixture Lambda for the WebSocket e2e test:
42 *
43 * - `$connect` → accept the connection
44 * - `$disconnect` → no-op
45 * - `$default` → echo the received frame back to the sender via the
46 * ManageConnections binding (`postToConnection`)
47 *
48 * The function URL exposes the stage's `wss://` URL so the test can
49 * discover it after deploy.
50 */
51export default WebSocketTestFunction.make(
52 {
53 main,
54 functionUrl: true,
55 timeout: Duration.seconds(30),
56 },
57 Effect.gen(function* () {
58 const { api, stage } = yield* WsApiAndStage;
59 const connections = yield* ApiGatewayV2.ManageConnections(stage);
60 const WsUrl = yield* stage.invokeUrl;
61 const CallbackUrl = yield* stage.callbackUrl;
62
63 yield* ApiGatewayV2.onWebSocketRoute(api, { routeKey: "$connect" }, () =>
64 Effect.succeed({ statusCode: 200 }),
65 );
66
67 yield* ApiGatewayV2.onWebSocketRoute(
68 api,
69 { routeKey: "$disconnect" },
70 () => Effect.void,
71 );
72
73 yield* ApiGatewayV2.onWebSocketRoute(
74 api,
75 { routeKey: "$default" },
76 (event) =>
77 connections
78 .postToConnection({
79 ConnectionId: event.requestContext.connectionId,
80 Data: `echo:${event.body ?? ""}`,
81 })
82 .pipe(
83 Effect.asVoid,
84 // The peer may disconnect between send and receive — not a
85 // failure for the echo fixture.
86 Effect.catchTag("GoneException", () => Effect.void),
87 Effect.orDie,
88 ),
89 );
90
91 return {
92 fetch: Effect.gen(function* () {
93 return yield* HttpServerResponse.json({
94 wsUrl: yield* WsUrl,
95 callbackUrl: yield* CallbackUrl,
96 });
97 }).pipe(Effect.orDie),
98 };
99 }).pipe(
100 Effect.provide(
101 Layer.provideMerge(
102 Layer.mergeAll(
103 Lambda.WebSocketEventSource,
104 ApiGatewayV2.ManageConnectionsHttp,
105 ),
106 WsApiAndStageLive,
107 ),
108 ),
109 ),
110);