Jay Brown
238 lines
8.0 KiB
Go
238 lines
8.0 KiB
Go
// Package documentinit tests for document initialization.
|
|
// Note: Mock-based tests have been removed. See remove_texttract_and_mocks_plan.md.
|
|
package documentinit
|
|
|
|
import (
|
|
"bytes"
|
|
"fmt"
|
|
"testing"
|
|
"time"
|
|
|
|
"queryorchestration/internal/database/repository"
|
|
"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"
|
|
)
|
|
|
|
// TestMeasureFileSize_Integration tests the measureFileSize function using a real
|
|
// S3 client (localstack) to verify that HeadObject correctly returns file sizes.
|
|
// This is an integration test that requires Docker.
|
|
func TestMeasureFileSize_Integration(t *testing.T) {
|
|
ctx := t.Context()
|
|
|
|
// Set up AWS container (localstack) and S3 client
|
|
cfg := &DocInitConfig{}
|
|
test.CreateAWSResources(t, cfg)
|
|
|
|
svc := New(cfg)
|
|
|
|
// Upload a test file with known content
|
|
bucket := cfg.GetBucket()
|
|
testKey := fmt.Sprintf("test/filesize/%s.txt", uuid.New().String())
|
|
testContent := []byte("Hello, World! This is a test file for measuring file size.")
|
|
expectedSize := int64(len(testContent))
|
|
|
|
// Upload the test file to S3
|
|
_, err := cfg.GetStoreClient().PutObject(ctx, &s3.PutObjectInput{
|
|
Bucket: aws.String(bucket),
|
|
Key: aws.String(testKey),
|
|
Body: bytes.NewReader(testContent),
|
|
})
|
|
require.NoError(t, err, "failed to upload test file to S3")
|
|
|
|
// Measure the file size using our function
|
|
measuredSize, err := svc.measureFileSize(ctx, bucket, testKey)
|
|
require.NoError(t, err, "measureFileSize should not return an error")
|
|
require.NotNil(t, measuredSize, "measured size should not be nil")
|
|
|
|
// Verify the measured size matches expected
|
|
assert.Equal(t, expectedSize, *measuredSize, "measured file size should match uploaded content size")
|
|
}
|
|
|
|
// TestMeasureFileSize_Integration_LargeFile tests file size measurement for larger files.
|
|
func TestMeasureFileSize_Integration_LargeFile(t *testing.T) {
|
|
ctx := t.Context()
|
|
|
|
// Set up AWS container (localstack) and S3 client
|
|
cfg := &DocInitConfig{}
|
|
test.CreateAWSResources(t, cfg)
|
|
|
|
svc := New(cfg)
|
|
|
|
// Upload a larger test file (1MB)
|
|
bucket := cfg.GetBucket()
|
|
testKey := fmt.Sprintf("test/filesize/large/%s.bin", uuid.New().String())
|
|
expectedSize := int64(1024 * 1024) // 1MB
|
|
testContent := make([]byte, expectedSize)
|
|
for i := range testContent {
|
|
testContent[i] = byte(i % 256)
|
|
}
|
|
|
|
// Upload the test file to S3
|
|
_, err := cfg.GetStoreClient().PutObject(ctx, &s3.PutObjectInput{
|
|
Bucket: aws.String(bucket),
|
|
Key: aws.String(testKey),
|
|
Body: bytes.NewReader(testContent),
|
|
})
|
|
require.NoError(t, err, "failed to upload large test file to S3")
|
|
|
|
// Measure the file size using our function
|
|
measuredSize, err := svc.measureFileSize(ctx, bucket, testKey)
|
|
require.NoError(t, err, "measureFileSize should not return an error for large files")
|
|
require.NotNil(t, measuredSize, "measured size should not be nil")
|
|
|
|
// Verify the measured size matches expected
|
|
assert.Equal(t, expectedSize, *measuredSize, "measured file size should match 1MB")
|
|
}
|
|
|
|
// TestMeasureFileSize_Integration_NonExistentKey tests error handling for non-existent S3 keys.
|
|
func TestMeasureFileSize_Integration_NonExistentKey(t *testing.T) {
|
|
ctx := t.Context()
|
|
|
|
// Set up AWS container (localstack) and S3 client
|
|
cfg := &DocInitConfig{}
|
|
test.CreateAWSResources(t, cfg)
|
|
|
|
svc := New(cfg)
|
|
|
|
bucket := cfg.GetBucket()
|
|
nonExistentKey := fmt.Sprintf("test/nonexistent/%s.txt", uuid.New().String())
|
|
|
|
// Try to measure size of non-existent file
|
|
_, err := svc.measureFileSize(ctx, bucket, nonExistentKey)
|
|
assert.Error(t, err, "measureFileSize should return an error for non-existent key")
|
|
assert.Contains(t, err.Error(), "HeadObject failed", "error should indicate HeadObject failure")
|
|
}
|
|
|
|
// TestDocumentInitCreate_ReturnsCreateResult verifies that Create() returns a
|
|
// *CreateResult with correct IsDuplicate, BatchID, and Filename fields for both
|
|
// new and duplicate documents.
|
|
func TestDocumentInitCreate_ReturnsCreateResult(t *testing.T) {
|
|
ctx := t.Context()
|
|
|
|
// Set up DB, S3, and SQS
|
|
cfg := &DocInitConfig{}
|
|
test.CreateDB(t, cfg)
|
|
acfg := test.CreateAWSContainer(t, cfg)
|
|
|
|
test.SetQueueClient(t, ctx, cfg, acfg.ExternalEndpoint)
|
|
cfg.DocumentSyncURL = test.CreateQueue(t, cfg, test.DocSyncRunnerName)
|
|
|
|
test.SetStoreClient(t, ctx, cfg, acfg.ExternalEndpoint)
|
|
test.CreateBucket(t, cfg)
|
|
|
|
svc := New(cfg)
|
|
|
|
clientID := "create_result_" + uuid.New().String()[:8]
|
|
err := cfg.GetDBQueries().CreateClient(ctx, &repository.CreateClientParams{
|
|
Clientid: clientID,
|
|
Name: "Test CreateResult",
|
|
})
|
|
require.NoError(t, err)
|
|
|
|
// Create a batch
|
|
batchID, err := cfg.GetDBQueries().CreateBatchUpload(ctx, &repository.CreateBatchUploadParams{
|
|
ClientID: clientID,
|
|
OriginalFilename: "test.zip",
|
|
TotalDocuments: 2,
|
|
})
|
|
require.NoError(t, err)
|
|
|
|
// Upload a file to S3
|
|
location := objectstore.BucketKey{
|
|
Location: objectstore.Import,
|
|
ClientID: clientID,
|
|
CreatedAt: time.Now().UTC(),
|
|
EntityID: uuid.New(),
|
|
BatchID: &batchID,
|
|
}
|
|
testContent := []byte("test content for create result")
|
|
_, err = cfg.GetStoreClient().PutObject(ctx, &s3.PutObjectInput{
|
|
Bucket: aws.String(cfg.GetBucket()),
|
|
Key: aws.String(location.String()),
|
|
Body: bytes.NewReader(testContent),
|
|
})
|
|
require.NoError(t, err)
|
|
|
|
filename := "invoice.pdf"
|
|
uniqueHash := "create_result_hash_" + uuid.New().String()[:8]
|
|
|
|
// Test 1: New document returns IsDuplicate=false with correct fields
|
|
result, err := svc.Create(ctx, &Create{
|
|
Bucket: cfg.GetBucket(),
|
|
Key: location,
|
|
Hash: uniqueHash,
|
|
Filename: &filename,
|
|
})
|
|
require.NoError(t, err)
|
|
require.NotNil(t, result)
|
|
|
|
assert.NotEqual(t, uuid.Nil, result.ID, "new document should have a valid ID")
|
|
assert.False(t, result.IsDuplicate, "new document should not be a duplicate")
|
|
assert.NotNil(t, result.BatchID, "batch document should have BatchID set")
|
|
assert.Equal(t, batchID, *result.BatchID, "BatchID should match the batch")
|
|
assert.NotNil(t, result.Filename, "Filename should be set")
|
|
assert.Equal(t, filename, *result.Filename, "Filename should match")
|
|
|
|
// Test 2: Duplicate document (same hash) returns IsDuplicate=true
|
|
location2 := objectstore.BucketKey{
|
|
Location: objectstore.Import,
|
|
ClientID: clientID,
|
|
CreatedAt: time.Now().UTC(),
|
|
EntityID: uuid.New(),
|
|
BatchID: &batchID,
|
|
}
|
|
_, err = cfg.GetStoreClient().PutObject(ctx, &s3.PutObjectInput{
|
|
Bucket: aws.String(cfg.GetBucket()),
|
|
Key: aws.String(location2.String()),
|
|
Body: bytes.NewReader(testContent),
|
|
})
|
|
require.NoError(t, err)
|
|
|
|
dupResult, err := svc.Create(ctx, &Create{
|
|
Bucket: cfg.GetBucket(),
|
|
Key: location2,
|
|
Hash: uniqueHash, // Same hash triggers duplicate
|
|
Filename: &filename,
|
|
})
|
|
require.NoError(t, err)
|
|
require.NotNil(t, dupResult)
|
|
|
|
assert.Equal(t, result.ID, dupResult.ID, "duplicate should reference the original document ID")
|
|
assert.True(t, dupResult.IsDuplicate, "second document with same hash should be duplicate")
|
|
assert.NotNil(t, dupResult.BatchID, "duplicate should still have BatchID")
|
|
assert.Equal(t, batchID, *dupResult.BatchID, "duplicate BatchID should match")
|
|
|
|
// Test 3: Non-batch document has nil BatchID
|
|
location3 := objectstore.BucketKey{
|
|
Location: objectstore.Import,
|
|
ClientID: clientID,
|
|
CreatedAt: time.Now().UTC(),
|
|
EntityID: uuid.New(),
|
|
// No BatchID
|
|
}
|
|
_, err = cfg.GetStoreClient().PutObject(ctx, &s3.PutObjectInput{
|
|
Bucket: aws.String(cfg.GetBucket()),
|
|
Key: aws.String(location3.String()),
|
|
Body: bytes.NewReader(testContent),
|
|
})
|
|
require.NoError(t, err)
|
|
|
|
nonBatchResult, err := svc.Create(ctx, &Create{
|
|
Bucket: cfg.GetBucket(),
|
|
Key: location3,
|
|
Hash: "nonbatch_hash_" + uuid.New().String()[:8],
|
|
})
|
|
require.NoError(t, err)
|
|
require.NotNil(t, nonBatchResult)
|
|
|
|
assert.False(t, nonBatchResult.IsDuplicate, "non-batch new document should not be duplicate")
|
|
assert.Nil(t, nonBatchResult.BatchID, "non-batch document should have nil BatchID")
|
|
assert.Nil(t, nonBatchResult.Filename, "non-batch document should have nil Filename")
|
|
}
|