1import * as Batch from "@/AWS/Batch";
2import * as IAM from "@/AWS/IAM";
3import * as Lambda from "@/AWS/Lambda";
4import type * as batch from "@distilled.cloud/aws/batch";
5import * as Duration from "effect/Duration";
6import * as Effect from "effect/Effect";
7import * as Layer from "effect/Layer";
8import * as Stream from "effect/Stream";
9import { HttpServerRequest } from "effect/http/HttpServerRequest";
10import * as HttpServerResponse from "effect/http/HttpServerResponse";
11import path from "pathe";
12import { BatchTestNetwork } from "./TestNetwork.ts";
13
14const main = path.resolve(import.meta.dirname, "handler.ts");
15
16export class BatchTestFunction extends Lambda.Function<Lambda.Function>()(
17 "BatchTestFunction",
18) {}
19
20export default BatchTestFunction.make(
21 {
22 main,
23 functionUrl: true,
24 // submit/describe fan out SDK calls — AWS's 3s default intermittently
25 // times out under cold starts.
26 timeout: Duration.seconds(30),
27 },
28 Effect.gen(function* () {
29 const unmanagedServiceRole = process.env.AWS_TEST_SLOW
30 ? undefined
31 : yield* IAM.Role("BatchServiceRole", {
32 assumeRolePolicyDocument: {
33 Version: "2012-10-17",
34 Statement: [
35 {
36 Effect: "Allow",
37 Principal: { Service: "batch.amazonaws.com" },
38 Action: ["sts:AssumeRole"],
39 },
40 ],
41 },
42 managedPolicyArns: [
43 "arn:aws:iam::aws:policy/service-role/AWSBatchServiceRole",
44 ],
45 });
46
47 // Fast capability tests only need jobs to reach RUNNABLE, so use an
48 // unmanaged CE without VPC resources. The slow SUCCEEDED round-trip keeps
49 // the managed Fargate path and its stack-owned network.
50 const computeEnvironment = process.env.AWS_TEST_SLOW
51 ? yield* Effect.gen(function* () {
52 const network = yield* BatchTestNetwork;
53 return yield* Batch.ComputeEnvironment("TestCE", {
54 subnets: network.subnetIds,
55 securityGroupIds: network.securityGroupIds,
56 });
57 })
58 : yield* Batch.ComputeEnvironment("TestCE", {
59 managementType: "UNMANAGED",
60 unmanagedvCpus: 4,
61 serviceRole: unmanagedServiceRole!.roleArn,
62 });
63 const queue = yield* Batch.JobQueue("TestQueue", {
64 computeEnvironments: [computeEnvironment.computeEnvironmentArn],
65 });
66 const executionRole = yield* IAM.Role("BatchExecutionRole", {
67 assumeRolePolicyDocument: {
68 Version: "2012-10-17",
69 Statement: [
70 {
71 Effect: "Allow",
72 Principal: { Service: "ecs-tasks.amazonaws.com" },
73 Action: ["sts:AssumeRole"],
74 },
75 ],
76 },
77 managedPolicyArns: [
78 "arn:aws:iam::aws:policy/service-role/AmazonECSTaskExecutionRolePolicy",
79 ],
80 });
81 const jobDefinition = yield* Batch.JobDefinition("EchoJob", {
82 image: "public.ecr.aws/docker/library/busybox:latest",
83 command: ["echo", "hello-from-batch"],
84 executionRoleArn: executionRole.roleArn,
85 platformCapabilities: process.env.AWS_TEST_SLOW ? ["FARGATE"] : ["EC2"],
86 vcpus: process.env.AWS_TEST_SLOW ? 0.25 : 1,
87 timeout: Duration.minutes(5),
88 });
89
90 const submitJob = yield* Batch.SubmitJob(queue, jobDefinition);
91 const describeJobs = yield* Batch.DescribeJobs(queue);
92 const terminateJob = yield* Batch.TerminateJob(queue);
93 const cancelJob = yield* Batch.CancelJob(queue);
94 const listJobs = yield* Batch.ListJobs(queue);
95 const getJobQueueSnapshot = yield* Batch.GetJobQueueSnapshot(queue);
96
97 // Deploy-time: creates the EventBridge rule (default bus, source
98 // aws.batch) targeting this Function. Runtime firing rides on real job
99 // submissions from the suite; the test verifies the rule deploys.
100 yield* Batch.consumeJobEvents({ kinds: ["job-state"] }, (events) =>
101 Stream.runForEach(events, (event) =>
102 Effect.log(
103 `batch job event: ${event.detail.jobId} -> ${event.detail.status}`,
104 ),
105 ),
106 );
107
108 return {
109 fetch: Effect.gen(function* () {
110 const request = yield* HttpServerRequest;
111 const url = new URL(request.originalUrl);
112 const pathname = url.pathname;
113
114 if (request.method === "POST" && pathname === "/submit") {
115 const body = (yield* request.json) as unknown as {
116 jobName: string;
117 };
118 const result = yield* submitJob({ jobName: body.jobName });
119 return yield* HttpServerResponse.json({
120 jobId: result.jobId,
121 jobName: result.jobName,
122 jobArn: result.jobArn,
123 });
124 }
125
126 if (request.method === "GET" && pathname === "/status") {
127 const jobId = url.searchParams.get("jobId")!;
128 const result = yield* describeJobs({ jobs: [jobId] });
129 const job = result.jobs?.[0];
130 return yield* HttpServerResponse.json({
131 status: job?.status,
132 statusReason: job?.statusReason,
133 jobQueue: job?.jobQueue,
134 });
135 }
136
137 if (request.method === "POST" && pathname === "/terminate") {
138 const body = (yield* request.json) as unknown as {
139 jobId: string;
140 reason: string;
141 };
142 yield* terminateJob({ jobId: body.jobId, reason: body.reason });
143 return yield* HttpServerResponse.json({ terminated: true });
144 }
145
146 if (request.method === "POST" && pathname === "/cancel") {
147 const body = (yield* request.json) as unknown as {
148 jobId: string;
149 reason: string;
150 };
151 yield* cancelJob({ jobId: body.jobId, reason: body.reason });
152 return yield* HttpServerResponse.json({ cancelled: true });
153 }
154
155 if (request.method === "GET" && pathname === "/jobs") {
156 const jobStatus = url.searchParams.get("jobStatus") ?? undefined;
157 const result = yield* listJobs(
158 jobStatus ? { jobStatus: jobStatus as batch.JobStatus } : {},
159 );
160 return yield* HttpServerResponse.json({
161 jobs: (result.jobSummaryList ?? []).map((job) => ({
162 jobId: job.jobId,
163 jobName: job.jobName,
164 status: job.status,
165 })),
166 });
167 }
168
169 if (request.method === "GET" && pathname === "/snapshot") {
170 const result = yield* getJobQueueSnapshot();
171 return yield* HttpServerResponse.json({
172 jobs: (result.frontOfQueue?.jobs ?? []).map((job) => job.jobArn),
173 });
174 }
175
176 return yield* HttpServerResponse.json(
177 { error: "Not found", method: request.method, pathname },
178 { status: 404 },
179 );
180 }).pipe(Effect.orDie),
181 };
182 }).pipe(
183 Effect.provide(
184 Layer.mergeAll(
185 Lambda.EventSource,
186 Batch.SubmitJobHttp,
187 Batch.DescribeJobsHttp,
188 Batch.TerminateJobHttp,
189 Batch.CancelJobHttp,
190 Batch.ListJobsHttp,
191 Batch.GetJobQueueSnapshotHttp,
192 ),
193 ),
194 ),
195);