Jay Brown
134 lines
4.2 KiB
Go
134 lines
4.2 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")
|
|
}
|