package docinitrunner_test import ( "bytes" "regexp" "testing" "time" docinitrunner "queryorchestration/api/docInitRunner" "queryorchestration/internal/database/repository" 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(), }) 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":".+"}`)) }