package api import ( "archive/zip" "bytes" "context" "crypto/sha256" "encoding/hex" "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" ) // TestProcessBatchWork tests the main batch processing work function func TestProcessBatchWork(t *testing.T) { if testing.Short() { t.SkipNow() } ctx := t.Context() cfg := &BaseConfig{} _ = serviceconfig.InitializeConfig(cfg) test.CreateDB(t, cfg) test.CreateAWSResources(t, cfg) // Create test client with root folder clientID := "test_batch_worker" test.CreateTestClient(t, cfg, clientID, "Test Batch Worker Client") // 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 := &BaseConfig{} _ = serviceconfig.InitializeConfig(cfg) test.CreateDB(t, cfg) test.CreateAWSResources(t, cfg) // Create test client with root folder clientID := "test_single_batch" test.CreateTestClient(t, cfg, clientID, "Test Single Batch Client") // 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 := &BaseConfig{} _ = serviceconfig.InitializeConfig(cfg) test.CreateDB(t, cfg) test.CreateAWSResources(t, cfg) // Create test client with root folder clientID := "test_batch_failure" test.CreateTestClient(t, cfg, clientID, "Test Batch Failure Client") // 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 := &BaseConfig{} _ = 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 := &BaseConfig{} _ = 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() } // Helper function to create a ZIP with nested folder structure func createNestedFolderZIP(t *testing.T) []byte { var buf bytes.Buffer zipWriter := zip.NewWriter(&buf) // Create files in various nested folder structures testFiles := []struct { path string content string }{ {"document1.pdf", "%%PDF-1.4\n%%Root level PDF\n%%%%EOF"}, {"foldera/file1.pdf", "%%PDF-1.4\n%%FolderA File1\n%%%%EOF"}, {"foldera/file2.pdf", "%%PDF-1.4\n%%FolderA File2\n%%%%EOF"}, {"folderb/folderc/file3.pdf", "%%PDF-1.4\n%%Nested File3\n%%%%EOF"}, {"folderb/folderc/file4.pdf", "%%PDF-1.4\n%%Nested File4\n%%%%EOF"}, {"folderb/file5.pdf", "%%PDF-1.4\n%%FolderB File5\n%%%%EOF"}, {"deep/nested/structure/test/file6.pdf", "%%PDF-1.4\n%%Deep nested\n%%%%EOF"}, // Add a non-PDF to verify it's skipped {"foldera/readme.txt", "This should be skipped"}, } for _, tf := range testFiles { fileWriter, err := zipWriter.Create(tf.path) require.NoError(t, err) _, err = fileWriter.Write([]byte(tf.content)) require.NoError(t, err) } err := zipWriter.Close() require.NoError(t, err) return buf.Bytes() } // Helper function to create a ZIP with path traversal attempts func createMaliciousPathZIP(t *testing.T) []byte { var buf bytes.Buffer zipWriter := zip.NewWriter(&buf) // Create files with various path traversal attempts testFiles := []struct { path string content string }{ {"../escape.pdf", "%%PDF-1.4\n%%Escape attempt\n%%%%EOF"}, {"../../etc/passwd", "malicious content"}, {"./valid.pdf", "%%PDF-1.4\n%%Valid PDF\n%%%%EOF"}, {"folder/../sibling.pdf", "%%PDF-1.4\n%%Sibling\n%%%%EOF"}, {"normal/file.pdf", "%%PDF-1.4\n%%Normal\n%%%%EOF"}, } for _, tf := range testFiles { // Note: zip.Writer doesn't validate paths, so we can create these entries header := &zip.FileHeader{ Name: tf.path, Method: zip.Deflate, } fileWriter, err := zipWriter.CreateHeader(header) require.NoError(t, err) _, err = fileWriter.Write([]byte(tf.content)) require.NoError(t, err) } err := zipWriter.Close() require.NoError(t, err) return buf.Bytes() } // TestSanitizeZipPath tests the path sanitization function func TestSanitizeZipPath(t *testing.T) { tests := []struct { name string input string expected string }{ {"simple file", "document.pdf", "document.pdf"}, {"file in folder", "folder/document.pdf", "folder/document.pdf"}, {"nested folders", "folder/subfolder/document.pdf", "folder/subfolder/document.pdf"}, {"deep nesting", "a/b/c/d/e/f/document.pdf", "a/b/c/d/e/f/document.pdf"}, {"dot prefix cleaned", "./document.pdf", "document.pdf"}, {"dot in path cleaned", "folder/./document.pdf", "folder/document.pdf"}, {"parent traversal rejected", "../document.pdf", ""}, {"parent in middle cleaned", "folder/../document.pdf", "document.pdf"}, {"absolute path rejected", "/etc/passwd", ""}, {"backslash rejected", "folder\\document.pdf", ""}, {"double parent rejected", "../../document.pdf", ""}, {"complex traversal attempt", "folder/../../etc/passwd", ""}, } for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { result := sanitizeZipPath(tt.input) assert.Equal(t, tt.expected, result, "sanitizeZipPath(%q) = %q, want %q", tt.input, result, tt.expected) }) } } // TestProcessBatchWithNestedFolders tests processing a ZIP with nested folder structure func TestProcessBatchWithNestedFolders(t *testing.T) { if testing.Short() { t.SkipNow() } ctx := t.Context() cfg := &BaseConfig{} _ = serviceconfig.InitializeConfig(cfg) test.CreateDB(t, cfg) test.CreateAWSResources(t, cfg) // Create test client with root folder clientID := "test_nested_folders" test.CreateTestClient(t, cfg, clientID, "Test Nested Folders Client") // 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 nested folders zipContent := createNestedFolderZIP(t) archiveKey := fmt.Sprintf("test/%s/nested.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, "nested.zip", 7, bucket, archiveKey, int64(len(zipContent))) require.NoError(t, err) // Use the real upload handler that creates documents in the database uploadHandler := createDocumentUploadHandler(cfg) // Create batch details batchDetails := &batch.BatchUploadDetails{ BatchUploadSummary: batch.BatchUploadSummary{ ID: batchID, ClientID: clientID, OriginalFilename: "nested.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 batch status batchInfo, err := batchService.Get(ctx, clientID, batchID) require.NoError(t, err) assert.Equal(t, int32(7), batchInfo.ProcessedDocuments) // 7 PDFs assert.Equal(t, int32(1), batchInfo.InvalidTypeDocuments) // 1 txt file assert.Equal(t, "completed", batchInfo.Status) // Manually trigger document initialization for uploaded files since S3 notifications // don't work in test environment. Get all document uploads for this batch. uploads, err := cfg.GetDBQueries().ListDocumentUploads(ctx, clientID) require.NoError(t, err) // Process each uploaded file through the document initialization pipeline for _, upload := range uploads { if upload.BatchID != nil && *upload.BatchID == batchID { // Parse the bucket key key, err := objectstore.ParseBucketKey(upload.Key) require.NoError(t, err) // Get file hash from S3 object metadata (normally provided by S3 event) headResult, err := s3Client.HeadObject(ctx, &s3.HeadObjectInput{ Bucket: aws.String(upload.Bucket), Key: aws.String(upload.Key), }) require.NoError(t, err) // Create document directly in database since full pipeline requires additional setup _, err = cfg.GetDBQueries().CreateDocument(ctx, &repository.CreateDocumentParams{ Clientid: key.ClientID, Hash: *headResult.ETag, // Use ETag as hash for testing BatchID: upload.BatchID, Filename: upload.Filename, }) require.NoError(t, err) } } // Get documents for this specific batch from database batchDocs, err := cfg.GetDBQueries().ListDocumentsByBatch(ctx, &batchID) require.NoError(t, err) // Verify we have the expected number of documents assert.Equal(t, 7, len(batchDocs), "Should have exactly 7 documents in database") // Verify filenames are preserved with their paths foundFilenames := make(map[string]bool) for _, doc := range batchDocs { if doc.Filename != nil { foundFilenames[*doc.Filename] = true } } // Check that we found all expected filenames with their paths expectedFilenames := []string{ "document1.pdf", "foldera/file1.pdf", "foldera/file2.pdf", "folderb/folderc/file3.pdf", "folderb/folderc/file4.pdf", "folderb/file5.pdf", "deep/nested/structure/test/file6.pdf", } for _, expectedFilename := range expectedFilenames { assert.True(t, foundFilenames[expectedFilename], "Should have preserved filename with path: %s", expectedFilename) } } // TestProcessBatchWithPathTraversal tests that path traversal attempts are handled safely func TestProcessBatchWithPathTraversal(t *testing.T) { if testing.Short() { t.SkipNow() } ctx := t.Context() cfg := &BaseConfig{} _ = serviceconfig.InitializeConfig(cfg) test.CreateDB(t, cfg) test.CreateAWSResources(t, cfg) // Create test client with root folder clientID := "test_path_traversal" test.CreateTestClient(t, cfg, clientID, "Test Path Traversal Client") // 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 path traversal attempts zipContent := createMaliciousPathZIP(t) archiveKey := fmt.Sprintf("test/%s/malicious.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, "malicious.zip", 5, 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 batch details batchDetails := &batch.BatchUploadDetails{ BatchUploadSummary: batch.BatchUploadSummary{ ID: batchID, ClientID: clientID, OriginalFilename: "malicious.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 that path traversal attempts were blocked assert.NotContains(t, uploadedDocs, "../escape.pdf", "Should reject parent traversal") assert.NotContains(t, uploadedDocs, "../../etc/passwd", "Should reject multiple parent traversal") assert.NotContains(t, uploadedDocs, "/etc/passwd", "Should reject absolute paths") // Verify that safe files were processed correctly assert.Contains(t, uploadedDocs, "valid.pdf", "Should process ./valid.pdf as valid.pdf") assert.Contains(t, uploadedDocs, "sibling.pdf", "Should clean folder/../sibling.pdf to sibling.pdf") assert.Contains(t, uploadedDocs, "normal/file.pdf", "Should process normal paths correctly") // Verify only safe PDFs were processed assert.Equal(t, 3, len(uploadedDocs), "Should have processed exactly 3 safe PDFs") } // TestProcessZipFile_RecordsOutcomes tests that processZipFile records per-file // outcomes for each extracted file. func TestProcessZipFile_RecordsOutcomes(t *testing.T) { if testing.Short() { t.SkipNow() } ctx := t.Context() cfg := &BaseConfig{} _ = serviceconfig.InitializeConfig(cfg) test.CreateDB(t, cfg) test.CreateAWSResources(t, cfg) clientID := "test_outcomes_record" test.CreateTestClient(t, cfg, clientID, "Test Outcomes Record Client") batchService := batch.New(cfg) s3ClientInterface := cfg.GetStoreClient() s3Client, ok := s3ClientInterface.(*s3.Client) require.True(t, ok) bucket := cfg.GetBucket() // Create ZIP with 2 PDFs zipContent := createTestZIPForWorker(t, 2) archiveKey := fmt.Sprintf("test/%s/outcomes.zip", clientID) _, err := s3Client.PutObject(ctx, &s3.PutObjectInput{ Bucket: aws.String(bucket), Key: aws.String(archiveKey), Body: bytes.NewReader(zipContent), }) require.NoError(t, err) batchID, err := batchService.CreateWithStorage(ctx, clientID, "outcomes.zip", 2, bucket, archiveKey, int64(len(zipContent))) require.NoError(t, err) uploadHandler := func(ctx context.Context, clientID string, docData io.Reader, filename string, batchID *uuid.UUID) error { _, _ = io.ReadAll(docData) return nil } workerConfig := map[string]any{ ConfigKeyBatchService: batchService, ConfigKeyS3Client: s3Client, ConfigKeyBucket: bucket, ConfigKeyUploadHandler: uploadHandler, } logger := slog.New(slog.NewTextHandler(os.Stdout, nil)) err = processBatchWork(ctx, logger, workerConfig) require.NoError(t, err) // Verify outcomes were recorded for each file outcomes, err := batchService.ListOutcomes(ctx, batchID) require.NoError(t, err) assert.Len(t, outcomes, 2) for _, o := range outcomes { assert.Equal(t, "submitted", o.Outcome, "each PDF should be recorded as submitted") assert.Contains(t, o.Filename, ".pdf") } } // TestProcessZipFile_DetectsDuplicates verifies that when a file in the batch // has the same content hash as an existing document in the system, it is // recorded as a duplicate with the existing document's ID. func TestProcessZipFile_DetectsDuplicates(t *testing.T) { if testing.Short() { t.SkipNow() } ctx := t.Context() cfg := &BaseConfig{} _ = serviceconfig.InitializeConfig(cfg) test.CreateDB(t, cfg) test.CreateAWSResources(t, cfg) clientID := "test_outcomes_dup" test.CreateTestClient(t, cfg, clientID, "Test Outcomes Dup Client") batchService := batch.New(cfg) s3ClientInterface := cfg.GetStoreClient() s3Client, ok := s3ClientInterface.(*s3.Client) require.True(t, ok) bucket := cfg.GetBucket() // Pre-create an existing document with a known hash knownContent := []byte("%%PDF-1.4\n%%Already existing content\n%%%%EOF") knownHash := sha256Hash(knownContent) existingDocID, err := cfg.GetDBQueries().CreateDocument(ctx, &repository.CreateDocumentParams{ Clientid: clientID, Hash: knownHash, }) require.NoError(t, err) // Create a ZIP with one file that has the same content as the existing doc var buf bytes.Buffer zipWriter := zip.NewWriter(&buf) w, err := zipWriter.Create("duplicate_of_existing.pdf") require.NoError(t, err) _, err = w.Write(knownContent) require.NoError(t, err) require.NoError(t, zipWriter.Close()) zipContent := buf.Bytes() archiveKey := fmt.Sprintf("test/%s/dup.zip", clientID) _, err = s3Client.PutObject(ctx, &s3.PutObjectInput{ Bucket: aws.String(bucket), Key: aws.String(archiveKey), Body: bytes.NewReader(zipContent), }) require.NoError(t, err) batchID, err := batchService.CreateWithStorage(ctx, clientID, "dup.zip", 1, bucket, archiveKey, int64(len(zipContent))) require.NoError(t, err) uploadHandler := func(ctx context.Context, cID string, docData io.Reader, filename string, bID *uuid.UUID) error { _, _ = io.ReadAll(docData) return nil } workerConfig := map[string]any{ ConfigKeyBatchService: batchService, ConfigKeyS3Client: s3Client, ConfigKeyBucket: bucket, ConfigKeyUploadHandler: uploadHandler, } logger := slog.New(slog.NewTextHandler(os.Stdout, nil)) err = processBatchWork(ctx, logger, workerConfig) require.NoError(t, err) // Verify outcome is "duplicate" with the existing document's ID outcomes, err := batchService.ListOutcomes(ctx, batchID) require.NoError(t, err) require.Len(t, outcomes, 1) assert.Equal(t, "duplicate", outcomes[0].Outcome) require.NotNil(t, outcomes[0].DocumentID) assert.Equal(t, existingDocID, *outcomes[0].DocumentID) } // TestProcessZipFile_MixedOutcomes tests a ZIP with PDF + non-PDF + paths to // verify all outcome types (submitted, invalid_type) are correctly recorded. func TestProcessZipFile_MixedOutcomes(t *testing.T) { if testing.Short() { t.SkipNow() } ctx := t.Context() cfg := &BaseConfig{} _ = serviceconfig.InitializeConfig(cfg) test.CreateDB(t, cfg) test.CreateAWSResources(t, cfg) clientID := "test_outcomes_mixed" test.CreateTestClient(t, cfg, clientID, "Test Outcomes Mixed Client") batchService := batch.New(cfg) s3ClientInterface := cfg.GetStoreClient() s3Client, ok := s3ClientInterface.(*s3.Client) require.True(t, ok) bucket := cfg.GetBucket() // Create mixed ZIP zipContent := createMixedContentZIP(t) archiveKey := fmt.Sprintf("test/%s/mixed_outcomes.zip", clientID) _, err := s3Client.PutObject(ctx, &s3.PutObjectInput{ Bucket: aws.String(bucket), Key: aws.String(archiveKey), Body: bytes.NewReader(zipContent), }) require.NoError(t, err) batchID, err := batchService.CreateWithStorage(ctx, clientID, "mixed_outcomes.zip", 3, bucket, archiveKey, int64(len(zipContent))) require.NoError(t, err) uploadHandler := func(ctx context.Context, cID string, docData io.Reader, filename string, bID *uuid.UUID) error { _, _ = io.ReadAll(docData) return nil } workerConfig := map[string]any{ ConfigKeyBatchService: batchService, ConfigKeyS3Client: s3Client, ConfigKeyBucket: bucket, ConfigKeyUploadHandler: uploadHandler, } logger := slog.New(slog.NewTextHandler(os.Stdout, nil)) err = processBatchWork(ctx, logger, workerConfig) require.NoError(t, err) // Verify outcomes outcomes, err := batchService.ListOutcomes(ctx, batchID) require.NoError(t, err) assert.Len(t, outcomes, 3) // 1 PDF + 1 txt + 1 docx outcomeMap := map[string]string{} for _, o := range outcomes { outcomeMap[o.Filename] = o.Outcome } assert.Equal(t, "submitted", outcomeMap["document.pdf"]) assert.Equal(t, "invalid_type", outcomeMap["notes.txt"]) assert.Equal(t, "invalid_type", outcomeMap["report.docx"]) } // sha256Hash computes the hex-encoded SHA-256 hash of data. func sha256Hash(data []byte) string { h := sha256.Sum256(data) return hex.EncodeToString(h[:]) } // TestProcessBatchWithDuplicateFilenames tests handling of duplicate filenames in different folders func TestProcessBatchWithDuplicateFilenames(t *testing.T) { if testing.Short() { t.SkipNow() } ctx := t.Context() cfg := &BaseConfig{} _ = serviceconfig.InitializeConfig(cfg) test.CreateDB(t, cfg) test.CreateAWSResources(t, cfg) // Create test client with root folder clientID := "test_duplicate_names" test.CreateTestClient(t, cfg, clientID, "Test Duplicate Names Client") // 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() // Create ZIP with duplicate filenames in different folders var buf bytes.Buffer zipWriter := zip.NewWriter(&buf) // Create files with same names in different folders testFiles := []struct { path string content string }{ {"document.pdf", "%%PDF-1.4\n%%Root document\n%%%%EOF"}, {"folder1/document.pdf", "%%PDF-1.4\n%%Folder1 document\n%%%%EOF"}, {"folder2/document.pdf", "%%PDF-1.4\n%%Folder2 document\n%%%%EOF"}, {"folder1/subfolder/document.pdf", "%%PDF-1.4\n%%Subfolder document\n%%%%EOF"}, } for _, tf := range testFiles { fileWriter, err := zipWriter.Create(tf.path) require.NoError(t, err) _, err = fileWriter.Write([]byte(tf.content)) require.NoError(t, err) } err := zipWriter.Close() require.NoError(t, err) zipContent := buf.Bytes() // Upload the ZIP archiveKey := fmt.Sprintf("test/%s/duplicates.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, "duplicates.zip", 4, bucket, archiveKey, int64(len(zipContent))) require.NoError(t, err) // Track uploaded documents uploadedDocs := make(map[string]string) 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 } // Extract the content identifier from the PDF content := string(data) uploadedDocs[filename] = content return nil } // Create batch details batchDetails := &batch.BatchUploadDetails{ BatchUploadSummary: batch.BatchUploadSummary{ ID: batchID, ClientID: clientID, OriginalFilename: "duplicates.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 all files were uploaded with correct paths assert.Equal(t, 4, len(uploadedDocs), "Should have uploaded all 4 PDFs") // Verify each file has unique path and correct content assert.Contains(t, uploadedDocs["document.pdf"], "Root document") assert.Contains(t, uploadedDocs["folder1/document.pdf"], "Folder1 document") assert.Contains(t, uploadedDocs["folder2/document.pdf"], "Folder2 document") assert.Contains(t, uploadedDocs["folder1/subfolder/document.pdf"], "Subfolder document") }