package docinitrunner_test import ( "bytes" "regexp" "testing" "time" docinitrunner "queryorchestration/api/docInitRunner" "queryorchestration/internal/database/repository" "queryorchestration/internal/document/batch/outcome" documentinit "queryorchestration/internal/document/init" "queryorchestration/internal/server/runner" "queryorchestration/internal/serviceconfig/objectstore" "queryorchestration/internal/serviceconfig/queue" "queryorchestration/internal/serviceconfig/queue/documentsync" "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/jackc/pgx/v5/pgtype" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" ) type DocInitConfig struct { runner.BaseConfig[docinitrunner.Body] documentsync.DocSyncConfig queue.QueueConfig objectstore.ObjectStoreConfig } // stringPtr returns a pointer to a string func stringPtr(s string) *string { return &s } func TestDocInitRunner(t *testing.T) { cfg := &DocInitConfig{} test.CreateDB(t, cfg) acfg := test.CreateAWSContainer(t, cfg) // Set up SQS queue test.SetQueueClient(t, t.Context(), cfg, acfg.ExternalEndpoint) cfg.DocumentSyncURL = test.CreateQueue(t, cfg, test.DocSyncRunnerName) // Set up S3 bucket (required for file size measurement) test.SetStoreClient(t, t.Context(), cfg, acfg.ExternalEndpoint) test.CreateBucket(t, cfg) runnerInstance := docinitrunner.New(&docinitrunner.Services{ Document: documentinit.New(cfg), Queries: cfg.GetDBQueries(), Outcome: outcome.New(cfg.GetDBQueries()), }) err := cfg.GetDBQueries().CreateClient(t.Context(), &repository.CreateClientParams{ Clientid: "clientid", Name: "client_name", }) require.NoError(t, err) location := objectstore.BucketKey{ Location: objectstore.Import, ClientID: "clientid", CreatedAt: time.Now().UTC(), } // Upload a test file to S3 (required for HeadObject to work) testContent := []byte("test document content for file size measurement") _, err = cfg.GetStoreClient().PutObject(t.Context(), &s3.PutObjectInput{ Bucket: aws.String(cfg.GetBucket()), Key: aws.String(location.String()), Body: bytes.NewReader(testContent), }) require.NoError(t, err) // Create a document upload record to simulate the upload pipeline uploadID := uuid.New() err = cfg.GetDBQueries().AddDocumentUpload(t.Context(), &repository.AddDocumentUploadParams{ ID: uploadID, Clientid: "clientid", Bucket: cfg.GetBucket(), Key: location.String(), Part: 1, Createdat: pgtype.Timestamp{ Valid: true, Time: time.Now().UTC(), }, Filename: stringPtr("test-document.pdf"), BatchID: nil, }) require.NoError(t, err) doc := docinitrunner.Body{ Bucket: cfg.GetBucket(), Key: location.String(), Hash: "hash", } assert.True(t, runnerInstance.Process(t.Context(), doc)) test.AssertMessageBody(t, cfg, cfg.GetDocumentSyncURL(), regexp.MustCompile(`{"id":".+"}`)) } // setupDocInitTest creates shared infrastructure (DB, S3, SQS) and returns // the config, runner, and a cleanup-style helper. func setupDocInitTest(t *testing.T) *DocInitConfig { t.Helper() cfg := &DocInitConfig{} test.CreateDB(t, cfg) acfg := test.CreateAWSContainer(t, cfg) test.SetQueueClient(t, t.Context(), cfg, acfg.ExternalEndpoint) cfg.DocumentSyncURL = test.CreateQueue(t, cfg, test.DocSyncRunnerName) test.SetStoreClient(t, t.Context(), cfg, acfg.ExternalEndpoint) test.CreateBucket(t, cfg) return cfg } // uploadTestFileToS3 uploads a test file and returns the BucketKey. // If batchID is non-nil, it is embedded in the key so ParseBucketKey can extract it. func uploadTestFileToS3(t *testing.T, cfg *DocInitConfig, clientID string, batchID *uuid.UUID) objectstore.BucketKey { t.Helper() location := objectstore.BucketKey{ Location: objectstore.Import, ClientID: clientID, CreatedAt: time.Now().UTC(), EntityID: uuid.New(), BatchID: batchID, } testContent := []byte("test doc content " + uuid.New().String()) _, err := cfg.GetStoreClient().PutObject(t.Context(), &s3.PutObjectInput{ Bucket: aws.String(cfg.GetBucket()), Key: aws.String(location.String()), Body: bytes.NewReader(testContent), }) require.NoError(t, err) return location } // TestDocInitRunner_ReportsInitComplete processes a batch document and verifies // the outcome row is updated to init_complete with the assigned document_id. func TestDocInitRunner_ReportsInitComplete(t *testing.T) { cfg := setupDocInitTest(t) ctx := t.Context() clientID := "batch_init_complete_" + uuid.New().String()[:8] err := cfg.GetDBQueries().CreateClient(ctx, &repository.CreateClientParams{ Clientid: clientID, Name: "Batch Init Complete", }) require.NoError(t, err) // Create batch batchID, err := cfg.GetDBQueries().CreateBatchUpload(ctx, &repository.CreateBatchUploadParams{ ClientID: clientID, OriginalFilename: "test.zip", TotalDocuments: 1, }) require.NoError(t, err) filename := "invoice.pdf" // Insert submitted outcome err = cfg.GetDBQueries().InsertBatchDocumentOutcome(ctx, &repository.InsertBatchDocumentOutcomeParams{ BatchID: batchID, Filename: filename, Column3: repository.BatchOutcomeStatusSubmitted, }) require.NoError(t, err) // Upload test file and create upload record with batch ID embedded in key location := uploadTestFileToS3(t, cfg, clientID, &batchID) uploadID := uuid.New() err = cfg.GetDBQueries().AddDocumentUpload(ctx, &repository.AddDocumentUploadParams{ ID: uploadID, Clientid: clientID, Bucket: cfg.GetBucket(), Key: location.String(), Part: 1, Createdat: pgtype.Timestamp{ Valid: true, Time: time.Now().UTC(), }, Filename: &filename, BatchID: &batchID, }) require.NoError(t, err) runnerInstance := docinitrunner.New(&docinitrunner.Services{ Document: documentinit.New(cfg), Queries: cfg.GetDBQueries(), Outcome: outcome.New(cfg.GetDBQueries()), }) doc := docinitrunner.Body{ Bucket: cfg.GetBucket(), Key: location.String(), Hash: "unique_hash_init_complete_" + uuid.New().String()[:8], } assert.True(t, runnerInstance.Process(ctx, doc)) // Verify outcome was updated to init_complete with document_id set rows, err := cfg.GetDBQueries().ListBatchDocumentOutcomes(ctx, batchID) require.NoError(t, err) require.Len(t, rows, 1) assert.Equal(t, repository.BatchOutcomeStatusInitComplete, rows[0].Outcome) assert.NotNil(t, rows[0].DocumentID, "document_id should be set after init") } // TestDocInitRunner_ReportsInitDuplicate processes a batch document whose hash // matches an existing document and verifies the outcome is init_duplicate. func TestDocInitRunner_ReportsInitDuplicate(t *testing.T) { cfg := setupDocInitTest(t) ctx := t.Context() clientID := "batch_init_dup_" + uuid.New().String()[:8] err := cfg.GetDBQueries().CreateClient(ctx, &repository.CreateClientParams{ Clientid: clientID, Name: "Batch Init Duplicate", }) require.NoError(t, err) // Create an existing document with a known hash existingHash := "existing_hash_" + uuid.New().String()[:8] _, err = cfg.GetDBQueries().CreateDocument(ctx, &repository.CreateDocumentParams{ Clientid: clientID, Hash: existingHash, }) require.NoError(t, err) // Create batch batchID, err := cfg.GetDBQueries().CreateBatchUpload(ctx, &repository.CreateBatchUploadParams{ ClientID: clientID, OriginalFilename: "test.zip", TotalDocuments: 1, }) require.NoError(t, err) filename := "duplicate.pdf" // Insert submitted outcome err = cfg.GetDBQueries().InsertBatchDocumentOutcome(ctx, &repository.InsertBatchDocumentOutcomeParams{ BatchID: batchID, Filename: filename, Column3: repository.BatchOutcomeStatusSubmitted, }) require.NoError(t, err) // Upload test file and create upload record with batch ID in key location := uploadTestFileToS3(t, cfg, clientID, &batchID) uploadID := uuid.New() err = cfg.GetDBQueries().AddDocumentUpload(ctx, &repository.AddDocumentUploadParams{ ID: uploadID, Clientid: clientID, Bucket: cfg.GetBucket(), Key: location.String(), Part: 1, Createdat: pgtype.Timestamp{ Valid: true, Time: time.Now().UTC(), }, Filename: &filename, BatchID: &batchID, }) require.NoError(t, err) runnerInstance := docinitrunner.New(&docinitrunner.Services{ Document: documentinit.New(cfg), Queries: cfg.GetDBQueries(), Outcome: outcome.New(cfg.GetDBQueries()), }) // Use the same hash as the existing document to trigger duplicate doc := docinitrunner.Body{ Bucket: cfg.GetBucket(), Key: location.String(), Hash: existingHash, } assert.True(t, runnerInstance.Process(ctx, doc)) // Verify outcome was updated to init_duplicate rows, err := cfg.GetDBQueries().ListBatchDocumentOutcomes(ctx, batchID) require.NoError(t, err) require.Len(t, rows, 1) assert.Equal(t, repository.BatchOutcomeStatusInitDuplicate, rows[0].Outcome) assert.NotNil(t, rows[0].DocumentID, "document_id should be set for duplicate") } // TestDocInitRunner_NoOpForNonBatchDocument verifies that processing a non-batch // document does not cause outcome errors (the Resolve call is a no-op for nil batchID). func TestDocInitRunner_NoOpForNonBatchDocument(t *testing.T) { cfg := setupDocInitTest(t) ctx := t.Context() clientID := "nonbatch_init_" + uuid.New().String()[:8] err := cfg.GetDBQueries().CreateClient(ctx, &repository.CreateClientParams{ Clientid: clientID, Name: "NonBatch Init", }) require.NoError(t, err) location := uploadTestFileToS3(t, cfg, clientID, nil) uploadID := uuid.New() fname := "regular.pdf" err = cfg.GetDBQueries().AddDocumentUpload(ctx, &repository.AddDocumentUploadParams{ ID: uploadID, Clientid: clientID, Bucket: cfg.GetBucket(), Key: location.String(), Part: 1, Createdat: pgtype.Timestamp{ Valid: true, Time: time.Now().UTC(), }, Filename: &fname, BatchID: nil, // Non-batch }) require.NoError(t, err) runnerInstance := docinitrunner.New(&docinitrunner.Services{ Document: documentinit.New(cfg), Queries: cfg.GetDBQueries(), Outcome: outcome.New(cfg.GetDBQueries()), }) doc := docinitrunner.Body{ Bucket: cfg.GetBucket(), Key: location.String(), Hash: "nonbatch_hash_" + uuid.New().String()[:8], } // Should process successfully without any outcome errors assert.True(t, runnerInstance.Process(ctx, doc)) }