package foldercleanup_test import ( "errors" "testing" "time" "queryorchestration/internal/database/repository" "queryorchestration/internal/document/batch/foldercleanup" "queryorchestration/internal/serviceconfig" "queryorchestration/internal/test" "github.com/google/uuid" "github.com/jackc/pgx/v5" "github.com/jackc/pgx/v5/pgtype" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" ) type TestConfig struct { serviceconfig.BaseConfig } func createClientRoot(t *testing.T, cfg *TestConfig, clientID string) *repository.Folder { t.Helper() ctx := t.Context() err := cfg.GetDBQueries().CreateClient(ctx, &repository.CreateClientParams{ Clientid: clientID, Name: "Test Client", }) require.NoError(t, err) root, err := cfg.GetDBQueries().CreateFolder(ctx, &repository.CreateFolderParams{ Path: "/", Clientid: clientID, Createdby: "test", }) require.NoError(t, err) return root } func createBatch(t *testing.T, cfg *TestConfig, clientID string, total int32) uuid.UUID { t.Helper() batchID, err := cfg.GetDBQueries().CreateBatchUpload(t.Context(), &repository.CreateBatchUploadParams{ ClientID: clientID, OriginalFilename: "batch.zip", TotalDocuments: total, }) require.NoError(t, err) return batchID } func TestFinalizeDeletesNewRejectedFolders(t *testing.T) { cfg := &TestConfig{} test.CreateDB(t, cfg) ctx := t.Context() clientID := "cleanup_rejected_" + uuid.New().String()[:8] root := createClientRoot(t, cfg, clientID) batchID := createBatch(t, cfg, clientID, 1) folder, err := cfg.GetDBQueries().CreateFolder(ctx, &repository.CreateFolderParams{ Path: "/auto", Parentid: &root.ID, Clientid: clientID, Createdby: "test", }) require.NoError(t, err) err = cfg.GetDBQueries().RecordBatchFolderCandidate(ctx, &repository.RecordBatchFolderCandidateParams{ BatchID: batchID, FolderID: folder.ID, Path: folder.Path, ExistedBefore: false, }) require.NoError(t, err) filename := "duplicate.pdf" err = cfg.GetDBQueries().AddDocumentUpload(ctx, &repository.AddDocumentUploadParams{ ID: uuid.New(), Clientid: clientID, Bucket: "bucket", Key: "key", Part: 1, Createdat: pgtype.Timestamp{Valid: true}, Filename: &filename, BatchID: &batchID, FolderID: &folder.ID, }) require.NoError(t, err) err = cfg.GetDBQueries().InsertBatchDocumentOutcome(ctx, &repository.InsertBatchDocumentOutcomeParams{ BatchID: batchID, Filename: filename, Column3: repository.BatchOutcomeStatusDuplicate, }) require.NoError(t, err) require.NoError(t, cfg.GetDBQueries().UpdateBatchStatus(ctx, &repository.UpdateBatchStatusParams{ ID: batchID, Column2: repository.BatchStatusFailed, })) result, err := foldercleanup.New(cfg).TryFinalizeBatch(ctx, batchID) require.NoError(t, err) require.True(t, result.Ready) assert.Equal(t, int64(1), result.DetachedUploads) assert.Equal(t, int64(1), result.DeletedFolders) upload, err := cfg.GetDBQueries().GetDocumentUploadByKey(ctx, &repository.GetDocumentUploadByKeyParams{ Bucket: "bucket", Key: "key", }) require.NoError(t, err) assert.Nil(t, upload.FolderID) _, err = cfg.GetDBQueries().GetFolderByID(ctx, folder.ID) assert.True(t, errors.Is(err, pgx.ErrNoRows)) } func TestFinalizePreservesExistingFolders(t *testing.T) { cfg := &TestConfig{} test.CreateDB(t, cfg) ctx := t.Context() clientID := "cleanup_existing_" + uuid.New().String()[:8] root := createClientRoot(t, cfg, clientID) batchID := createBatch(t, cfg, clientID, 1) folder, err := cfg.GetDBQueries().CreateFolder(ctx, &repository.CreateFolderParams{ Path: "/existing", Parentid: &root.ID, Clientid: clientID, Createdby: "test", }) require.NoError(t, err) err = cfg.GetDBQueries().RecordBatchFolderCandidate(ctx, &repository.RecordBatchFolderCandidateParams{ BatchID: batchID, FolderID: folder.ID, Path: folder.Path, ExistedBefore: true, }) require.NoError(t, err) err = cfg.GetDBQueries().InsertBatchDocumentOutcome(ctx, &repository.InsertBatchDocumentOutcomeParams{ BatchID: batchID, Filename: "duplicate.pdf", Column3: repository.BatchOutcomeStatusDuplicate, }) require.NoError(t, err) require.NoError(t, cfg.GetDBQueries().UpdateBatchStatus(ctx, &repository.UpdateBatchStatusParams{ ID: batchID, Column2: repository.BatchStatusFailed, })) result, err := foldercleanup.New(cfg).TryFinalizeBatch(ctx, batchID) require.NoError(t, err) require.True(t, result.Ready) assert.Zero(t, result.DeletedFolders) persisted, err := cfg.GetDBQueries().GetFolderByID(ctx, folder.ID) require.NoError(t, err) assert.Equal(t, "/existing", persisted.Path) } func TestFinalizeReadyBatchesRetriesReadyCleanup(t *testing.T) { cfg := &TestConfig{} test.CreateDB(t, cfg) ctx := t.Context() clientID := "cleanup_retry_" + uuid.New().String()[:8] root := createClientRoot(t, cfg, clientID) batchID := createBatch(t, cfg, clientID, 1) folder, err := cfg.GetDBQueries().CreateFolder(ctx, &repository.CreateFolderParams{ Path: "/retry", Parentid: &root.ID, Clientid: clientID, Createdby: "test", }) require.NoError(t, err) require.NoError(t, cfg.GetDBQueries().RecordBatchFolderCandidate(ctx, &repository.RecordBatchFolderCandidateParams{ BatchID: batchID, FolderID: folder.ID, Path: folder.Path, ExistedBefore: false, })) require.NoError(t, cfg.GetDBQueries().InsertBatchDocumentOutcome(ctx, &repository.InsertBatchDocumentOutcomeParams{ BatchID: batchID, Filename: "invalid.exe", Column3: repository.BatchOutcomeStatusInvalidType, })) require.NoError(t, cfg.GetDBQueries().UpdateBatchStatus(ctx, &repository.UpdateBatchStatusParams{ ID: batchID, Column2: repository.BatchStatusFailed, })) service := foldercleanup.New(cfg) finalized, err := service.FinalizeReadyBatches(ctx, 10) require.NoError(t, err) assert.Equal(t, 1, finalized) batch, err := cfg.GetDBQueries().LockBatchUploadForFolderCleanup(ctx, batchID) require.NoError(t, err) assert.True(t, batch.FolderCleanupCompletedAt.Valid) _, err = cfg.GetDBQueries().GetFolderByID(ctx, folder.ID) assert.True(t, errors.Is(err, pgx.ErrNoRows)) finalized, err = service.FinalizeReadyBatches(ctx, 10) require.NoError(t, err) assert.Zero(t, finalized) } func TestFinalizeDeclinesFailedBatchWithoutOutcomes(t *testing.T) { cfg := &TestConfig{} test.CreateDB(t, cfg) ctx := t.Context() clientID := "cleanup_zero_" + uuid.New().String()[:8] root := createClientRoot(t, cfg, clientID) batchID := createBatch(t, cfg, clientID, 0) folder, err := cfg.GetDBQueries().CreateFolder(ctx, &repository.CreateFolderParams{ Path: "/zero", Parentid: &root.ID, Clientid: clientID, Createdby: "test", }) require.NoError(t, err) require.NoError(t, cfg.GetDBQueries().RecordBatchFolderCandidate(ctx, &repository.RecordBatchFolderCandidateParams{ BatchID: batchID, FolderID: folder.ID, Path: folder.Path, ExistedBefore: false, })) require.NoError(t, cfg.GetDBQueries().UpdateBatchStatus(ctx, &repository.UpdateBatchStatusParams{ ID: batchID, Column2: repository.BatchStatusFailed, })) service := foldercleanup.New(cfg) result, err := service.TryFinalizeBatch(ctx, batchID) require.NoError(t, err) require.False(t, result.Ready) batch, err := cfg.GetDBQueries().LockBatchUploadForFolderCleanup(ctx, batchID) require.NoError(t, err) assert.False(t, batch.FolderCleanupCompletedAt.Valid) _, err = cfg.GetDBQueries().GetFolderByID(ctx, folder.ID) require.NoError(t, err) finalized, err := service.FinalizeReadyBatches(ctx, 10) require.NoError(t, err) assert.Zero(t, finalized) } func TestFinalizeStaleFailedBatchWithOutcomeEvidence(t *testing.T) { cfg := &TestConfig{} test.CreateDB(t, cfg) ctx := t.Context() clientID := "cleanup_stale_" + uuid.New().String()[:8] root := createClientRoot(t, cfg, clientID) batchID := createBatch(t, cfg, clientID, 2) folder, err := cfg.GetDBQueries().CreateFolder(ctx, &repository.CreateFolderParams{ Path: "/stale", Parentid: &root.ID, Clientid: clientID, Createdby: "test", }) require.NoError(t, err) require.NoError(t, cfg.GetDBQueries().RecordBatchFolderCandidate(ctx, &repository.RecordBatchFolderCandidateParams{ BatchID: batchID, FolderID: folder.ID, Path: folder.Path, ExistedBefore: false, })) require.NoError(t, cfg.GetDBQueries().InsertBatchDocumentOutcome(ctx, &repository.InsertBatchDocumentOutcomeParams{ BatchID: batchID, Filename: "stuck.pdf", Column3: repository.BatchOutcomeStatusSubmitted, })) require.NoError(t, cfg.GetDBQueries().UpdateBatchStatus(ctx, &repository.UpdateBatchStatusParams{ ID: batchID, Column2: repository.BatchStatusFailed, })) require.NoError(t, cfg.GetDBQueries().SetBatchCompletedAt(ctx, &repository.SetBatchCompletedAtParams{ ID: batchID, CompletedAt: pgtype.Timestamp{Time: time.Now().Add(-20 * time.Minute), Valid: true}, })) result, err := foldercleanup.New(cfg).TryFinalizeBatch(ctx, batchID) require.NoError(t, err) require.True(t, result.Ready) assert.Equal(t, int64(1), result.DeletedFolders) assert.Equal(t, int32(1), result.RejectedDocuments) }