@@ -61,6 +61,56 @@ func TestCheckpointerResumption(t *testing.T) {
6161 assert .Equal (t , "key-11" , finalResumeInfo .StartAfter )
6262}
6363
64+ func TestCheckpointerResumptionWithRole (t * testing.T ) {
65+ ctx := context .Background ()
66+
67+ // First scan - process 6 objects then interrupt.
68+ initialProgress := & sources.Progress {}
69+ tracker := NewCheckpointer (ctx , initialProgress )
70+ role := "test-role"
71+
72+ firstPage := & s3.ListObjectsV2Output {
73+ Contents : make ([]s3types.Object , 12 ), // Total of 12 objects
74+ }
75+ for i := range 12 {
76+ key := fmt .Sprintf ("key-%d" , i )
77+ firstPage .Contents [i ] = s3types.Object {Key : & key }
78+ }
79+
80+ // Process first 6 objects.
81+ for i := range 6 {
82+ err := tracker .UpdateObjectCompletion (ctx , i , "test-bucket" , role , firstPage .Contents )
83+ assert .NoError (t , err )
84+ }
85+
86+ // Verify resume info is set correctly.
87+ resumeInfo , err := tracker .ResumePoint (ctx )
88+ require .NoError (t , err )
89+ assert .Equal (t , "test-bucket" , resumeInfo .CurrentBucket )
90+ assert .Equal (t , "key-5" , resumeInfo .StartAfter )
91+ assert .Equal (t , role , resumeInfo .Role )
92+
93+ // Resume scan with existing progress.
94+ resumeTracker := NewCheckpointer (ctx , initialProgress )
95+
96+ resumePage := & s3.ListObjectsV2Output {
97+ Contents : firstPage .Contents [6 :], // Remaining 6 objects
98+ }
99+
100+ // Process remaining objects.
101+ for i := range len (resumePage .Contents ) {
102+ err := resumeTracker .UpdateObjectCompletion (ctx , i , "test-bucket" , role , resumePage .Contents )
103+ assert .NoError (t , err )
104+ }
105+
106+ // Verify final resume info.
107+ finalResumeInfo , err := resumeTracker .ResumePoint (ctx )
108+ require .NoError (t , err )
109+ assert .Equal (t , "test-bucket" , finalResumeInfo .CurrentBucket )
110+ assert .Equal (t , "key-11" , finalResumeInfo .StartAfter )
111+ assert .Equal (t , role , finalResumeInfo .Role )
112+ }
113+
64114func TestCheckpointerReset (t * testing.T ) {
65115 tests := []struct {
66116 name string
@@ -111,6 +161,13 @@ func TestGetResumePoint(t *testing.T) {
111161 },
112162 expectedResumeInfo : ResumeInfo {CurrentBucket : "test-bucket" , StartAfter : "test-key" },
113163 },
164+ {
165+ name : "valid resume info with role" ,
166+ progress : & sources.Progress {
167+ EncodedResumeInfo : `{"current_bucket":"test-bucket","start_after":"test-key","role":"test-role"}` ,
168+ },
169+ expectedResumeInfo : ResumeInfo {CurrentBucket : "test-bucket" , StartAfter : "test-key" , Role : "test-role" },
170+ },
114171 {
115172 name : "empty encoded resume info" ,
116173 progress : & sources.Progress {EncodedResumeInfo : "" },
@@ -121,6 +178,13 @@ func TestGetResumePoint(t *testing.T) {
121178 EncodedResumeInfo : `{"current_bucket":"","start_after":"test-key"}` ,
122179 },
123180 },
181+ {
182+ name : "no role in resume info" ,
183+ progress : & sources.Progress {
184+ EncodedResumeInfo : `{"current_bucket":"test-bucket","start_after":"test-key"}` ,
185+ },
186+ expectedResumeInfo : ResumeInfo {CurrentBucket : "test-bucket" , StartAfter : "test-key" , Role : "" },
187+ },
124188 {
125189 name : "unmarshal error" ,
126190 progress : & sources.Progress {
@@ -257,6 +321,122 @@ func TestCheckpointerUpdate(t *testing.T) {
257321 })
258322 }
259323}
324+ func TestCheckpointerUpdateWithRole (t * testing.T ) {
325+ role := "test-role"
326+ tests := []struct {
327+ name string
328+ description string
329+ completedIdx int
330+ pageSize int
331+ preCompleted []int
332+ expectedKey string
333+ expectedRole string
334+ expectedLowestIncomplete int
335+ }{
336+ {
337+ name : "first object completed" ,
338+ description : "Basic case - completing first object" ,
339+ completedIdx : 0 ,
340+ pageSize : 3 ,
341+ expectedKey : "key-0" ,
342+ expectedRole : role ,
343+ expectedLowestIncomplete : 1 ,
344+ },
345+ {
346+ name : "completing missing middle" ,
347+ description : "Completing object when previous is done" ,
348+ completedIdx : 1 ,
349+ pageSize : 3 ,
350+ preCompleted : []int {0 },
351+ expectedKey : "key-1" ,
352+ expectedRole : role ,
353+ expectedLowestIncomplete : 2 ,
354+ },
355+ {
356+ name : "all objects completed in order" ,
357+ description : "Completing final object in sequence" ,
358+ completedIdx : 2 ,
359+ pageSize : 3 ,
360+ preCompleted : []int {0 , 1 },
361+ expectedKey : "key-2" ,
362+ expectedRole : role ,
363+ expectedLowestIncomplete : 3 ,
364+ },
365+ {
366+ name : "out of order completion before lowest" ,
367+ description : "Completing object before current lowest incomplete - should not affect checkpoint" ,
368+ completedIdx : 1 ,
369+ pageSize : 4 ,
370+ preCompleted : []int {0 , 2 , 3 },
371+ expectedKey : "key-3" ,
372+ expectedRole : role ,
373+ expectedLowestIncomplete : 4 ,
374+ },
375+ {
376+ name : "last index in max page" ,
377+ description : "Edge case - maximum page size boundary" ,
378+ completedIdx : 999 ,
379+ pageSize : 1000 ,
380+ preCompleted : func () []int {
381+ indices := make ([]int , 999 )
382+ for i := range indices {
383+ indices [i ] = i
384+ }
385+ return indices
386+ }(),
387+ expectedKey : "key-999" ,
388+ expectedRole : role ,
389+ expectedLowestIncomplete : 1000 ,
390+ },
391+ }
392+
393+ for _ , tt := range tests {
394+ t .Run (tt .name , func (t * testing.T ) {
395+ t .Parallel ()
396+
397+ ctx := context .Background ()
398+ progress := new (sources.Progress )
399+ tracker := & Checkpointer {
400+ progress : progress ,
401+ completedObjects : make ([]bool , tt .pageSize ),
402+ completionOrder : make ([]int , 0 , tt .pageSize ),
403+ lowestIncompleteIdx : 0 ,
404+ }
405+
406+ page := & s3.ListObjectsV2Output {Contents : make ([]s3types.Object , tt .pageSize )}
407+ for i := range tt .pageSize {
408+ key := fmt .Sprintf ("key-%d" , i )
409+ page .Contents [i ] = s3types.Object {Key : & key }
410+ }
411+
412+ // Setup pre-completed objects.
413+ for _ , idx := range tt .preCompleted {
414+ tracker .completedObjects [idx ] = true
415+ tracker .completionOrder = append (tracker .completionOrder , idx )
416+ }
417+
418+ // Find the correct lowest incomplete index after pre-completion.
419+ for i := range tt .pageSize {
420+ if ! tracker .completedObjects [i ] {
421+ tracker .lowestIncompleteIdx = i
422+ break
423+ }
424+ }
425+
426+ err := tracker .UpdateObjectCompletion (ctx , tt .completedIdx , "test-bucket" , role , page .Contents )
427+ assert .NoError (t , err , "Unexpected error updating progress" )
428+
429+ var info ResumeInfo
430+ err = json .Unmarshal ([]byte (progress .EncodedResumeInfo ), & info )
431+ assert .NoError (t , err , "Failed to decode resume info" )
432+ assert .Equal (t , tt .expectedKey , info .StartAfter , "Incorrect resume point" )
433+ assert .Equal (t , tt .expectedRole , info .Role , "Incorrect role" )
434+
435+ assert .Equal (t , tt .expectedLowestIncomplete , tracker .lowestIncompleteIdx ,
436+ "Incorrect lowest incomplete index" )
437+ })
438+ }
439+ }
260440
261441func TestCheckpointerUpdateUnitScan (t * testing.T ) {
262442 ctx := context .Background ()
0 commit comments