query-orchestration/internal/database/repository/batch_test.go

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