query-orchestration/internal/document/batch/foldercleanup/service_test.go

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)
}