@@ -20,13 +20,26 @@ function createService() {
2020 snapshot,
2121 } as unknown as ComputeWorkloadManager ;
2222
23+ const submitSuspendCompletion = vi . fn ( async ( ) => ( { success : true } ) ) ;
24+
2325 const service = new ComputeSnapshotService ( {
2426 computeManager,
25- workerClient : { } as SupervisorHttpClient ,
27+ workerClient : { submitSuspendCompletion } as unknown as SupervisorHttpClient ,
2628 wideEventOpts : { service : "supervisor-test" , env : { } , enabled : false } ,
29+ snapshotCallbackSecret : "test-secret" ,
2730 } ) ;
2831
29- return { service, snapshot } ;
32+ return { service, snapshot, submitSuspendCompletion } ;
33+ }
34+
35+ function dispatchedMetadata ( snapshot : {
36+ mock : { calls : Array < Array < { metadata ?: Record < string , string > } > > } ;
37+ } ) {
38+ const metadata = snapshot . mock . calls [ 0 ] ?. [ 0 ] ?. metadata ;
39+ if ( ! metadata ) {
40+ throw new Error ( "Snapshot was not dispatched" ) ;
41+ }
42+ return metadata ;
3043}
3144
3245function delayedSnapshot ( runnerId = "runner-1" ) {
@@ -38,6 +51,24 @@ function delayedSnapshot(runnerId = "runner-1") {
3851}
3952
4053describe ( "ComputeSnapshotService" , ( ) => {
54+ it ( "refuses to construct with an empty callback secret" , ( ) => {
55+ const computeManager = {
56+ snapshotDelayMs : DELAY_MS ,
57+ snapshotDispatchLimit : 1 ,
58+ snapshot : vi . fn ( async ( ) => true ) ,
59+ } as unknown as ComputeWorkloadManager ;
60+
61+ expect (
62+ ( ) =>
63+ new ComputeSnapshotService ( {
64+ computeManager,
65+ workerClient : { } as SupervisorHttpClient ,
66+ wideEventOpts : { service : "supervisor-test" , env : { } , enabled : false } ,
67+ snapshotCallbackSecret : "" ,
68+ } )
69+ ) . toThrow ( ) ;
70+ } ) ;
71+
4172 it ( "dispatches a scheduled snapshot after the delay" , async ( ) => {
4273 const { service, snapshot } = createService ( ) ;
4374 try {
@@ -46,7 +77,12 @@ describe("ComputeSnapshotService", () => {
4677 await vi . waitFor ( ( ) => expect ( snapshot ) . toHaveBeenCalledTimes ( 1 ) , { timeout : 2_000 } ) ;
4778 expect ( snapshot ) . toHaveBeenCalledWith ( {
4879 runnerId : "runner-1" ,
49- metadata : { runId : "run_1" , snapshotFriendlyId : "snapshot_1" } ,
80+ metadata : expect . objectContaining ( {
81+ runId : "run_1" ,
82+ snapshotFriendlyId : "snapshot_1" ,
83+ snapshotCallbackNonce : expect . any ( String ) ,
84+ snapshotCallbackToken : expect . any ( String ) ,
85+ } ) ,
5086 } ) ;
5187 } finally {
5288 service . stop ( ) ;
@@ -121,8 +157,84 @@ describe("ComputeSnapshotService", () => {
121157 expect ( snapshot ) . toHaveBeenCalledTimes ( 1 ) ;
122158 expect ( snapshot ) . toHaveBeenCalledWith ( {
123159 runnerId : "runner-1" ,
124- metadata : { runId : "run_1" , snapshotFriendlyId : "snapshot_2" } ,
160+ metadata : expect . objectContaining ( {
161+ runId : "run_1" ,
162+ snapshotFriendlyId : "snapshot_2" ,
163+ snapshotCallbackNonce : expect . any ( String ) ,
164+ snapshotCallbackToken : expect . any ( String ) ,
165+ } ) ,
166+ } ) ;
167+ } finally {
168+ service . stop ( ) ;
169+ }
170+ } ) ;
171+
172+ it ( "accepts a snapshot callback with the dispatched token" , async ( ) => {
173+ const { service, snapshot, submitSuspendCompletion } = createService ( ) ;
174+ try {
175+ service . schedule ( "run_1" , delayedSnapshot ( ) ) ;
176+
177+ await vi . waitFor ( ( ) => expect ( snapshot ) . toHaveBeenCalledTimes ( 1 ) , { timeout : 2_000 } ) ;
178+ const metadata = dispatchedMetadata ( snapshot ) ;
179+
180+ const result = await service . handleCallback ( {
181+ status : "completed" ,
182+ instance_id : "instance_1" ,
183+ snapshot_id : "compute_snapshot_1" ,
184+ metadata,
185+ } ) ;
186+
187+ expect ( result ) . toEqual ( { ok : true , status : 200 } ) ;
188+ expect ( submitSuspendCompletion ) . toHaveBeenCalledWith ( {
189+ runId : "run_1" ,
190+ snapshotId : "snapshot_1" ,
191+ body : {
192+ success : true ,
193+ checkpoint : {
194+ type : "COMPUTE" ,
195+ location : "compute_snapshot_1" ,
196+ } ,
197+ } ,
198+ } ) ;
199+ } finally {
200+ service . stop ( ) ;
201+ }
202+ } ) ;
203+
204+ it ( "rejects a snapshot callback without a valid token" , async ( ) => {
205+ const { service, submitSuspendCompletion } = createService ( ) ;
206+ try {
207+ const result = await service . handleCallback ( {
208+ status : "completed" ,
209+ instance_id : "instance_1" ,
210+ snapshot_id : "compute_snapshot_1" ,
211+ metadata : { runId : "run_1" , snapshotFriendlyId : "snapshot_1" } ,
212+ } ) ;
213+
214+ expect ( result ) . toEqual ( { ok : false , status : 401 } ) ;
215+ expect ( submitSuspendCompletion ) . not . toHaveBeenCalled ( ) ;
216+ } finally {
217+ service . stop ( ) ;
218+ }
219+ } ) ;
220+
221+ it ( "rejects a snapshot callback whose token is for a different snapshot" , async ( ) => {
222+ const { service, snapshot, submitSuspendCompletion } = createService ( ) ;
223+ try {
224+ service . schedule ( "run_1" , delayedSnapshot ( ) ) ;
225+
226+ await vi . waitFor ( ( ) => expect ( snapshot ) . toHaveBeenCalledTimes ( 1 ) , { timeout : 2_000 } ) ;
227+ const metadata = dispatchedMetadata ( snapshot ) ;
228+
229+ const result = await service . handleCallback ( {
230+ status : "completed" ,
231+ instance_id : "instance_1" ,
232+ snapshot_id : "compute_snapshot_1" ,
233+ metadata : { ...metadata , snapshotFriendlyId : "snapshot_2" } ,
125234 } ) ;
235+
236+ expect ( result ) . toEqual ( { ok : false , status : 401 } ) ;
237+ expect ( submitSuspendCompletion ) . not . toHaveBeenCalled ( ) ;
126238 } finally {
127239 service . stop ( ) ;
128240 }
0 commit comments