// Package documentclean_test contains tests for the document clean service. // Note: Text extraction has been removed. The document pipeline ends after cleaning. package documentclean_test import ( "bytes" "testing" "time" "queryorchestration/internal/database/repository" documentclean "queryorchestration/internal/document/clean" "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" ) type DocCleanConfig struct { serviceconfig.BaseConfig objectstore.ObjectStoreConfig } func TestService(t *testing.T) { cfg := &DocCleanConfig{} svc := documentclean.New(cfg) assert.NotNil(t, svc) } // TestDocumentCleanClean_ReturnsCleanResult verifies that Clean() returns a // *CleanResult with correct Passed and FailReason fields. // // Two scenarios are tested: // 1. Already-cleaned document (HasDocumentCleanEntry=true) returns Passed=true // 2. Non-PDF file returns Passed=false with a non-nil FailReason func TestDocumentCleanClean_ReturnsCleanResult(t *testing.T) { cfg := &DocCleanConfig{} test.CreateDB(t, cfg) test.CreateAWSResources(t, cfg) ctx := t.Context() svc := documentclean.New(cfg) // Test 1: Already cleaned document returns Passed=true (shortcut path) clientID := "clean_result_" + uuid.New().String()[:8] err := cfg.GetDBQueries().CreateClient(ctx, &repository.CreateClientParams{ Clientid: clientID, Name: "Test Clean Result", }) require.NoError(t, err) docID, err := cfg.GetDBQueries().CreateDocument(ctx, &repository.CreateDocumentParams{ Clientid: clientID, Hash: "hash_clean_result_" + uuid.New().String()[:8], }) require.NoError(t, err) // Create a clean entry so HasDocumentCleanEntry returns true bucket := "bucket" key := "key" hash := "hash" cleanID, err := cfg.GetDBQueries().AddDocumentClean(ctx, &repository.AddDocumentCleanParams{ Documentid: docID, Bucket: &bucket, Key: &key, Hash: &hash, Mimetype: repository.NullCleanmimetype{ Cleanmimetype: repository.CleanmimetypeApplicationPdf, Valid: true, }, }) require.NoError(t, err) err = cfg.GetDBQueries().AddDocumentCleanEntry(ctx, &repository.AddDocumentCleanEntryParams{ Cleanid: cleanID, Version: 1, }) require.NoError(t, err) result, err := svc.Clean(ctx, docID) require.NoError(t, err) require.NotNil(t, result) assert.True(t, result.Passed, "already-cleaned document should return Passed=true") assert.Nil(t, result.FailReason, "Passed document should have nil FailReason") // Test 2: Non-PDF file returns Passed=false with FailReason // Upload a non-PDF file to S3 first to get its ETag fileContent := []byte("This is NOT a PDF file. It has no PDF signature.") location := objectstore.BucketKey{ Location: objectstore.Import, ClientID: clientID, CreatedAt: time.Now().UTC(), EntityID: uuid.New(), } _, err = cfg.GetStoreClient().PutObject(ctx, &s3.PutObjectInput{ Bucket: aws.String(cfg.GetBucket()), Key: aws.String(location.String()), Body: bytes.NewReader(fileContent), ContentType: aws.String("text/plain"), }) require.NoError(t, err) // Get ETag to use as the document hash (IfMatch requires it) headResp, err := cfg.GetStoreClient().HeadObject(ctx, &s3.HeadObjectInput{ Bucket: aws.String(cfg.GetBucket()), Key: aws.String(location.String()), }) require.NoError(t, err) etag := *headResp.ETag // Create document with hash matching the ETag docID2, err := cfg.GetDBQueries().CreateDocument(ctx, &repository.CreateDocumentParams{ Clientid: clientID, Hash: etag, }) require.NoError(t, err) err = cfg.GetDBQueries().AddDocumentEntry(ctx, &repository.AddDocumentEntryParams{ Documentid: docID2, Bucket: cfg.GetBucket(), Key: location.String(), }) require.NoError(t, err) failResult, err := svc.Clean(ctx, docID2) require.NoError(t, err) require.NotNil(t, failResult) assert.False(t, failResult.Passed, "non-PDF document should fail validation") assert.NotNil(t, failResult.FailReason, "failed document should have a FailReason") } // TestClean_RejectsNonPDFWithPDFContentType verifies that files with // application/pdf content type metadata but non-PDF content are rejected. // This is the core scenario where S3 metadata lies about the file type. func TestClean_RejectsNonPDFWithPDFContentType(t *testing.T) { cfg := &DocCleanConfig{} test.CreateDB(t, cfg) test.CreateAWSResources(t, cfg) ctx := t.Context() svc := documentclean.New(cfg) uniqueSuffix := uuid.New().String()[:8] clientID := "mimetype_check_" + uniqueSuffix err := cfg.GetDBQueries().CreateClient(ctx, &repository.CreateClientParams{ Clientid: clientID, Name: "Test MIME Check " + uniqueSuffix, }) require.NoError(t, err) tests := []struct { name string content []byte }{ { name: "xlsx with PDF content type", content: []byte("PK\x03\x04fake xlsx content that is definitely not a PDF file at all"), }, { name: "OLE2/msg with PDF content type", content: []byte("\xD0\xCF\x11\xE0\xA1\xB1\x1A\xE1fake OLE2 msg content not a PDF"), }, } for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { location := objectstore.BucketKey{ Location: objectstore.Import, ClientID: clientID, CreatedAt: time.Now().UTC(), EntityID: uuid.New(), } // Upload with application/pdf content type despite non-PDF content _, err := cfg.GetStoreClient().PutObject(ctx, &s3.PutObjectInput{ Bucket: aws.String(cfg.GetBucket()), Key: aws.String(location.String()), Body: bytes.NewReader(tt.content), ContentType: aws.String("application/pdf"), }) require.NoError(t, err) headResp, err := cfg.GetStoreClient().HeadObject(ctx, &s3.HeadObjectInput{ Bucket: aws.String(cfg.GetBucket()), Key: aws.String(location.String()), }) require.NoError(t, err) etag := *headResp.ETag docID, err := cfg.GetDBQueries().CreateDocument(ctx, &repository.CreateDocumentParams{ Clientid: clientID, Hash: etag, }) require.NoError(t, err) err = cfg.GetDBQueries().AddDocumentEntry(ctx, &repository.AddDocumentEntryParams{ Documentid: docID, Bucket: cfg.GetBucket(), Key: location.String(), }) require.NoError(t, err) result, err := svc.Clean(ctx, docID) require.NoError(t, err) require.NotNil(t, result) assert.False(t, result.Passed, "file with non-PDF content should be rejected even if content type says application/pdf") require.NotNil(t, result.FailReason, "rejected document should have a FailReason") assert.Equal(t, "invalid_mimetype", *result.FailReason, "non-PDF content should be rejected at MIME type check, not at PDF parse stage") }) } }