1import * as AWS from "@/AWS";
2import * as Core from "@/Test/Core";
3import * as Test from "@/Test/Alchemy";
4import * as eventbridge from "@distilled.cloud/aws/eventbridge";
5import { describe, expect } from "alchemy-test";
6import * as Data from "effect/Data";
7import * as Effect from "effect/Effect";
8import * as Schedule from "effect/Schedule";
9import * as HttpClient from "effect/http/HttpClient";
10import * as HttpClientRequest from "effect/http/HttpClientRequest";
11import SageMakerTestFunctionLive, { SageMakerTestFunction } from "./handler";
12
13const testOptions = { providers: AWS.providers() };
14const { test, beforeAll, afterAll } = Test.make(testOptions);
15const sharedStack = Core.scratchStack(testOptions, "SageMakerBindings");
16
17const readinessPolicy = Schedule.max([
18 Schedule.fixed("2 seconds"),
19 Schedule.recurs(75),
20]);
21
22let baseUrl: string;
23let functionArn: string;
24
25class TransientUpstream extends Data.TaggedError("TransientUpstream")<{
26 readonly status: number;
27 readonly body: string;
28}> {}
29
30// Retry only transient 5xx (cold re-inits under parallel-suite load);
31// genuine 4xx/assertion failures surface immediately.
32const send = (request: HttpClientRequest.HttpClientRequest) =>
33 HttpClient.execute(request).pipe(
34 Effect.flatMap((response) =>
35 response.status >= 500
36 ? response.text.pipe(
37 Effect.flatMap((body) =>
38 Effect.fail(
39 new TransientUpstream({ status: response.status, body }),
40 ),
41 ),
42 )
43 : Effect.succeed(response),
44 ),
45 Effect.retry({
46 while: (e) => e._tag === "TransientUpstream",
47 schedule: Schedule.max([
48 Schedule.exponential("500 millis"),
49 Schedule.recurs(6),
50 ]),
51 }),
52 );
53
54const postJson = (path: string, body: unknown) =>
55 send(
56 HttpClientRequest.post(`${baseUrl}${path}`).pipe(
57 HttpClientRequest.bodyJsonUnsafe(body),
58 ),
59 ).pipe(Effect.flatMap((r) => r.json));
60
61const getJson = (path: string) =>
62 send(HttpClientRequest.get(`${baseUrl}${path}`)).pipe(
63 Effect.flatMap((r) => r.json),
64 );
65
66describe(
67 "SageMaker FeatureStore Bindings",
68 {
69 tags: [
70 "provider:aws",
71 "provider:aws:batch",
72 "provider:aws:lambda",
73 "provider:aws:sagemaker",
74 "live",
75 ],
76 },
77 () => {
78 beforeAll(
79 Effect.gen(function* () {
80 yield* Effect.logInfo("SageMaker bindings: destroying previous stack");
81 yield* sharedStack.destroy();
82
83 yield* Effect.logInfo("SageMaker bindings: deploying fixture");
84 const attrs = yield* sharedStack.deploy(
85 Effect.gen(function* () {
86 return yield* SageMakerTestFunction;
87 }).pipe(Effect.provide(SageMakerTestFunctionLive)),
88 );
89
90 expect(attrs.functionUrl).toBeTruthy();
91 baseUrl = attrs.functionUrl!.replace(/\/+$/, "");
92 functionArn = attrs.functionArn;
93
94 yield* Effect.logInfo(
95 `SageMaker bindings: probing readiness at ${baseUrl}/health`,
96 );
97 yield* HttpClient.get(`${baseUrl}/health`).pipe(
98 Effect.flatMap((response) =>
99 response.status === 200
100 ? Effect.succeed(response)
101 : Effect.fail(
102 new Error(`Function not ready: ${response.status}`),
103 ),
104 ),
105 Effect.retry({ schedule: readinessPolicy }),
106 );
107 }),
108 { timeout: 480_000 },
109 );
110
111 afterAll(sharedStack.destroy(), { timeout: 300_000 });
112
113 describe("binding registration", () => {
114 test.provider("all six capabilities initialize in the runtime", () =>
115 Effect.gen(function* () {
116 const response = (yield* getJson("/bindings")) as { bound: string[] };
117 expect(response.bound.sort()).toEqual(
118 [
119 "batchGetRecord",
120 "batchWriteRecord",
121 "deleteRecord",
122 "getRecord",
123 "listRecords",
124 "putRecord",
125 ].sort(),
126 );
127 }),
128 );
129 });
130
131 describe("PutRecord", () => {
132 test.provider(
133 "writes a record to the online store",
134 () =>
135 Effect.gen(function* () {
136 const body = (yield* postJson("/put-record", {
137 userId: "user-put-1",
138 clicks: 7,
139 })) as { success: boolean };
140 expect(body.success).toBe(true);
141 }),
142 { timeout: 60_000 },
143 );
144 });
145
146 describe("GetRecord", () => {
147 test.provider(
148 "reads back a written record",
149 () =>
150 Effect.gen(function* () {
151 yield* postJson("/put-record", {
152 userId: "user-roundtrip-1",
153 clicks: 42,
154 });
155
156 const body = (yield* getJson(
157 "/get-record?userId=user-roundtrip-1",
158 )) as {
159 record: { FeatureName?: string; ValueAsString?: string }[];
160 };
161 const byName = Object.fromEntries(
162 body.record.map((f) => [f.FeatureName, f.ValueAsString]),
163 );
164 expect(byName.user_id).toBe("user-roundtrip-1");
165 expect(byName.clicks).toBe("42");
166 }),
167 { timeout: 60_000 },
168 );
169
170 test.provider(
171 "returns an empty record for an unknown identifier",
172 () =>
173 Effect.gen(function* () {
174 const body = (yield* getJson(
175 "/get-record?userId=user-never-written",
176 )) as { record: unknown[] };
177 expect(body.record).toEqual([]);
178 }),
179 { timeout: 60_000 },
180 );
181 });
182
183 describe("DeleteRecord", () => {
184 test.provider(
185 "soft-deletes a written record (GetRecord no longer returns it)",
186 () =>
187 Effect.gen(function* () {
188 yield* postJson("/put-record", {
189 userId: "user-delete-1",
190 clicks: 1,
191 });
192 yield* postJson("/delete-record", { userId: "user-delete-1" });
193
194 // Online-store deletes are read-after-write consistent in practice;
195 // allow a short bounded window regardless.
196 const body = yield* getJson(
197 "/get-record?userId=user-delete-1",
198 ).pipe(
199 Effect.repeat({
200 schedule: Schedule.spaced("2 seconds"),
201 until: (b): boolean =>
202 (b as { record: unknown[] }).record.length === 0,
203 times: 8,
204 }),
205 );
206 expect((body as { record: unknown[] }).record).toEqual([]);
207 }),
208 { timeout: 60_000 },
209 );
210 });
211
212 describe("BatchWriteRecord + BatchGetRecord", () => {
213 test.provider(
214 "bulk-writes records and reads them back in one batch",
215 () =>
216 Effect.gen(function* () {
217 const write = (yield* postJson("/batch-write-record", {
218 records: [
219 { userId: "user-batch-1", clicks: 11 },
220 { userId: "user-batch-2", clicks: 22 },
221 ],
222 })) as { errors: unknown[]; unprocessed: unknown[] };
223 expect(write.errors).toEqual([]);
224 expect(write.unprocessed).toEqual([]);
225
226 const read = (yield* postJson("/batch-get-record", {
227 userIds: ["user-batch-1", "user-batch-2", "user-batch-missing"],
228 })) as {
229 records: {
230 userId: string;
231 record: { FeatureName?: string; ValueAsString?: string }[];
232 }[];
233 errors: unknown[];
234 };
235 expect(read.errors).toEqual([]);
236 const byUser = Object.fromEntries(
237 read.records.map((r) => [
238 r.userId,
239 Object.fromEntries(
240 r.record.map((f) => [f.FeatureName, f.ValueAsString]),
241 ),
242 ]),
243 );
244 expect(byUser["user-batch-1"]?.clicks).toBe("11");
245 expect(byUser["user-batch-2"]?.clicks).toBe("22");
246 expect(byUser["user-batch-missing"]).toBeUndefined();
247 }),
248 { timeout: 60_000 },
249 );
250 });
251
252 describe("ListRecords", () => {
253 test.provider(
254 "lists the identifiers of stored records",
255 () =>
256 Effect.gen(function* () {
257 yield* postJson("/put-record", {
258 userId: "user-list-1",
259 clicks: 3,
260 });
261
262 const body = yield* getJson("/list-records").pipe(
263 Effect.repeat({
264 schedule: Schedule.spaced("2 seconds"),
265 until: (b): boolean =>
266 (b as { identifiers: string[] }).identifiers.includes(
267 "user-list-1",
268 ),
269 times: 8,
270 }),
271 );
272 expect((body as { identifiers: string[] }).identifiers).toContain(
273 "user-list-1",
274 );
275 }),
276 { timeout: 60_000 },
277 );
278 });
279
280 describe(
281 "consumeSageMakerEvents",
282 { tags: ["provider:aws:eventbridge"] },
283 () => {
284 test.provider(
285 "the deploy created an EventBridge rule targeting the function",
286 () =>
287 Effect.gen(function* () {
288 // Out-of-band via distilled: the fixture's consumeSageMakerEvents
289 // must have materialized as a rule on the default bus with the
290 // Lambda as target.
291 const { RuleNames } = yield* eventbridge.listRuleNamesByTarget({
292 TargetArn: functionArn,
293 });
294 expect((RuleNames ?? []).length).toBeGreaterThanOrEqual(1);
295 }),
296 );
297 },
298 );
299 },
300);