query-orchestration/internal/document/batch/storage/service_test.go

package storage

import (
	"archive/zip"
	"bytes"
	"fmt"
	"testing"

	"github.com/aws/aws-sdk-go-v2/service/s3"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"

	"queryorchestration/internal/serviceconfig"
	"queryorchestration/internal/serviceconfig/objectstore"
	"queryorchestration/internal/test"
)

// TestConfig provides a test configuration that implements both service config and AWS interfaces
type TestConfig struct {
	serviceconfig.BaseConfig
	objectstore.ObjectStoreConfig
}

func TestStoreZIPFile_Success(t *testing.T) {
	// Setup: Create testcontainer with localstack
	cfg := &TestConfig{}
	test.CreateDB(t, cfg)
	test.CreateAWSResources(t, cfg)

	// Create storage service
	service := New(cfg)

	// Create test ZIP content with 3 PDFs
	zipContent := createTestZIPContent(t)

	// Test: Store ZIP file
	result, err := service.StoreZIPFile(t.Context(), "test-client", "test.zip", zipContent)

	// Verify: Storage result
	require.NoError(t, err)
	assert.NotEmpty(t, result.Bucket)
	assert.NotEmpty(t, result.Key)
	assert.Greater(t, result.SizeBytes, int64(0))
	assert.Contains(t, result.Key, "test-client")
	assert.Contains(t, result.Key, "test.zip")

	// Verify: File exists in S3 (using actual S3 client)
	s3Client := cfg.GetStoreClient()
	_, err = s3Client.HeadObject(t.Context(), &s3.HeadObjectInput{
		Bucket: &result.Bucket,
		Key:    &result.Key,
	})
	require.NoError(t, err, "ZIP file should exist in S3")
}

func TestStoreZIPFile_LargeFile(t *testing.T) {
	// Setup: Create testcontainer with localstack
	cfg := &TestConfig{}
	test.CreateDB(t, cfg)
	test.CreateAWSResources(t, cfg)

	// Create storage service
	service := New(cfg)

	// Create a large test ZIP content (but still under 100MB limit)
	var buf bytes.Buffer
	zipWriter := zip.NewWriter(&buf)

	// Create multiple PDFs to make a larger ZIP
	for i := 0; i < 10; i++ {
		filename := fmt.Sprintf("document_%d.pdf", i)
		fileWriter, err := zipWriter.Create(filename)
		require.NoError(t, err)

		// Write a larger PDF content
		pdfContent := fmt.Sprintf("%%PDF-1.4\n")
		// Add some bulk content
		for j := 0; j < 1000; j++ {
			pdfContent += fmt.Sprintf("%%Page %d Line %d content\n", i, j)
		}
		pdfContent += "%%EOF"
		_, err = fileWriter.Write([]byte(pdfContent))
		require.NoError(t, err)
	}

	err := zipWriter.Close()
	require.NoError(t, err)

	// Test: Store larger ZIP file
	result, err := service.StoreZIPFile(t.Context(), "test-client", "large-batch.zip", bytes.NewReader(buf.Bytes()))

	// Verify: Storage result
	require.NoError(t, err)
	assert.NotEmpty(t, result.Bucket)
	assert.NotEmpty(t, result.Key)
	assert.Greater(t, result.SizeBytes, int64(0)) // Should have some size
}

func TestValidateZIPFile_Success(t *testing.T) {
	service := New(nil) // No config needed for validation

	// Test: Valid ZIP file
	err := service.ValidateZIPFile("test.zip", 1024*1024) // 1MB

	// Verify: No error
	assert.NoError(t, err)
}

func TestValidateZIPFile_InvalidExtension(t *testing.T) {
	service := New(nil)

	// Test: Invalid file extension
	err := service.ValidateZIPFile("test.pdf", 1024*1024)

	// Verify: Error returned
	require.Error(t, err)
	assert.Contains(t, err.Error(), "invalid file type")
}

func TestValidateZIPFile_TooLarge(t *testing.T) {
	service := New(nil)

	// Test: File too large (101MB)
	err := service.ValidateZIPFile("test.zip", 101*1024*1024)

	// Verify: Error returned
	require.Error(t, err)
	assert.Contains(t, err.Error(), "file too large")
}

func TestValidateZIPFile_Empty(t *testing.T) {
	service := New(nil)

	// Test: Empty file
	err := service.ValidateZIPFile("test.zip", 0)

	// Verify: Error returned
	require.Error(t, err)
	assert.Contains(t, err.Error(), "empty file not allowed")
}

// Helper function to create test ZIP content with 3 PDF files
func createTestZIPContent(t *testing.T) *bytes.Reader {
	var buf bytes.Buffer
	zipWriter := zip.NewWriter(&buf)

	// Create 3 PDF files in the ZIP
	pdfFiles := []string{"test1.pdf", "test2.pdf", "test3.pdf"}

	for _, filename := range pdfFiles {
		fileWriter, err := zipWriter.Create(filename)
		require.NoError(t, err)

		// Write minimal PDF content (PDF magic header)
		// Each file has slightly different content for identification
		pdfContent := fmt.Sprintf("%%PDF-1.4\n%%Test content for %s\n%%%%EOF", filename)
		_, err = fileWriter.Write([]byte(pdfContent))
		require.NoError(t, err)
	}

	err := zipWriter.Close()
	require.NoError(t, err)

	return bytes.NewReader(buf.Bytes())
}