package repository_test import ( "fmt" "testing" "queryorchestration/internal/database/repository" "queryorchestration/internal/serviceconfig" "queryorchestration/internal/test" "github.com/google/uuid" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" ) func TestBatchUpload(t *testing.T) { if testing.Short() { t.SkipNow() } t.Run("CreateBatchUpload", func(t *testing.T) { t.Parallel() ctx := t.Context() cfg := &serviceconfig.BaseConfig{} test.CreateDB(t, cfg) queries := cfg.GetDBQueries() // Create a client first clientID := fmt.Sprintf("TEST_CLIENT_%s", uuid.New().String()[:8]) err := queries.CreateClient(ctx, &repository.CreateClientParams{ Name: fmt.Sprintf("Test Client %s", uuid.New().String()[:8]), Clientid: clientID, }) require.NoError(t, err) // Create batch upload params := &repository.CreateBatchUploadParams{ ClientID: clientID, OriginalFilename: "test_batch.zip", TotalDocuments: 10, } batchID, err := queries.CreateBatchUpload(ctx, params) require.NoError(t, err) assert.NotEqual(t, uuid.Nil, batchID) }) t.Run("GetBatchUpload", func(t *testing.T) { t.Parallel() ctx := t.Context() cfg := &serviceconfig.BaseConfig{} test.CreateDB(t, cfg) queries := cfg.GetDBQueries() // Create a client clientID := fmt.Sprintf("TEST_CLIENT_GET_%s", uuid.New().String()[:8]) err := queries.CreateClient(ctx, &repository.CreateClientParams{ Name: fmt.Sprintf("Test Client Get %s", uuid.New().String()[:8]), Clientid: clientID, }) require.NoError(t, err) // Create batch upload batchID, err := queries.CreateBatchUpload(ctx, &repository.CreateBatchUploadParams{ ClientID: clientID, OriginalFilename: "test_get.zip", TotalDocuments: 5, }) require.NoError(t, err) // Get batch upload batch, err := queries.GetBatchUpload(ctx, &repository.GetBatchUploadParams{ ID: batchID, ClientID: clientID, }) require.NoError(t, err) assert.Equal(t, batchID, batch.ID) assert.Equal(t, clientID, batch.ClientID) assert.Equal(t, "test_get.zip", batch.OriginalFilename) assert.Equal(t, int32(5), batch.TotalDocuments) assert.Equal(t, repository.BatchStatusProcessing, batch.Status) assert.Equal(t, int32(0), batch.ProcessedDocuments) assert.Equal(t, int32(0), batch.FailedDocuments) assert.NotNil(t, batch.CreatedAt) assert.False(t, batch.CompletedAt.Valid) // pgtype.Timestamp has Valid field }) t.Run("ListBatchUploads", func(t *testing.T) { t.Parallel() ctx := t.Context() cfg := &serviceconfig.BaseConfig{} test.CreateDB(t, cfg) queries := cfg.GetDBQueries() // Create a client clientID := fmt.Sprintf("TEST_CLIENT_LIST_%s", uuid.New().String()[:8]) err := queries.CreateClient(ctx, &repository.CreateClientParams{ Name: fmt.Sprintf("Test Client List %s", uuid.New().String()[:8]), Clientid: clientID, }) require.NoError(t, err) // Create multiple batch uploads for i := 0; i < 3; i++ { totalDocs := int32(1) + int32(i) //nolint:gosec // safe conversion for test data _, err := queries.CreateBatchUpload(ctx, &repository.CreateBatchUploadParams{ ClientID: clientID, OriginalFilename: fmt.Sprintf("batch%d.zip", i), TotalDocuments: totalDocs, }) require.NoError(t, err) } // List batch uploads batches, err := queries.ListBatchUploads(ctx, &repository.ListBatchUploadsParams{ ClientID: clientID, Limit: 10, Offset: 0, }) require.NoError(t, err) assert.Len(t, batches, 3) }) t.Run("UpdateBatchProgress", func(t *testing.T) { t.Parallel() ctx := t.Context() cfg := &serviceconfig.BaseConfig{} test.CreateDB(t, cfg) queries := cfg.GetDBQueries() // Create a client clientID := fmt.Sprintf("TEST_CLIENT_PROGRESS_%s", uuid.New().String()[:8]) err := queries.CreateClient(ctx, &repository.CreateClientParams{ Name: fmt.Sprintf("Test Client Progress %s", uuid.New().String()[:8]), Clientid: clientID, }) require.NoError(t, err) // Create batch upload batchID, err := queries.CreateBatchUpload(ctx, &repository.CreateBatchUploadParams{ ClientID: clientID, OriginalFilename: "progress.zip", TotalDocuments: 100, }) require.NoError(t, err) // Update progress err = queries.UpdateBatchProgress(ctx, &repository.UpdateBatchProgressParams{ ID: batchID, ProcessedDocuments: 42, FailedDocuments: 2, InvalidTypeDocuments: 1, ProgressPercent: 45, }) require.NoError(t, err) // Verify update batch, err := queries.GetBatchUpload(ctx, &repository.GetBatchUploadParams{ ID: batchID, ClientID: clientID, }) require.NoError(t, err) assert.Equal(t, int32(42), batch.ProcessedDocuments) assert.Equal(t, int32(2), batch.FailedDocuments) assert.Equal(t, int32(1), batch.InvalidTypeDocuments) assert.Equal(t, int32(45), batch.ProgressPercent) }) t.Run("UpdateBatchStatus", func(t *testing.T) { t.Parallel() ctx := t.Context() cfg := &serviceconfig.BaseConfig{} test.CreateDB(t, cfg) queries := cfg.GetDBQueries() // Create a client clientID := fmt.Sprintf("TEST_CLIENT_STATUS_%s", uuid.New().String()[:8]) err := queries.CreateClient(ctx, &repository.CreateClientParams{ Name: fmt.Sprintf("Test Client Status %s", uuid.New().String()[:8]), Clientid: clientID, }) require.NoError(t, err) // Create batch upload batchID, err := queries.CreateBatchUpload(ctx, &repository.CreateBatchUploadParams{ ClientID: clientID, OriginalFilename: "status.zip", TotalDocuments: 10, }) require.NoError(t, err) // Update status to completed err = queries.UpdateBatchStatus(ctx, &repository.UpdateBatchStatusParams{ ID: batchID, Column2: repository.BatchStatusCompleted, }) require.NoError(t, err) // Verify update batch, err := queries.GetBatchUpload(ctx, &repository.GetBatchUploadParams{ ID: batchID, ClientID: clientID, }) require.NoError(t, err) assert.Equal(t, repository.BatchStatusCompleted, batch.Status) assert.True(t, batch.CompletedAt.Valid) assert.True(t, batch.CompletedAt.Time.After(batch.CreatedAt.Time)) }) t.Run("AddFailedFilename", func(t *testing.T) { t.Parallel() ctx := t.Context() cfg := &serviceconfig.BaseConfig{} test.CreateDB(t, cfg) queries := cfg.GetDBQueries() // Create a client clientID := fmt.Sprintf("TEST_CLIENT_FAILED_%s", uuid.New().String()[:8]) err := queries.CreateClient(ctx, &repository.CreateClientParams{ Name: fmt.Sprintf("Test Client Failed %s", uuid.New().String()[:8]), Clientid: clientID, }) require.NoError(t, err) // Create batch upload batchID, err := queries.CreateBatchUpload(ctx, &repository.CreateBatchUploadParams{ ClientID: clientID, OriginalFilename: "failed.zip", TotalDocuments: 10, }) require.NoError(t, err) // Add failed filenames err = queries.AddFailedFilename(ctx, &repository.AddFailedFilenameParams{ ID: batchID, Column2: "document1.pdf", }) require.NoError(t, err) err = queries.AddFailedFilename(ctx, &repository.AddFailedFilenameParams{ ID: batchID, Column2: "document2.pdf", }) require.NoError(t, err) // Verify failed filenames were added batch, err := queries.GetBatchUpload(ctx, &repository.GetBatchUploadParams{ ID: batchID, ClientID: clientID, }) require.NoError(t, err) // Verify failed filenames were added as JSON array assert.NotNil(t, batch.FailedFilenames) // The failed filenames should contain our test filenames failedFilenamesStr := string(batch.FailedFilenames) assert.Contains(t, failedFilenamesStr, "document1.pdf") assert.Contains(t, failedFilenamesStr, "document2.pdf") }) t.Run("CancelBatchUpload", func(t *testing.T) { t.Parallel() ctx := t.Context() cfg := &serviceconfig.BaseConfig{} test.CreateDB(t, cfg) queries := cfg.GetDBQueries() // Create a client clientID := fmt.Sprintf("TEST_CLIENT_CANCEL_%s", uuid.New().String()[:8]) err := queries.CreateClient(ctx, &repository.CreateClientParams{ Name: fmt.Sprintf("Test Client Cancel %s", uuid.New().String()[:8]), Clientid: clientID, }) require.NoError(t, err) // Create batch upload batchID, err := queries.CreateBatchUpload(ctx, &repository.CreateBatchUploadParams{ ClientID: clientID, OriginalFilename: "cancel.zip", TotalDocuments: 10, }) require.NoError(t, err) // Cancel batch upload err = queries.CancelBatchUpload(ctx, &repository.CancelBatchUploadParams{ ID: batchID, ClientID: clientID, }) require.NoError(t, err) // Verify cancellation batch, err := queries.GetBatchUpload(ctx, &repository.GetBatchUploadParams{ ID: batchID, ClientID: clientID, }) require.NoError(t, err) assert.Equal(t, repository.BatchStatusCancelled, batch.Status) assert.True(t, batch.CompletedAt.Valid) }) t.Run("GetDocumentsByBatchId", func(t *testing.T) { t.Parallel() ctx := t.Context() cfg := &serviceconfig.BaseConfig{} test.CreateDB(t, cfg) queries := cfg.GetDBQueries() // Create a client clientID := fmt.Sprintf("TEST_CLIENT_DOCS_%s", uuid.New().String()[:8]) err := queries.CreateClient(ctx, &repository.CreateClientParams{ Name: fmt.Sprintf("Test Client Docs %s", uuid.New().String()[:8]), Clientid: clientID, }) require.NoError(t, err) // Create batch upload batchID, err := queries.CreateBatchUpload(ctx, &repository.CreateBatchUploadParams{ ClientID: clientID, OriginalFilename: "docs.zip", TotalDocuments: 3, }) require.NoError(t, err) // Create documents with batch ID var docIDs []uuid.UUID for i := 0; i < 3; i++ { docID, err := queries.CreateDocument(ctx, &repository.CreateDocumentParams{ Clientid: clientID, Hash: fmt.Sprintf("hash%d", i), BatchID: &batchID, }) require.NoError(t, err) docIDs = append(docIDs, docID) } // Get documents by batch ID docs, err := queries.GetDocumentsByBatchId(ctx, &batchID) require.NoError(t, err) assert.Len(t, docs, 3) // Verify all documents are returned for _, doc := range docs { assert.Contains(t, docIDs, doc.ID) assert.Equal(t, clientID, doc.Clientid) } }) t.Run("CountDocumentsByBatchId", func(t *testing.T) { t.Parallel() ctx := t.Context() cfg := &serviceconfig.BaseConfig{} test.CreateDB(t, cfg) queries := cfg.GetDBQueries() // Create a client clientID := fmt.Sprintf("TEST_CLIENT_COUNT_%s", uuid.New().String()[:8]) err := queries.CreateClient(ctx, &repository.CreateClientParams{ Name: fmt.Sprintf("Test Client Count %s", uuid.New().String()[:8]), Clientid: clientID, }) require.NoError(t, err) // Create batch upload batchID, err := queries.CreateBatchUpload(ctx, &repository.CreateBatchUploadParams{ ClientID: clientID, OriginalFilename: "count.zip", TotalDocuments: 5, }) require.NoError(t, err) // Create documents with batch ID for i := 0; i < 4; i++ { _, err := queries.CreateDocument(ctx, &repository.CreateDocumentParams{ Clientid: clientID, Hash: fmt.Sprintf("hash%d", i), BatchID: &batchID, }) require.NoError(t, err) } // Count documents by batch ID count, err := queries.CountDocumentsByBatchId(ctx, &batchID) require.NoError(t, err) assert.Equal(t, int64(4), count) }) t.Run("GetUnprocessedBatches", func(t *testing.T) { t.Parallel() ctx := t.Context() cfg := &serviceconfig.BaseConfig{} test.CreateDB(t, cfg) queries := cfg.GetDBQueries() // Create multiple clients clientID1 := fmt.Sprintf("TEST_CLIENT_UNPROCESSED1_%s", uuid.New().String()[:8]) clientID2 := fmt.Sprintf("TEST_CLIENT_UNPROCESSED2_%s", uuid.New().String()[:8]) for _, clientID := range []string{clientID1, clientID2} { err := queries.CreateClient(ctx, &repository.CreateClientParams{ Name: fmt.Sprintf("Test Client %s", clientID), Clientid: clientID, }) require.NoError(t, err) } // Create batches with different statuses var processingBatchIDs []uuid.UUID // Client 1: 2 processing batches for i := 0; i < 2; i++ { batchID, err := queries.CreateBatchUpload(ctx, &repository.CreateBatchUploadParams{ ClientID: clientID1, OriginalFilename: fmt.Sprintf("processing%d.zip", i), TotalDocuments: 10, }) require.NoError(t, err) processingBatchIDs = append(processingBatchIDs, batchID) } // Client 2: 1 processing batch batchID, err := queries.CreateBatchUpload(ctx, &repository.CreateBatchUploadParams{ ClientID: clientID2, OriginalFilename: "processing_client2.zip", TotalDocuments: 5, }) require.NoError(t, err) processingBatchIDs = append(processingBatchIDs, batchID) // Create completed batch (should not be returned) completedBatchID, err := queries.CreateBatchUpload(ctx, &repository.CreateBatchUploadParams{ ClientID: clientID1, OriginalFilename: "completed.zip", TotalDocuments: 8, }) require.NoError(t, err) err = queries.UpdateBatchStatus(ctx, &repository.UpdateBatchStatusParams{ ID: completedBatchID, Column2: repository.BatchStatusCompleted, }) require.NoError(t, err) // Create failed batch (should not be returned) failedBatchID, err := queries.CreateBatchUpload(ctx, &repository.CreateBatchUploadParams{ ClientID: clientID2, OriginalFilename: "failed.zip", TotalDocuments: 3, }) require.NoError(t, err) err = queries.UpdateBatchStatus(ctx, &repository.UpdateBatchStatusParams{ ID: failedBatchID, Column2: repository.BatchStatusFailed, }) require.NoError(t, err) // Get unprocessed batches unprocessedBatches, err := queries.GetUnprocessedBatches(ctx) require.NoError(t, err) // Should return exactly 3 processing batches assert.Len(t, unprocessedBatches, 3) // Verify all returned batches have processing status returnedBatchIDs := make([]uuid.UUID, len(unprocessedBatches)) for i, batch := range unprocessedBatches { assert.Equal(t, repository.BatchStatusProcessing, batch.Status) returnedBatchIDs[i] = batch.ID } // Verify all processing batches are returned for _, expectedID := range processingBatchIDs { assert.Contains(t, returnedBatchIDs, expectedID) } // Verify completed and failed batches are NOT returned assert.NotContains(t, returnedBatchIDs, completedBatchID) assert.NotContains(t, returnedBatchIDs, failedBatchID) // Verify batches are ordered by created_at ASC (oldest first) if len(unprocessedBatches) > 1 { for i := 0; i < len(unprocessedBatches)-1; i++ { assert.True(t, unprocessedBatches[i].CreatedAt.Time.Before(unprocessedBatches[i+1].CreatedAt.Time) || unprocessedBatches[i].CreatedAt.Time.Equal(unprocessedBatches[i+1].CreatedAt.Time), "batches should be ordered by created_at ASC") } } }) } func TestBatchStatusEnum(t *testing.T) { t.Run("Scan method", func(t *testing.T) { // Test valid scan with string var status repository.BatchStatus err := status.Scan("processing") require.NoError(t, err) assert.Equal(t, repository.BatchStatusProcessing, status) // Test valid scan with []byte var statusBytes repository.BatchStatus err = statusBytes.Scan([]byte("completed")) require.NoError(t, err) assert.Equal(t, repository.BatchStatusCompleted, statusBytes) // Test all enum values testCases := []struct { str string expected repository.BatchStatus }{ {"processing", repository.BatchStatusProcessing}, {"completed", repository.BatchStatusCompleted}, {"failed", repository.BatchStatusFailed}, {"cancelled", repository.BatchStatusCancelled}, //nolint:misspell // matches database enum value } for _, tc := range testCases { var s repository.BatchStatus err := s.Scan(tc.str) require.NoError(t, err) assert.Equal(t, tc.expected, s) } // Test scan with unsupported type var invalidStatus repository.BatchStatus err = invalidStatus.Scan(123) // int is not supported assert.Error(t, err) assert.Contains(t, err.Error(), "unsupported scan type") }) t.Run("NullBatchStatus", func(t *testing.T) { // Test scan with nil var nullStatus repository.NullBatchStatus err := nullStatus.Scan(nil) require.NoError(t, err) assert.False(t, nullStatus.Valid) assert.Equal(t, repository.BatchStatus(""), nullStatus.BatchStatus) // Test scan with valid value err = nullStatus.Scan("processing") require.NoError(t, err) assert.True(t, nullStatus.Valid) assert.Equal(t, repository.BatchStatusProcessing, nullStatus.BatchStatus) // Test Value method when Valid is false nullStatus.Valid = false // reset to false val, err := nullStatus.Value() require.NoError(t, err) assert.Nil(t, val) // Test Value method when Valid is true nullStatus.Valid = true nullStatus.BatchStatus = repository.BatchStatusCompleted val, err = nullStatus.Value() require.NoError(t, err) assert.Equal(t, "completed", val) }) t.Run("BatchStatus Valid method", func(t *testing.T) { // Test valid enum values assert.True(t, repository.BatchStatusProcessing.Valid()) assert.True(t, repository.BatchStatusCompleted.Valid()) assert.True(t, repository.BatchStatusFailed.Valid()) assert.True(t, repository.BatchStatusCancelled.Valid()) //nolint:misspell // matches database enum value // Test invalid enum value invalidStatus := repository.BatchStatus("invalid") assert.False(t, invalidStatus.Valid()) }) }