batch upload impl
This commit is contained in:
jay brown
@@ -1,19 +1,27 @@
|
||||
package queryapi_test
|
||||
|
||||
import (
|
||||
"archive/zip"
|
||||
"bytes"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"mime/multipart"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"path/filepath"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"queryorchestration/internal/database/repository"
|
||||
"queryorchestration/internal/test"
|
||||
|
||||
queryapi "queryorchestration/api/queryAPI"
|
||||
|
||||
"github.com/aws/aws-sdk-go-v2/service/s3"
|
||||
"github.com/google/uuid"
|
||||
"github.com/labstack/echo/v4"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
|
||||
queryapi "queryorchestration/api/queryAPI"
|
||||
"queryorchestration/internal/database/repository"
|
||||
"queryorchestration/internal/test"
|
||||
)
|
||||
|
||||
func TestUploadDocument(t *testing.T) {
|
||||
@@ -175,6 +183,7 @@ func TestUploadDocumentBatch(t *testing.T) {
|
||||
cfg := &ControllerConfig{}
|
||||
test.CreateDB(t, cfg)
|
||||
initializeTestConfig(t, cfg)
|
||||
test.CreateAWSResources(t, cfg)
|
||||
|
||||
svc := createControllerServices(cfg)
|
||||
cons := queryapi.NewControllers(svc, cfg)
|
||||
@@ -283,3 +292,227 @@ func TestCancelDocumentBatch(t *testing.T) {
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, repository.BatchStatusCancelled, batch.Status)
|
||||
}
|
||||
|
||||
func TestUploadDocumentBatch_Part1_Success(t *testing.T) {
|
||||
// Setup: Create testcontainers (database + localstack)
|
||||
cfg := &ControllerConfig{}
|
||||
test.CreateDB(t, cfg)
|
||||
initializeTestConfig(t, cfg)
|
||||
test.CreateAWSResources(t, cfg)
|
||||
|
||||
svc := createControllerServices(cfg)
|
||||
cons := queryapi.NewControllers(svc, cfg)
|
||||
|
||||
// Create test client with unique ID
|
||||
clientID := fmt.Sprintf("test_client_part1_%d", time.Now().UnixNano())
|
||||
err := cfg.GetDBQueries().CreateClient(t.Context(), &repository.CreateClientParams{
|
||||
Clientid: clientID,
|
||||
Name: "Test Client Part1",
|
||||
})
|
||||
require.NoError(t, err)
|
||||
|
||||
// Create test ZIP file with 3 PDFs (test1.pdf, test2.pdf, test3.pdf)
|
||||
zipContent := createTestZIPFile(t, 3)
|
||||
|
||||
// Create multipart form request
|
||||
body := &bytes.Buffer{}
|
||||
writer := multipart.NewWriter(body)
|
||||
|
||||
part, err := writer.CreateFormFile("archive", "test-batch.zip")
|
||||
require.NoError(t, err)
|
||||
|
||||
_, err = io.Copy(part, zipContent)
|
||||
require.NoError(t, err)
|
||||
|
||||
err = writer.Close()
|
||||
require.NoError(t, err)
|
||||
|
||||
// Test: Upload ZIP batch
|
||||
e := echo.New()
|
||||
rec := httptest.NewRecorder()
|
||||
ctx := e.NewContext(
|
||||
httptest.NewRequest(http.MethodPost, "/", body),
|
||||
rec,
|
||||
)
|
||||
ctx.Request().Header.Set("Content-Type", writer.FormDataContentType())
|
||||
|
||||
err = cons.UploadDocumentBatch(ctx, clientID)
|
||||
require.NoError(t, err)
|
||||
|
||||
// Verify: HTTP response
|
||||
assert.Equal(t, http.StatusAccepted, rec.Code)
|
||||
|
||||
var response queryapi.BatchUploadResponse
|
||||
err = json.Unmarshal(rec.Body.Bytes(), &response)
|
||||
require.NoError(t, err)
|
||||
|
||||
assert.NotEmpty(t, response.BatchId)
|
||||
assert.Equal(t, queryapi.BatchStatusProcessing, response.Status)
|
||||
assert.NotEmpty(t, response.StatusUrl)
|
||||
|
||||
// Verify: Batch exists in database with storage metadata
|
||||
batch, err := cfg.GetDBQueries().GetBatchUploadWithStorage(t.Context(), &repository.GetBatchUploadWithStorageParams{
|
||||
ID: uuid.UUID(response.BatchId),
|
||||
ClientID: clientID,
|
||||
})
|
||||
require.NoError(t, err)
|
||||
|
||||
assert.Equal(t, "test-batch.zip", batch.OriginalFilename)
|
||||
assert.Equal(t, repository.BatchStatusProcessing, batch.Status)
|
||||
assert.NotNil(t, batch.ArchiveBucket)
|
||||
assert.NotEmpty(t, *batch.ArchiveBucket)
|
||||
assert.NotNil(t, batch.ArchiveKey)
|
||||
assert.NotEmpty(t, *batch.ArchiveKey)
|
||||
assert.Contains(t, *batch.ArchiveKey, clientID)
|
||||
assert.Contains(t, *batch.ArchiveKey, "test-batch.zip")
|
||||
assert.NotNil(t, batch.FileSizeBytes)
|
||||
assert.Greater(t, *batch.FileSizeBytes, int64(0))
|
||||
|
||||
// Verify: ZIP file exists in S3
|
||||
s3Client := cfg.GetStoreClient()
|
||||
_, err = s3Client.HeadObject(t.Context(), &s3.HeadObjectInput{
|
||||
Bucket: batch.ArchiveBucket,
|
||||
Key: batch.ArchiveKey,
|
||||
})
|
||||
require.NoError(t, err, "ZIP file should exist in S3")
|
||||
}
|
||||
|
||||
func TestUploadDocumentBatch_Part1_InvalidFile(t *testing.T) {
|
||||
cfg := &ControllerConfig{}
|
||||
test.CreateDB(t, cfg)
|
||||
initializeTestConfig(t, cfg)
|
||||
|
||||
svc := createControllerServices(cfg)
|
||||
cons := queryapi.NewControllers(svc, cfg)
|
||||
|
||||
// Create test client with unique ID
|
||||
clientID := fmt.Sprintf("test_client_invalid_%d", time.Now().UnixNano())
|
||||
err := cfg.GetDBQueries().CreateClient(t.Context(), &repository.CreateClientParams{
|
||||
Clientid: clientID,
|
||||
Name: "Test Client Invalid",
|
||||
})
|
||||
require.NoError(t, err)
|
||||
|
||||
// Create invalid file (not ZIP)
|
||||
body := &bytes.Buffer{}
|
||||
writer := multipart.NewWriter(body)
|
||||
|
||||
part, err := writer.CreateFormFile("archive", "test.pdf")
|
||||
require.NoError(t, err)
|
||||
|
||||
_, err = part.Write([]byte("not a zip file"))
|
||||
require.NoError(t, err)
|
||||
|
||||
err = writer.Close()
|
||||
require.NoError(t, err)
|
||||
|
||||
// Test: Upload invalid file
|
||||
ctx := echo.New().NewContext(
|
||||
httptest.NewRequest(http.MethodPost, "/", body),
|
||||
httptest.NewRecorder(),
|
||||
)
|
||||
ctx.Request().Header.Set("Content-Type", writer.FormDataContentType())
|
||||
|
||||
err = cons.UploadDocumentBatch(ctx, clientID)
|
||||
|
||||
// Verify: Error response
|
||||
require.Error(t, err)
|
||||
httpErr, ok := err.(*echo.HTTPError)
|
||||
require.True(t, ok)
|
||||
assert.Equal(t, http.StatusBadRequest, httpErr.Code)
|
||||
assert.Contains(t, httpErr.Message, "invalid file type")
|
||||
}
|
||||
|
||||
func TestUploadDocumentBatch_Part1_TooLarge(t *testing.T) {
|
||||
cfg := &ControllerConfig{}
|
||||
test.CreateDB(t, cfg)
|
||||
initializeTestConfig(t, cfg)
|
||||
|
||||
svc := createControllerServices(cfg)
|
||||
cons := queryapi.NewControllers(svc, cfg)
|
||||
|
||||
// Create test client with unique ID
|
||||
clientID := fmt.Sprintf("test_client_toolarge_%d", time.Now().UnixNano())
|
||||
err := cfg.GetDBQueries().CreateClient(t.Context(), &repository.CreateClientParams{
|
||||
Clientid: clientID,
|
||||
Name: "Test Client TooLarge",
|
||||
})
|
||||
require.NoError(t, err)
|
||||
|
||||
// Create oversized ZIP file
|
||||
body := &bytes.Buffer{}
|
||||
writer := multipart.NewWriter(body)
|
||||
|
||||
// Set a fake Content-Length header to simulate large file
|
||||
part, err := writer.CreateFormFile("archive", "huge.zip")
|
||||
require.NoError(t, err)
|
||||
|
||||
// Write minimal content but set file size in header
|
||||
_, err = part.Write([]byte("PK")) // ZIP file header
|
||||
require.NoError(t, err)
|
||||
|
||||
err = writer.Close()
|
||||
require.NoError(t, err)
|
||||
|
||||
// Test: Upload with modified request to simulate large file
|
||||
req := httptest.NewRequest(http.MethodPost, "/", body)
|
||||
req.Header.Set("Content-Type", writer.FormDataContentType())
|
||||
|
||||
// Override the file size in the form data
|
||||
ctx := echo.New().NewContext(req, httptest.NewRecorder())
|
||||
|
||||
// Manually create form with large file size
|
||||
form, _ := ctx.MultipartForm()
|
||||
if len(form.File["archive"]) > 0 {
|
||||
// Simulate 101MB file
|
||||
form.File["archive"][0].Size = 101 * 1024 * 1024
|
||||
}
|
||||
|
||||
err = cons.UploadDocumentBatch(ctx, clientID)
|
||||
|
||||
// Verify: Error response
|
||||
require.Error(t, err)
|
||||
httpErr, ok := err.(*echo.HTTPError)
|
||||
require.True(t, ok)
|
||||
assert.Equal(t, http.StatusBadRequest, httpErr.Code)
|
||||
assert.Contains(t, httpErr.Message, "file too large")
|
||||
}
|
||||
|
||||
// Helper function to create test ZIP with multiple PDFs
|
||||
func createTestZIPFile(t *testing.T, pdfCount int) *bytes.Reader {
|
||||
var buf bytes.Buffer
|
||||
zipWriter := zip.NewWriter(&buf)
|
||||
|
||||
// Create standard test PDF files
|
||||
standardFiles := []string{"test1.pdf", "test2.pdf", "test3.pdf"}
|
||||
|
||||
// If requested count is <= 3, use standard names
|
||||
if pdfCount <= 3 {
|
||||
for i := 0; i < pdfCount; i++ {
|
||||
fileWriter, err := zipWriter.Create(standardFiles[i])
|
||||
require.NoError(t, err)
|
||||
|
||||
// Write minimal PDF content
|
||||
pdfContent := fmt.Sprintf("%%PDF-1.4\n%%Test content for %s\n%%%%EOF", standardFiles[i])
|
||||
_, err = fileWriter.Write([]byte(pdfContent))
|
||||
require.NoError(t, err)
|
||||
}
|
||||
} else {
|
||||
// For more than 3 files, use numbered names
|
||||
for i := 0; i < pdfCount; i++ {
|
||||
filename := fmt.Sprintf("document_%d.pdf", i)
|
||||
fileWriter, err := zipWriter.Create(filename)
|
||||
require.NoError(t, err)
|
||||
|
||||
// Write minimal PDF content
|
||||
pdfContent := fmt.Sprintf("%%PDF-1.4\n%%Document %d content\n%%%%EOF", i)
|
||||
_, err = fileWriter.Write([]byte(pdfContent))
|
||||
require.NoError(t, err)
|
||||
}
|
||||
}
|
||||
|
||||
err := zipWriter.Close()
|
||||
require.NoError(t, err)
|
||||
|
||||
return bytes.NewReader(buf.Bytes())
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user