query-orchestration/internal/server/api/batch_worker_test.go

package api

import (
	"archive/zip"
	"bytes"
	"context"
	"fmt"
	"io"
	"log/slog"
	"os"
	"testing"

	"queryorchestration/internal/database/repository"
	"queryorchestration/internal/document/batch"
	"queryorchestration/internal/serviceconfig"
	"queryorchestration/internal/serviceconfig/objectstore"
	"queryorchestration/internal/test"

	"github.com/aws/aws-sdk-go-v2/aws"
	"github.com/aws/aws-sdk-go-v2/service/s3"
	"github.com/google/uuid"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
)

// TestConfig embeds objectstore for S3 operations
type TestConfig struct {
	serviceconfig.BaseConfig
	objectstore.ObjectStoreConfig
}

// TestProcessBatchWork tests the main batch processing work function
func TestProcessBatchWork(t *testing.T) {
	if testing.Short() {
		t.SkipNow()
	}

	ctx := t.Context()
	cfg := &TestConfig{}
	_ = serviceconfig.InitializeConfig(cfg)
	test.CreateDB(t, cfg)
	test.CreateAWSResources(t, cfg)

	// Create test client
	clientID := "test_batch_worker"
	err := cfg.GetDBQueries().CreateClient(ctx, &repository.CreateClientParams{
		Clientid: clientID,
		Name:     "Test Batch Worker Client",
	})
	require.NoError(t, err)

	// Create a test batch service
	batchService := batch.New(cfg)

	// Upload a test ZIP file to S3
	zipContent := createTestZIPForWorker(t, 2)
	s3ClientInterface := cfg.GetStoreClient()
	s3Client, ok := s3ClientInterface.(*s3.Client)
	require.True(t, ok, "s3Client should be *s3.Client")
	bucket := cfg.GetBucket()
	archiveKey := fmt.Sprintf("test/%s/batch.zip", clientID)

	_, err = s3Client.PutObject(ctx, &s3.PutObjectInput{
		Bucket: aws.String(bucket),
		Key:    aws.String(archiveKey),
		Body:   bytes.NewReader(zipContent),
	})
	require.NoError(t, err)

	// Create a batch record in the database
	batchID, err := batchService.CreateWithStorage(ctx, clientID, "batch.zip", 2, bucket, archiveKey, int64(len(zipContent)))
	require.NoError(t, err)

	// Track uploaded documents
	uploadedDocs := make(map[string][]byte)
	uploadHandler := func(ctx context.Context, clientID string, docData io.Reader, filename string, batchID *uuid.UUID) error {
		data, err := io.ReadAll(docData)
		if err != nil {
			return err
		}
		uploadedDocs[filename] = data
		return nil
	}

	// Create worker config
	workerConfig := map[string]any{
		ConfigKeyBatchService:  batchService,
		ConfigKeyS3Client:      s3Client,
		ConfigKeyBucket:        bucket,
		ConfigKeyUploadHandler: uploadHandler,
	}

	// Execute the batch work
	logger := slog.New(slog.NewTextHandler(os.Stdout, nil))
	err = processBatchWork(ctx, logger, workerConfig)
	require.NoError(t, err)

	// Verify that documents were uploaded
	assert.Len(t, uploadedDocs, 2)
	assert.Contains(t, uploadedDocs, "document_1.pdf")
	assert.Contains(t, uploadedDocs, "document_2.pdf")

	// Verify batch status was updated
	batchInfo, err := batchService.Get(ctx, clientID, batchID)
	require.NoError(t, err)
	assert.Equal(t, "completed", batchInfo.Status)
	assert.Equal(t, int32(2), batchInfo.ProcessedDocuments)
	assert.Equal(t, int32(0), batchInfo.FailedDocuments)
}

