Jay Brown
754 lines
23 KiB
Go
754 lines
23 KiB
Go
package batch_test
|
|
|
|
import (
|
|
"archive/zip"
|
|
"bytes"
|
|
"context"
|
|
"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/google/uuid"
|
|
"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, documentbatch.StatusProcessing, 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())
|
|
}
|
|
|
|
func TestListUnprocessed(t *testing.T) {
|
|
cfg := &TestConfig{}
|
|
test.CreateDB(t, cfg)
|
|
|
|
service := documentbatch.New(cfg)
|
|
ctx := t.Context()
|
|
|
|
// Create multiple test clients
|
|
clientID1 := "test_client_unprocessed1"
|
|
clientID2 := "test_client_unprocessed2"
|
|
|
|
for _, clientID := range []string{clientID1, clientID2} {
|
|
err := cfg.GetDBQueries().CreateClient(ctx, &repository.CreateClientParams{
|
|
Clientid: clientID,
|
|
Name: fmt.Sprintf("Test Client %s", clientID),
|
|
})
|
|
require.NoError(t, err)
|
|
}
|
|
|
|
// Create processing batches (should be returned)
|
|
var expectedBatchIDs []string
|
|
|
|
// Client 1: 2 processing batches
|
|
for i := 0; i < 2; i++ {
|
|
batchID, err := service.Create(ctx, clientID1, fmt.Sprintf("processing%d.zip", i), 10)
|
|
require.NoError(t, err)
|
|
expectedBatchIDs = append(expectedBatchIDs, batchID.String())
|
|
}
|
|
|
|
// Client 2: 1 processing batch
|
|
batchID, err := service.Create(ctx, clientID2, "processing_client2.zip", 5)
|
|
require.NoError(t, err)
|
|
expectedBatchIDs = append(expectedBatchIDs, batchID.String())
|
|
|
|
// Create completed batch (should NOT be returned)
|
|
completedBatchID, err := service.Create(ctx, clientID1, "completed.zip", 8)
|
|
require.NoError(t, err)
|
|
err = service.MarkCompleted(ctx, completedBatchID)
|
|
require.NoError(t, err)
|
|
|
|
// Create failed batch (should NOT be returned)
|
|
failedBatchID, err := service.Create(ctx, clientID2, "failed.zip", 3)
|
|
require.NoError(t, err)
|
|
err = service.MarkFailed(ctx, failedBatchID)
|
|
require.NoError(t, err)
|
|
|
|
// Test: Get unprocessed batches
|
|
unprocessedBatches, err := service.ListUnprocessed(ctx)
|
|
require.NoError(t, err)
|
|
|
|
// Verify: Should return exactly 3 processing batches
|
|
assert.Len(t, unprocessedBatches, 3)
|
|
|
|
// Verify: All returned batches have processing status
|
|
actualBatchIDs := make([]string, len(unprocessedBatches))
|
|
for i, batch := range unprocessedBatches {
|
|
assert.Equal(t, documentbatch.StatusProcessing, batch.Status)
|
|
actualBatchIDs[i] = batch.ID.String()
|
|
}
|
|
|
|
// Verify: All processing batches are returned
|
|
for _, expectedID := range expectedBatchIDs {
|
|
assert.Contains(t, actualBatchIDs, expectedID)
|
|
}
|
|
|
|
// Verify: Completed and failed batches are NOT returned
|
|
assert.NotContains(t, actualBatchIDs, completedBatchID.String())
|
|
assert.NotContains(t, actualBatchIDs, failedBatchID.String())
|
|
|
|
// Verify: Batches from both clients are returned
|
|
clientIDs := make(map[string]bool)
|
|
for _, batch := range unprocessedBatches {
|
|
clientIDs[batch.ClientID] = true
|
|
}
|
|
assert.True(t, clientIDs[clientID1], "should include batches from client1")
|
|
assert.True(t, clientIDs[clientID2], "should include batches from client2")
|
|
|
|
// 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.Before(unprocessedBatches[i+1].CreatedAt) ||
|
|
unprocessedBatches[i].CreatedAt.Equal(unprocessedBatches[i+1].CreatedAt),
|
|
"batches should be ordered by created_at ASC")
|
|
}
|
|
}
|
|
}
|
|
|
|
func TestListUnprocessed_EmptyResult(t *testing.T) {
|
|
cfg := &TestConfig{}
|
|
test.CreateDB(t, cfg)
|
|
|
|
service := documentbatch.New(cfg)
|
|
ctx := t.Context()
|
|
|
|
// Create test client
|
|
clientID := "test_client_empty"
|
|
err := cfg.GetDBQueries().CreateClient(ctx, &repository.CreateClientParams{
|
|
Clientid: clientID,
|
|
Name: "Test Client Empty",
|
|
})
|
|
require.NoError(t, err)
|
|
|
|
// Create only completed batches (none processing)
|
|
for i := 0; i < 2; i++ {
|
|
batchID, err := service.Create(ctx, clientID, fmt.Sprintf("completed%d.zip", i), 10)
|
|
require.NoError(t, err)
|
|
err = service.MarkCompleted(ctx, batchID)
|
|
require.NoError(t, err)
|
|
}
|
|
|
|
// Test: Get unprocessed batches when none exist
|
|
unprocessedBatches, err := service.ListUnprocessed(ctx)
|
|
require.NoError(t, err)
|
|
|
|
// Verify: Should return empty list
|
|
assert.Empty(t, unprocessedBatches)
|
|
}
|
|
|
|
func TestListUnprocessed_DatabaseError(t *testing.T) {
|
|
cfg := &TestConfig{}
|
|
test.CreateDB(t, cfg)
|
|
|
|
service := documentbatch.New(cfg)
|
|
|
|
// Use a canceled context to simulate database error
|
|
ctx, cancel := context.WithCancel(t.Context())
|
|
cancel() // Cancel immediately
|
|
|
|
// Test: Try to get unprocessed batches with canceled context
|
|
unprocessedBatches, err := service.ListUnprocessed(ctx)
|
|
|
|
// Verify: Should return error
|
|
require.Error(t, err)
|
|
assert.Nil(t, unprocessedBatches)
|
|
assert.Contains(t, err.Error(), "failed to query unprocessed batches")
|
|
}
|
|
|
|
// TestRecordOutcome verifies that RecordOutcome inserts a per-file outcome row
|
|
// and that it can be read back via ListOutcomes.
|
|
func TestRecordOutcome(t *testing.T) {
|
|
cfg := &TestConfig{}
|
|
test.CreateDB(t, cfg)
|
|
ctx := t.Context()
|
|
|
|
service := documentbatch.New(cfg)
|
|
|
|
clientID := "test_record_outcome_" + uuid.New().String()[:8]
|
|
err := cfg.GetDBQueries().CreateClient(ctx, &repository.CreateClientParams{
|
|
Clientid: clientID,
|
|
Name: "Test Client Record Outcome",
|
|
})
|
|
require.NoError(t, err)
|
|
|
|
batchID, err := service.Create(ctx, clientID, "outcomes.zip", 3)
|
|
require.NoError(t, err)
|
|
|
|
// Insert multiple outcomes
|
|
err = service.RecordOutcome(ctx, batchID, "good.pdf", repository.BatchOutcomeStatusSubmitted, nil, nil)
|
|
require.NoError(t, err)
|
|
|
|
errMsg := "file is not a PDF"
|
|
err = service.RecordOutcome(ctx, batchID, "notes.txt", repository.BatchOutcomeStatusInvalidType, &errMsg, nil)
|
|
require.NoError(t, err)
|
|
|
|
// Verify persistence via ListOutcomes
|
|
outcomes, err := service.ListOutcomes(ctx, batchID)
|
|
require.NoError(t, err)
|
|
require.Len(t, outcomes, 2)
|
|
|
|
// Verify the outcome details
|
|
found := map[string]string{}
|
|
for _, o := range outcomes {
|
|
found[o.Filename] = o.Outcome
|
|
}
|
|
assert.Equal(t, "submitted", found["good.pdf"])
|
|
assert.Equal(t, "invalid_type", found["notes.txt"])
|
|
}
|
|
|
|
// TestListOutcomes inserts outcomes for various states and verifies the listing.
|
|
func TestListOutcomes(t *testing.T) {
|
|
cfg := &TestConfig{}
|
|
test.CreateDB(t, cfg)
|
|
ctx := t.Context()
|
|
|
|
service := documentbatch.New(cfg)
|
|
|
|
clientID := "test_list_outcomes_" + uuid.New().String()[:8]
|
|
err := cfg.GetDBQueries().CreateClient(ctx, &repository.CreateClientParams{
|
|
Clientid: clientID,
|
|
Name: "Test Client List Outcomes",
|
|
})
|
|
require.NoError(t, err)
|
|
|
|
batchID, err := service.Create(ctx, clientID, "list.zip", 5)
|
|
require.NoError(t, err)
|
|
|
|
// Insert a variety of outcomes
|
|
err = service.RecordOutcome(ctx, batchID, "a.pdf", repository.BatchOutcomeStatusSubmitted, nil, nil)
|
|
require.NoError(t, err)
|
|
err = service.RecordOutcome(ctx, batchID, "b.pdf", repository.BatchOutcomeStatusSubmitted, nil, nil)
|
|
require.NoError(t, err)
|
|
|
|
errOpen := "cannot open"
|
|
err = service.RecordOutcome(ctx, batchID, "c.pdf", repository.BatchOutcomeStatusFailedOpen, &errOpen, nil)
|
|
require.NoError(t, err)
|
|
|
|
err = service.RecordOutcome(ctx, batchID, "d.txt", repository.BatchOutcomeStatusInvalidType, nil, nil)
|
|
require.NoError(t, err)
|
|
|
|
dupID := uuid.New()
|
|
// Create a real document for the duplicate reference
|
|
realDocID, createErr := cfg.GetDBQueries().CreateDocument(ctx, &repository.CreateDocumentParams{
|
|
Clientid: clientID,
|
|
Hash: "hash_dup",
|
|
BatchID: &batchID,
|
|
})
|
|
require.NoError(t, createErr)
|
|
_ = dupID // not needed since we use the real ID
|
|
err = service.RecordOutcome(ctx, batchID, "e.pdf", repository.BatchOutcomeStatusDuplicate, nil, &realDocID)
|
|
require.NoError(t, err)
|
|
|
|
outcomes, err := service.ListOutcomes(ctx, batchID)
|
|
require.NoError(t, err)
|
|
require.Len(t, outcomes, 5)
|
|
|
|
// Verify ordering is by created_at ASC
|
|
expectedFilenames := []string{"a.pdf", "b.pdf", "c.pdf", "d.txt", "e.pdf"}
|
|
for i, o := range outcomes {
|
|
assert.Equal(t, expectedFilenames[i], o.Filename)
|
|
}
|
|
|
|
// Verify the duplicate has document_id set
|
|
assert.NotNil(t, outcomes[4].DocumentID)
|
|
assert.Equal(t, realDocID, *outcomes[4].DocumentID)
|
|
}
|
|
|
|
// TestListOutcomes_WithCleanFail verifies that the clean_fail column from
|
|
// the currentCleanEntries view is joined when outcome is clean_failed.
|
|
func TestListOutcomes_WithCleanFail(t *testing.T) {
|
|
cfg := &TestConfig{}
|
|
test.CreateDB(t, cfg)
|
|
ctx := t.Context()
|
|
|
|
service := documentbatch.New(cfg)
|
|
|
|
clientID := "test_clean_fail_" + uuid.New().String()[:8]
|
|
err := cfg.GetDBQueries().CreateClient(ctx, &repository.CreateClientParams{
|
|
Clientid: clientID,
|
|
Name: "Test Client Clean Fail",
|
|
})
|
|
require.NoError(t, err)
|
|
|
|
batchID, err := service.Create(ctx, clientID, "cleanfail.zip", 1)
|
|
require.NoError(t, err)
|
|
|
|
// Create a real document
|
|
docID, err := cfg.GetDBQueries().CreateDocument(ctx, &repository.CreateDocumentParams{
|
|
Clientid: clientID,
|
|
Hash: "hash_clean_fail",
|
|
BatchID: &batchID,
|
|
})
|
|
require.NoError(t, err)
|
|
|
|
// Insert a clean_failed outcome
|
|
err = service.RecordOutcome(ctx, batchID, "corrupt.pdf", repository.BatchOutcomeStatusCleanFailed, nil, &docID)
|
|
require.NoError(t, err)
|
|
|
|
// Create a documentClean with fail type and a documentCleanEntry to populate
|
|
// the currentCleanEntries view
|
|
cleanID, err := cfg.GetDBQueries().AddDocumentClean(ctx, &repository.AddDocumentCleanParams{
|
|
Documentid: docID,
|
|
Fail: repository.NullCleanfailtype{
|
|
Cleanfailtype: repository.CleanfailtypeInvalidMimetype,
|
|
Valid: true,
|
|
},
|
|
})
|
|
require.NoError(t, err)
|
|
err = cfg.GetDBQueries().AddDocumentCleanEntry(ctx, &repository.AddDocumentCleanEntryParams{
|
|
Cleanid: cleanID,
|
|
Version: 1,
|
|
})
|
|
require.NoError(t, err)
|
|
|
|
outcomes, err := service.ListOutcomes(ctx, batchID)
|
|
require.NoError(t, err)
|
|
require.Len(t, outcomes, 1)
|
|
|
|
assert.Equal(t, "clean_failed", outcomes[0].Outcome)
|
|
require.NotNil(t, outcomes[0].CleanFail)
|
|
assert.Equal(t, "invalid_mimetype", *outcomes[0].CleanFail)
|
|
}
|
|
|
|
// TestCheckDuplicate creates a document, then verifies that CheckDuplicate
|
|
// returns its ID when the hash matches.
|
|
func TestCheckDuplicate(t *testing.T) {
|
|
cfg := &TestConfig{}
|
|
test.CreateDB(t, cfg)
|
|
ctx := t.Context()
|
|
|
|
service := documentbatch.New(cfg)
|
|
|
|
clientID := "test_check_dup_" + uuid.New().String()[:8]
|
|
err := cfg.GetDBQueries().CreateClient(ctx, &repository.CreateClientParams{
|
|
Clientid: clientID,
|
|
Name: "Test Client Check Dup",
|
|
})
|
|
require.NoError(t, err)
|
|
|
|
// Create a document with a known hash
|
|
knownHash := "sha256_known_hash_abcdef"
|
|
docID, err := cfg.GetDBQueries().CreateDocument(ctx, &repository.CreateDocumentParams{
|
|
Clientid: clientID,
|
|
Hash: knownHash,
|
|
})
|
|
require.NoError(t, err)
|
|
|
|
// Check duplicate with the same hash -- should find it
|
|
existingID, err := service.CheckDuplicate(ctx, clientID, knownHash)
|
|
require.NoError(t, err)
|
|
require.NotNil(t, existingID)
|
|
assert.Equal(t, docID, *existingID)
|
|
|
|
// Check duplicate with a different hash -- should return nil
|
|
noMatch, err := service.CheckDuplicate(ctx, clientID, "totally_different_hash")
|
|
require.NoError(t, err)
|
|
assert.Nil(t, noMatch)
|
|
}
|