Jay Brown
427 lines
14 KiB
Go
427 lines
14 KiB
Go
// 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")
|
|
})
|
|
}
|
|
}
|
|
|
|
// TestClean_AcceptsPDFWithTrailingNewline verifies that valid PDF files are
|
|
// accepted even when trailing whitespace (newlines) follows the %%EOF marker.
|
|
// Many PDF generators append a trailing newline; the cleaner must tolerate it.
|
|
func TestClean_AcceptsPDFWithTrailingNewline(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 := "pdf_trailing_nl_" + uniqueSuffix
|
|
err := cfg.GetDBQueries().CreateClient(ctx, &repository.CreateClientParams{
|
|
Clientid: clientID,
|
|
Name: "Test PDF Trailing Newline " + uniqueSuffix,
|
|
})
|
|
require.NoError(t, err)
|
|
|
|
// Structurally valid minimal PDF (cross-reference table, catalog, page).
|
|
// The base ends with %%EOF and no trailing newline; each sub-test appends its own suffix.
|
|
pdfBase := []byte("%PDF-1.4\n1 0 obj\n<</Type/Catalog/Pages 2 0 R>>\nendobj\n" +
|
|
"2 0 obj\n<</Type/Pages/Kids[3 0 R]/Count 1>>\nendobj\n" +
|
|
"3 0 obj\n<</Type/Page/Parent 2 0 R/MediaBox[0 0 612 792]>>\nendobj\n" +
|
|
"xref\n0 4\n0000000000 65535 f \n0000000009 00000 n \n0000000058 00000 n \n0000000115 00000 n \n" +
|
|
"trailer\n<</Size 4/Root 1 0 R>>\nstartxref\n190\n%%EOF")
|
|
|
|
tests := []struct {
|
|
name string
|
|
suffix string
|
|
}{
|
|
{name: "%%EOF\\n (Unix)", suffix: "\n"},
|
|
{name: "%%EOF\\r\\n (Windows)", suffix: "\r\n"},
|
|
{name: "%%EOF (no newline)", suffix: ""},
|
|
}
|
|
|
|
for _, tt := range tests {
|
|
t.Run(tt.name, func(t *testing.T) {
|
|
content := append([]byte{}, pdfBase...)
|
|
content = append(content, []byte(tt.suffix)...)
|
|
|
|
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(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.True(t, result.Passed, "valid PDF with trailing whitespace %q after %%%%EOF should be accepted", tt.suffix)
|
|
})
|
|
}
|
|
}
|
|
|
|
// TestClean_AcceptsTXTWithExtension verifies that a valid UTF-8 text file
|
|
// is accepted when the S3 key includes a .txt file extension.
|
|
func TestClean_AcceptsTXTWithExtension(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 := "txt_ext_" + uniqueSuffix
|
|
err := cfg.GetDBQueries().CreateClient(ctx, &repository.CreateClientParams{
|
|
Clientid: clientID,
|
|
Name: "Test TXT Extension " + uniqueSuffix,
|
|
})
|
|
require.NoError(t, err)
|
|
|
|
txtContent := []byte("Hello, this is a plain text file.\nIt has multiple lines.\nNo binary content here.")
|
|
|
|
fileType := "txt"
|
|
location := objectstore.BucketKey{
|
|
Location: objectstore.Import,
|
|
ClientID: clientID,
|
|
CreatedAt: time.Now().UTC(),
|
|
EntityID: uuid.New(),
|
|
FileType: &fileType,
|
|
}
|
|
|
|
_, err = cfg.GetStoreClient().PutObject(ctx, &s3.PutObjectInput{
|
|
Bucket: aws.String(cfg.GetBucket()),
|
|
Key: aws.String(location.String()),
|
|
Body: bytes.NewReader(txtContent),
|
|
ContentType: aws.String("text/plain"),
|
|
})
|
|
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.True(t, result.Passed, "valid TXT file with .txt extension should be accepted")
|
|
}
|
|
|
|
// TestClean_RejectsTXTWithoutExtension verifies that a text file without
|
|
// a .txt extension on the S3 key is rejected (regression guard for old uploads).
|
|
func TestClean_RejectsTXTWithoutExtension(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 := "txt_noext_" + uniqueSuffix
|
|
err := cfg.GetDBQueries().CreateClient(ctx, &repository.CreateClientParams{
|
|
Clientid: clientID,
|
|
Name: "Test TXT No Extension " + uniqueSuffix,
|
|
})
|
|
require.NoError(t, err)
|
|
|
|
txtContent := []byte("Hello, this is a plain text file.\nIt has multiple lines.\nNo binary content here.")
|
|
|
|
// No FileType set -- simulates old uploads without extension
|
|
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(txtContent),
|
|
ContentType: aws.String("text/plain"),
|
|
})
|
|
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, "TXT file without .txt extension should be rejected")
|
|
require.NotNil(t, result.FailReason)
|
|
assert.Equal(t, "invalid_mimetype", *result.FailReason)
|
|
}
|