// TestProcessSingleBatch tests processing a single batch
func TestProcessSingleBatch(t *testing.T) {
	if testing.Short() {
		t.SkipNow()
	}

	ctx := t.Context()
	cfg := &TestConfig{}
	_ = serviceconfig.InitializeConfig(cfg)
	test.CreateDB(t, cfg)
	test.CreateAWSResources(t, cfg)

	// Create test client
	clientID := "test_single_batch"
	err := cfg.GetDBQueries().CreateClient(ctx, &repository.CreateClientParams{
		Clientid: clientID,
		Name:     "Test Single Batch Client",
	})
	require.NoError(t, err)

	// Create services
	batchService := batch.New(cfg)
	s3ClientInterface := cfg.GetStoreClient()
	s3Client, ok := s3ClientInterface.(*s3.Client)
	require.True(t, ok, "s3Client should be *s3.Client")
	bucket := cfg.GetBucket()

	// Upload test ZIP with mixed content
	zipContent := createMixedContentZIP(t)
	archiveKey := fmt.Sprintf("test/%s/mixed.zip", clientID)

	_, err = s3Client.PutObject(ctx, &s3.PutObjectInput{
		Bucket: aws.String(bucket),
		Key:    aws.String(archiveKey),
		Body:   bytes.NewReader(zipContent),
	})
	require.NoError(t, err)

	// Create batch
	batchID, err := batchService.CreateWithStorage(ctx, clientID, "mixed.zip", 3, bucket, archiveKey, int64(len(zipContent)))
	require.NoError(t, err)

	// Mock upload handler
	uploadCount := 0
	uploadHandler := func(ctx context.Context, clientID string, docData io.Reader, filename string, batchID *uuid.UUID) error {
		uploadCount++
		return nil
	}

	// Create batch details
	batchDetails := &batch.BatchUploadDetails{
		BatchUploadSummary: batch.BatchUploadSummary{
			ID:               batchID,
			ClientID:         clientID,
			OriginalFilename: "mixed.zip",
			Status:           "processing",
		},
		ArchiveKey:    archiveKey,
		ArchiveBucket: bucket,
	}

	// Process the batch
	logger := slog.New(slog.NewTextHandler(os.Stdout, nil))
	err = processSingleBatch(ctx, logger, batchService, s3Client, bucket, uploadHandler, batchDetails)
	require.NoError(t, err)

	// Verify only PDFs were uploaded
	assert.Equal(t, 1, uploadCount) // Only 1 PDF file

	// Verify batch status
	batchInfo, err := batchService.Get(ctx, clientID, batchID)
	require.NoError(t, err)
	assert.Equal(t, int32(1), batchInfo.ProcessedDocuments)
	assert.Equal(t, int32(2), batchInfo.InvalidTypeDocuments) // 1 txt + 1 docx
}

// TestProcessSingleBatchWithFailure tests handling of upload failures
func TestProcessSingleBatchWithFailure(t *testing.T) {
	if testing.Short() {
		t.SkipNow()
	}

	ctx := t.Context()
	cfg := &TestConfig{}
	_ = serviceconfig.InitializeConfig(cfg)
	test.CreateDB(t, cfg)
	test.CreateAWSResources(t, cfg)

	// Create test client
	clientID := "test_batch_failure"
	err := cfg.GetDBQueries().CreateClient(ctx, &repository.CreateClientParams{
		Clientid: clientID,
		Name:     "Test Batch Failure Client",
	})
	require.NoError(t, err)

	// Create services
	batchService := batch.New(cfg)
	s3ClientInterface := cfg.GetStoreClient()
	s3Client, ok := s3ClientInterface.(*s3.Client)
	require.True(t, ok, "s3Client should be *s3.Client")
	bucket := cfg.GetBucket()

	// Upload test ZIP
	zipContent := createTestZIPForWorker(t, 3)
	archiveKey := fmt.Sprintf("test/%s/fail.zip", clientID)

	_, err = s3Client.PutObject(ctx, &s3.PutObjectInput{
		Bucket: aws.String(bucket),
		Key:    aws.String(archiveKey),
		Body:   bytes.NewReader(zipContent),
	})
	require.NoError(t, err)

	// Create batch
	batchID, err := batchService.CreateWithStorage(ctx, clientID, "fail.zip", 3, bucket, archiveKey, int64(len(zipContent)))
	require.NoError(t, err)

	// Upload handler that fails on second document
	uploadCount := 0
	uploadHandler := func(ctx context.Context, clientID string, docData io.Reader, filename string, batchID *uuid.UUID) error {
		uploadCount++
		if uploadCount == 2 {
			return fmt.Errorf("simulated upload failure")
		}
		return nil
	}

	// Create batch details
	batchDetails := &batch.BatchUploadDetails{
		BatchUploadSummary: batch.BatchUploadSummary{
			ID:               batchID,
			ClientID:         clientID,
			OriginalFilename: "fail.zip",
			Status:           "processing",
		},
		ArchiveKey:    archiveKey,
		ArchiveBucket: bucket,
	}

	// Process the batch
	logger := slog.New(slog.NewTextHandler(os.Stdout, nil))
	err = processSingleBatch(ctx, logger, batchService, s3Client, bucket, uploadHandler, batchDetails)
	require.NoError(t, err)

	// Verify results
	batchInfo, err := batchService.Get(ctx, clientID, batchID)
	require.NoError(t, err)
	assert.Equal(t, int32(2), batchInfo.ProcessedDocuments) // 1st and 3rd succeeded
	assert.Equal(t, int32(1), batchInfo.FailedDocuments)    // 2nd failed
	assert.Equal(t, "completed", batchInfo.Status)          // Status is still completed since some succeeded
	assert.Contains(t, batchInfo.FailedFilenames, "document_2.pdf")
}

