query-orchestration/internal/document/clean/service_test.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")
		})
	}
}