jay brown
423 lines
13 KiB
Go
423 lines
13 KiB
Go
package batch_test
|
|
|
|
import (
|
|
"archive/zip"
|
|
"bytes"
|
|
"fmt"
|
|
"io"
|
|
"testing"
|
|
|
|
"queryorchestration/internal/database/repository"
|
|
documentbatch "queryorchestration/internal/document/batch"
|
|
"queryorchestration/internal/document/batch/storage"
|
|
"queryorchestration/internal/serviceconfig"
|
|
"queryorchestration/internal/serviceconfig/objectstore"
|
|
"queryorchestration/internal/test"
|
|
|
|
"github.com/aws/aws-sdk-go-v2/service/s3"
|
|
"github.com/stretchr/testify/assert"
|
|
"github.com/stretchr/testify/require"
|
|
)
|
|
|
|
type TestConfig struct {
|
|
serviceconfig.BaseConfig
|
|
objectstore.ObjectStoreConfig
|
|
}
|
|
|
|
func TestUpdateProgress(t *testing.T) {
|
|
cfg := &TestConfig{}
|
|
test.CreateDB(t, cfg)
|
|
|
|
service := documentbatch.New(cfg)
|
|
|
|
// Create test client
|
|
clientID := "test_client_progress"
|
|
err := cfg.GetDBQueries().CreateClient(t.Context(), &repository.CreateClientParams{
|
|
Clientid: clientID,
|
|
Name: "Test Client Progress",
|
|
})
|
|
require.NoError(t, err)
|
|
|
|
// Create test batch with total documents
|
|
batchID, err := cfg.GetDBQueries().CreateBatchUpload(t.Context(), &repository.CreateBatchUploadParams{
|
|
ClientID: clientID,
|
|
OriginalFilename: "test.zip",
|
|
TotalDocuments: 10,
|
|
})
|
|
require.NoError(t, err)
|
|
|
|
// Test updating progress
|
|
err = service.UpdateProgress(t.Context(), clientID, batchID, 5, 2, 1)
|
|
require.NoError(t, err)
|
|
|
|
// Verify progress was updated
|
|
batch, err := cfg.GetDBQueries().GetBatchUpload(t.Context(), &repository.GetBatchUploadParams{
|
|
ID: batchID,
|
|
ClientID: clientID,
|
|
})
|
|
require.NoError(t, err)
|
|
assert.Equal(t, int32(5), batch.ProcessedDocuments)
|
|
assert.Equal(t, int32(2), batch.FailedDocuments)
|
|
assert.Equal(t, int32(1), batch.InvalidTypeDocuments)
|
|
assert.Equal(t, int32(80), batch.ProgressPercent) // (5+2+1)/10 * 100 = 80%
|
|
}
|
|
|
|
func TestUpdateProgressZeroTotalDocuments(t *testing.T) {
|
|
cfg := &TestConfig{}
|
|
test.CreateDB(t, cfg)
|
|
|
|
service := documentbatch.New(cfg)
|
|
|
|
// Create test client
|
|
clientID := "test_client_zero"
|
|
err := cfg.GetDBQueries().CreateClient(t.Context(), &repository.CreateClientParams{
|
|
Clientid: clientID,
|
|
Name: "Test Client Zero",
|
|
})
|
|
require.NoError(t, err)
|
|
|
|
// Create test batch with zero total documents
|
|
batchID, err := cfg.GetDBQueries().CreateBatchUpload(t.Context(), &repository.CreateBatchUploadParams{
|
|
ClientID: clientID,
|
|
OriginalFilename: "test.zip",
|
|
TotalDocuments: 0,
|
|
})
|
|
require.NoError(t, err)
|
|
|
|
// Test updating progress should fail with zero total documents
|
|
err = service.UpdateProgress(t.Context(), clientID, batchID, 5, 2, 1)
|
|
require.Error(t, err)
|
|
assert.Contains(t, err.Error(), "batch has zero total documents")
|
|
}
|
|
|
|
func TestMarkCompleted(t *testing.T) {
|
|
cfg := &TestConfig{}
|
|
test.CreateDB(t, cfg)
|
|
|
|
service := documentbatch.New(cfg)
|
|
|
|
// Create test client
|
|
clientID := "test_client_complete"
|
|
err := cfg.GetDBQueries().CreateClient(t.Context(), &repository.CreateClientParams{
|
|
Clientid: clientID,
|
|
Name: "Test Client Complete",
|
|
})
|
|
require.NoError(t, err)
|
|
|
|
// Create test batch
|
|
batchID, err := cfg.GetDBQueries().CreateBatchUpload(t.Context(), &repository.CreateBatchUploadParams{
|
|
ClientID: clientID,
|
|
OriginalFilename: "test.zip",
|
|
TotalDocuments: 10,
|
|
})
|
|
require.NoError(t, err)
|
|
|
|
// Test marking as completed
|
|
err = service.MarkCompleted(t.Context(), batchID)
|
|
require.NoError(t, err)
|
|
|
|
// Verify status was updated
|
|
batch, err := cfg.GetDBQueries().GetBatchUpload(t.Context(), &repository.GetBatchUploadParams{
|
|
ID: batchID,
|
|
ClientID: clientID,
|
|
})
|
|
require.NoError(t, err)
|
|
assert.Equal(t, repository.BatchStatusCompleted, batch.Status)
|
|
}
|
|
|
|
func TestMarkFailed(t *testing.T) {
|
|
cfg := &TestConfig{}
|
|
test.CreateDB(t, cfg)
|
|
|
|
service := documentbatch.New(cfg)
|
|
|
|
// Create test client
|
|
clientID := "test_client_failed"
|
|
err := cfg.GetDBQueries().CreateClient(t.Context(), &repository.CreateClientParams{
|
|
Clientid: clientID,
|
|
Name: "Test Client Failed",
|
|
})
|
|
require.NoError(t, err)
|
|
|
|
// Create test batch
|
|
batchID, err := cfg.GetDBQueries().CreateBatchUpload(t.Context(), &repository.CreateBatchUploadParams{
|
|
ClientID: clientID,
|
|
OriginalFilename: "test.zip",
|
|
TotalDocuments: 10,
|
|
})
|
|
require.NoError(t, err)
|
|
|
|
// Test marking as failed
|
|
err = service.MarkFailed(t.Context(), batchID)
|
|
require.NoError(t, err)
|
|
|
|
// Verify status was updated
|
|
batch, err := cfg.GetDBQueries().GetBatchUpload(t.Context(), &repository.GetBatchUploadParams{
|
|
ID: batchID,
|
|
ClientID: clientID,
|
|
})
|
|
require.NoError(t, err)
|
|
assert.Equal(t, repository.BatchStatusFailed, batch.Status)
|
|
}
|
|
|
|
func TestAddFailedFilename(t *testing.T) {
|
|
cfg := &TestConfig{}
|
|
test.CreateDB(t, cfg)
|
|
|
|
service := documentbatch.New(cfg)
|
|
|
|
// Create test client
|
|
clientID := "test_client_filename"
|
|
err := cfg.GetDBQueries().CreateClient(t.Context(), &repository.CreateClientParams{
|
|
Clientid: clientID,
|
|
Name: "Test Client Filename",
|
|
})
|
|
require.NoError(t, err)
|
|
|
|
// Create test batch
|
|
batchID, err := cfg.GetDBQueries().CreateBatchUpload(t.Context(), &repository.CreateBatchUploadParams{
|
|
ClientID: clientID,
|
|
OriginalFilename: "test.zip",
|
|
TotalDocuments: 10,
|
|
})
|
|
require.NoError(t, err)
|
|
|
|
// Test adding failed filename
|
|
err = service.AddFailedFilename(t.Context(), batchID, "failed_file.pdf")
|
|
require.NoError(t, err)
|
|
|
|
// Get the full details to check failed filenames
|
|
details, err := service.Get(t.Context(), clientID, batchID)
|
|
require.NoError(t, err)
|
|
assert.Contains(t, details.FailedFilenames, "failed_file.pdf")
|
|
}
|
|
|
|
func TestCreateWithStorage_Success(t *testing.T) {
|
|
// Setup: Create database testcontainer
|
|
cfg := &TestConfig{}
|
|
test.CreateDB(t, cfg)
|
|
service := documentbatch.New(cfg)
|
|
|
|
// Create test client
|
|
clientID := "test_client_storage"
|
|
err := cfg.GetDBQueries().CreateClient(t.Context(), &repository.CreateClientParams{
|
|
Clientid: clientID,
|
|
Name: "Test Client Storage",
|
|
})
|
|
require.NoError(t, err)
|
|
|
|
// Test: Create batch with storage metadata
|
|
batchID, err := service.CreateWithStorage(
|
|
t.Context(),
|
|
clientID,
|
|
"test.zip",
|
|
0, // totalDocs unknown in Part 1
|
|
"test-bucket",
|
|
"batches/test-client/2025/01/01/test.zip",
|
|
1024*1024, // 1MB
|
|
)
|
|
|
|
// Verify: Batch created successfully
|
|
require.NoError(t, err)
|
|
assert.NotEmpty(t, batchID)
|
|
|
|
// Verify: Batch can be retrieved with storage metadata
|
|
batch, err := service.Get(t.Context(), clientID, batchID)
|
|
require.NoError(t, err)
|
|
assert.Equal(t, "test.zip", batch.OriginalFilename)
|
|
assert.Equal(t, "processing", batch.Status)
|
|
assert.Equal(t, "test-bucket", batch.ArchiveBucket)
|
|
assert.Equal(t, "batches/test-client/2025/01/01/test.zip", batch.ArchiveKey)
|
|
assert.Equal(t, int64(1024*1024), batch.FileSizeBytes)
|
|
}
|
|
|
|
func TestCreateWithStorage_DatabaseFailure(t *testing.T) {
|
|
// Setup: Create database
|
|
cfg := &TestConfig{}
|
|
test.CreateDB(t, cfg)
|
|
_ = documentbatch.New(cfg) // service unused in skipped test
|
|
|
|
// Create test client
|
|
clientID := "test_client_storage_fail"
|
|
err := cfg.GetDBQueries().CreateClient(t.Context(), &repository.CreateClientParams{
|
|
Clientid: clientID,
|
|
Name: "Test Client Storage Fail",
|
|
})
|
|
require.NoError(t, err)
|
|
|
|
// Note: Skipping database connection close test as Pool interface doesn't expose Close method
|
|
// This test would require mocking the database layer to simulate connection failures
|
|
t.Skip("Database connection failure testing requires mocking")
|
|
}
|
|
|
|
// TestCreate_Success tests the Create function with a valid ZIP file uploaded to S3
|
|
func TestCreate_Success(t *testing.T) {
|
|
// Setup: Create database and AWS resources with testcontainers
|
|
cfg := &TestConfig{}
|
|
test.CreateDB(t, cfg)
|
|
test.CreateAWSResources(t, cfg)
|
|
|
|
// Create batch service and storage service
|
|
batchService := documentbatch.New(cfg)
|
|
storageService := storage.New(cfg)
|
|
|
|
// Create test client
|
|
clientID := "test_client_create"
|
|
err := cfg.GetDBQueries().CreateClient(t.Context(), &repository.CreateClientParams{
|
|
Clientid: clientID,
|
|
Name: "Test Client Create",
|
|
})
|
|
require.NoError(t, err)
|
|
|
|
// Create test ZIP content with 3 PDFs
|
|
zipContent := createTestZIPContent(t, 3)
|
|
|
|
// First upload the ZIP to S3 (simulating what would happen in the full flow)
|
|
storageResult, err := storageService.StoreZIPFile(t.Context(), clientID, "test.zip", zipContent)
|
|
require.NoError(t, err)
|
|
|
|
// Test: Use the Create function to create a batch record
|
|
batchID, err := batchService.Create(t.Context(), clientID, "test.zip", 3)
|
|
|
|
// Verify: Batch created successfully
|
|
require.NoError(t, err)
|
|
assert.NotEmpty(t, batchID)
|
|
|
|
// Verify: Batch can be retrieved
|
|
batch, err := cfg.GetDBQueries().GetBatchUpload(t.Context(), &repository.GetBatchUploadParams{
|
|
ID: batchID,
|
|
ClientID: clientID,
|
|
})
|
|
require.NoError(t, err)
|
|
assert.Equal(t, "test.zip", batch.OriginalFilename)
|
|
assert.Equal(t, int32(3), batch.TotalDocuments)
|
|
assert.Equal(t, clientID, batch.ClientID)
|
|
assert.Equal(t, repository.BatchStatusProcessing, batch.Status)
|
|
|
|
// Verify: File still exists in S3
|
|
s3Client := cfg.GetStoreClient()
|
|
_, err = s3Client.HeadObject(t.Context(), &s3.HeadObjectInput{
|
|
Bucket: &storageResult.Bucket,
|
|
Key: &storageResult.Key,
|
|
})
|
|
require.NoError(t, err, "ZIP file should exist in S3")
|
|
}
|
|
|
|
// TestCreate_InvalidZIP tests Create function with invalid ZIP data
|
|
func TestCreate_InvalidZIP(t *testing.T) {
|
|
// Setup: Create database and AWS resources
|
|
cfg := &TestConfig{}
|
|
test.CreateDB(t, cfg)
|
|
test.CreateAWSResources(t, cfg)
|
|
|
|
batchService := documentbatch.New(cfg)
|
|
storageService := storage.New(cfg)
|
|
|
|
// Create test client
|
|
clientID := "test_client_invalid_zip"
|
|
err := cfg.GetDBQueries().CreateClient(t.Context(), &repository.CreateClientParams{
|
|
Clientid: clientID,
|
|
Name: "Test Client Invalid ZIP",
|
|
})
|
|
require.NoError(t, err)
|
|
|
|
// Create invalid ZIP content (just plain text, not a ZIP)
|
|
invalidContent := bytes.NewReader([]byte("This is not a valid ZIP file"))
|
|
|
|
// Test: Try to validate before storing (storage service should validate)
|
|
err = storageService.ValidateZIPFile("notazip.txt", int64(invalidContent.Len()))
|
|
require.Error(t, err)
|
|
assert.Contains(t, err.Error(), "invalid file type")
|
|
|
|
// Even if we bypass validation and try with .zip extension, the content is invalid
|
|
// The Create function itself doesn't validate ZIP content, it just creates a DB record
|
|
// So we test that Create works but processing would fail later
|
|
batchID, err := batchService.Create(t.Context(), clientID, "fake.zip", 1)
|
|
|
|
// The Create function should succeed (it only creates DB record)
|
|
require.NoError(t, err)
|
|
assert.NotEmpty(t, batchID)
|
|
|
|
// Verify the batch exists but would fail during actual processing
|
|
batch, err := cfg.GetDBQueries().GetBatchUpload(t.Context(), &repository.GetBatchUploadParams{
|
|
ID: batchID,
|
|
ClientID: clientID,
|
|
})
|
|
require.NoError(t, err)
|
|
assert.Equal(t, "fake.zip", batch.OriginalFilename)
|
|
}
|
|
|
|
// TestCreate_S3UploadFailure tests Create function when S3 upload fails
|
|
func TestCreate_S3UploadFailure(t *testing.T) {
|
|
// Setup: Create database and AWS resources
|
|
cfg := &TestConfig{}
|
|
test.CreateDB(t, cfg)
|
|
test.CreateAWSResources(t, cfg)
|
|
|
|
batchService := documentbatch.New(cfg)
|
|
storageService := storage.New(cfg)
|
|
|
|
// Create test client
|
|
clientID := "test_client_s3_fail"
|
|
err := cfg.GetDBQueries().CreateClient(t.Context(), &repository.CreateClientParams{
|
|
Clientid: clientID,
|
|
Name: "Test Client S3 Fail",
|
|
})
|
|
require.NoError(t, err)
|
|
|
|
// Create valid ZIP content
|
|
zipContent := createTestZIPContent(t, 2)
|
|
|
|
// Temporarily break S3 configuration to simulate upload failure
|
|
// Save original bucket and set invalid one
|
|
originalBucket := cfg.GetBucket()
|
|
cfg.SetBucket("non-existent-bucket-that-should-not-exist")
|
|
|
|
// Test: Try to upload to S3 with invalid bucket
|
|
_, err = storageService.StoreZIPFile(t.Context(), clientID, "test.zip", zipContent)
|
|
|
|
// Verify: S3 upload should fail
|
|
require.Error(t, err)
|
|
assert.Contains(t, err.Error(), "failed to upload ZIP to S3")
|
|
|
|
// Restore original bucket for cleanup
|
|
cfg.SetBucket(originalBucket)
|
|
|
|
// The Create function itself is independent of S3, so it would still work
|
|
// This demonstrates separation of concerns
|
|
batchID, err := batchService.Create(t.Context(), clientID, "test.zip", 2)
|
|
require.NoError(t, err)
|
|
assert.NotEmpty(t, batchID)
|
|
|
|
// Verify batch exists in DB even though S3 upload failed
|
|
batch, err := cfg.GetDBQueries().GetBatchUpload(t.Context(), &repository.GetBatchUploadParams{
|
|
ID: batchID,
|
|
ClientID: clientID,
|
|
})
|
|
require.NoError(t, err)
|
|
assert.Equal(t, "test.zip", batch.OriginalFilename)
|
|
assert.Equal(t, int32(2), batch.TotalDocuments)
|
|
}
|
|
|
|
// Helper function to create test ZIP content with specified number of PDF files
|
|
func createTestZIPContent(t *testing.T, numFiles int) io.Reader {
|
|
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)
|
|
|
|
// Write minimal PDF content
|
|
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 bytes.NewReader(buf.Bytes())
|
|
}
|