// TestDownloadZipFromS3 tests downloading a ZIP file from S3
func TestDownloadZipFromS3(t *testing.T) {
	if testing.Short() {
		t.SkipNow()
	}

	ctx := t.Context()
	cfg := &TestConfig{}
	_ = serviceconfig.InitializeConfig(cfg)
	test.CreateAWSResources(t, cfg)

	s3ClientInterface := cfg.GetStoreClient()
	s3Client, ok := s3ClientInterface.(*s3.Client)
	require.True(t, ok, "s3Client should be *s3.Client")
	bucket := cfg.GetBucket()

	// Upload test data
	testData := []byte("test zip content")
	key := "test/download.zip"

	_, err := s3Client.PutObject(ctx, &s3.PutObjectInput{
		Bucket: aws.String(bucket),
		Key:    aws.String(key),
		Body:   bytes.NewReader(testData),
	})
	require.NoError(t, err)

	// Download the data
	downloaded, err := downloadZipFromS3(ctx, s3Client, bucket, key)
	require.NoError(t, err)
	assert.Equal(t, testData, downloaded)
}

// TestUploadToS3 tests uploading data to S3
func TestUploadToS3(t *testing.T) {
	if testing.Short() {
		t.SkipNow()
	}

	ctx := t.Context()
	cfg := &TestConfig{}
	_ = serviceconfig.InitializeConfig(cfg)
	test.CreateAWSResources(t, cfg)

	s3ClientInterface := cfg.GetStoreClient()
	s3Client, ok := s3ClientInterface.(*s3.Client)
	require.True(t, ok, "s3Client should be *s3.Client")
	bucket := cfg.GetBucket()

	// Upload test data
	testData := []byte("test upload content")
	key := "test/upload.txt"

	err := uploadToS3(ctx, s3Client, bucket, key, testData)
	require.NoError(t, err)

	// Verify the upload
	result, err := s3Client.GetObject(ctx, &s3.GetObjectInput{
		Bucket: aws.String(bucket),
		Key:    aws.String(key),
	})
	require.NoError(t, err)
	defer result.Body.Close()

	downloaded, err := io.ReadAll(result.Body)
	require.NoError(t, err)
	assert.Equal(t, testData, downloaded)
}

// Helper function to create a test ZIP file with PDFs
func createTestZIPForWorker(t *testing.T, numFiles int) []byte {
	var buf bytes.Buffer
	zipWriter := zip.NewWriter(&buf)

	for i := 1; i <= numFiles; i++ {
		filename := fmt.Sprintf("document_%d.pdf", i)
		fileWriter, err := zipWriter.Create(filename)
		require.NoError(t, err)

		pdfContent := fmt.Sprintf("%%PDF-1.4\n%%Test content for %s\n%%%%EOF", filename)
		_, err = fileWriter.Write([]byte(pdfContent))
		require.NoError(t, err)
	}

	err := zipWriter.Close()
	require.NoError(t, err)

	return buf.Bytes()
}

// Helper function to create a ZIP with mixed content types
func createMixedContentZIP(t *testing.T) []byte {
	var buf bytes.Buffer
	zipWriter := zip.NewWriter(&buf)

	// Add a PDF
	pdfWriter, err := zipWriter.Create("document.pdf")
	require.NoError(t, err)
	_, err = pdfWriter.Write([]byte("%%PDF-1.4\n%%Test PDF\n%%%%EOF"))
	require.NoError(t, err)

	// Add a TXT file
	txtWriter, err := zipWriter.Create("notes.txt")
	require.NoError(t, err)
	_, err = txtWriter.Write([]byte("This is a text file"))
	require.NoError(t, err)

	// Add a DOCX file
	docxWriter, err := zipWriter.Create("report.docx")
	require.NoError(t, err)
	_, err = docxWriter.Write([]byte("PK")) // Simplified DOCX header
	require.NoError(t, err)

	err = zipWriter.Close()
	require.NoError(t, err)

	return buf.Bytes()
}