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