package queryapi_test import ( "archive/zip" "bytes" "encoding/json" "fmt" "io" "mime/multipart" "net/http" "net/http/httptest" "path/filepath" "testing" "time" "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) { cfg := &ControllerConfig{} test.CreateDB(t, cfg) initializeTestConfig(t, cfg) acfg := test.CreateAWSContainer(t, cfg) test.SetStoreClient(t, t.Context(), cfg, acfg.ExternalEndpoint) test.CreateBucket(t, cfg) svc := createControllerServices(cfg) cons := queryapi.NewControllers(svc, cfg) err := cfg.GetDBQueries().CreateClient(t.Context(), &repository.CreateClientParams{ Clientid: "client_id", Name: "client_name", }) require.NoError(t, err) var buf bytes.Buffer writer := multipart.NewWriter(&buf) fileWriter, err := writer.CreateFormFile("file", filepath.Base("../../assets/original/helloWorld.pdf")) require.NoError(t, err) _, err = fileWriter.Write([]byte("hi")) require.NoError(t, err) err = writer.Close() require.NoError(t, err) ctx, rec := createContextWithBody(t, buf.Bytes()) ctx.Request().Header.Set("Content-Type", writer.FormDataContentType()) err = cons.UploadDocument(ctx, "client_id") require.NoError(t, err) assert.Equal(t, http.StatusOK, rec.Code) } func TestListDocumentsByClientId(t *testing.T) { cfg := &ControllerConfig{} test.CreateDB(t, cfg) initializeTestConfig(t, cfg) svc := createControllerServices(cfg) cons := queryapi.NewControllers(svc, cfg) err := cfg.GetDBQueries().CreateClient(t.Context(), &repository.CreateClientParams{ Clientid: "client_id", Name: "client_name", }) require.NoError(t, err) docId, err := cfg.GetDBQueries().CreateDocument(t.Context(), &repository.CreateDocumentParams{ Clientid: "client_id", Hash: "hash", BatchID: nil, }) require.NoError(t, err) ctx, rec := createContext(t) err = cons.ListDocumentsByClientId(ctx, "client_id") require.NoError(t, err) assert.Equal(t, http.StatusOK, rec.Code) assertBody(t, rec, queryapi.ListDocuments{ queryapi.DocumentSummary{ Id: docId, Hash: "hash", }, }) } func TestGetDocument(t *testing.T) { cfg := &ControllerConfig{} test.CreateDB(t, cfg) initializeTestConfig(t, cfg) svc := createControllerServices(cfg) cons := queryapi.NewControllers(svc, cfg) err := cfg.GetDBQueries().CreateClient(t.Context(), &repository.CreateClientParams{ Clientid: "client_id", Name: "client_name", }) require.NoError(t, err) docId, err := cfg.GetDBQueries().CreateDocument(t.Context(), &repository.CreateDocumentParams{ Clientid: "client_id", Hash: "hash", BatchID: nil, }) require.NoError(t, err) ctx, rec := createContext(t) err = cons.GetDocument(ctx, docId) require.NoError(t, err) assert.Equal(t, http.StatusOK, rec.Code) assertBody(t, rec, queryapi.Document{ Id: docId, ClientId: "client_id", Hash: "hash", Fields: map[string]interface{}{}, }) } func TestListDocumentBatches(t *testing.T) { cfg := &ControllerConfig{} test.CreateDB(t, cfg) initializeTestConfig(t, cfg) svc := createControllerServices(cfg) cons := queryapi.NewControllers(svc, cfg) // Create test client clientID := "test_client_list" err := cfg.GetDBQueries().CreateClient(t.Context(), &repository.CreateClientParams{ Clientid: clientID, Name: "Test Client Batch List", }) require.NoError(t, err) // Create test batch batchID, err := cfg.GetDBQueries().CreateBatchUpload(t.Context(), &repository.CreateBatchUploadParams{ ClientID: clientID, OriginalFilename: "test.zip", TotalDocuments: 10, }) require.NoError(t, err) // Test listing batches with default parameters ctx, rec := createContext(t) err = cons.ListDocumentBatches(ctx, clientID, queryapi.ListDocumentBatchesParams{}) require.NoError(t, err) assert.Equal(t, http.StatusOK, rec.Code) response := getBody[queryapi.BatchUploadList](t, rec) assert.Len(t, response.Batches, 1) assert.Equal(t, batchID.String(), response.Batches[0].BatchId.String()) // Test listing batches with custom limit and offset limit := int32(5) offset := int32(0) ctx2, rec2 := createContext(t) err = cons.ListDocumentBatches(ctx2, clientID, queryapi.ListDocumentBatchesParams{ Limit: &limit, Offset: &offset, }) require.NoError(t, err) assert.Equal(t, http.StatusOK, rec2.Code) response2 := getBody[queryapi.BatchUploadList](t, rec2) assert.Len(t, response2.Batches, 1) assert.Equal(t, batchID.String(), response2.Batches[0].BatchId.String()) } 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) // Create test client clientID := "test_client_upload" err := cfg.GetDBQueries().CreateClient(t.Context(), &repository.CreateClientParams{ Clientid: clientID, Name: "Test Client Batch Upload", }) require.NoError(t, err) // Create multipart form with ZIP file var buf bytes.Buffer writer := multipart.NewWriter(&buf) fileWriter, err := writer.CreateFormFile("archive", "test.zip") require.NoError(t, err) _, err = fileWriter.Write([]byte("PK")) // Minimal ZIP file header require.NoError(t, err) err = writer.Close() require.NoError(t, err) ctx, rec := createContextWithBody(t, buf.Bytes()) ctx.Request().Header.Set("Content-Type", writer.FormDataContentType()) // Test batch upload err = cons.UploadDocumentBatch(ctx, clientID) require.NoError(t, err) assert.Equal(t, http.StatusAccepted, rec.Code) response := getBody[queryapi.BatchUploadResponse](t, rec) assert.NotEmpty(t, response.BatchId) assert.Equal(t, queryapi.BatchStatusProcessing, response.Status) } func TestGetDocumentBatch(t *testing.T) { cfg := &ControllerConfig{} test.CreateDB(t, cfg) initializeTestConfig(t, cfg) svc := createControllerServices(cfg) cons := queryapi.NewControllers(svc, cfg) // Create test client clientID := "test_client_get" err := cfg.GetDBQueries().CreateClient(t.Context(), &repository.CreateClientParams{ Clientid: clientID, Name: "Test Client Batch Get", }) require.NoError(t, err) // Create test batch batchID, err := cfg.GetDBQueries().CreateBatchUpload(t.Context(), &repository.CreateBatchUploadParams{ ClientID: clientID, OriginalFilename: "test.zip", TotalDocuments: 10, }) require.NoError(t, err) // Test getting batch details ctx, rec := createContext(t) err = cons.GetDocumentBatch(ctx, clientID, queryapi.BatchID(batchID)) require.NoError(t, err) assert.Equal(t, http.StatusOK, rec.Code) response := getBody[queryapi.BatchUploadDetails](t, rec) assert.Equal(t, batchID.String(), response.BatchId.String()) assert.Equal(t, "test.zip", response.OriginalFilename) } func TestCancelDocumentBatch(t *testing.T) { cfg := &ControllerConfig{} test.CreateDB(t, cfg) initializeTestConfig(t, cfg) svc := createControllerServices(cfg) cons := queryapi.NewControllers(svc, cfg) // Create test client clientID := "test_client_cancel" err := cfg.GetDBQueries().CreateClient(t.Context(), &repository.CreateClientParams{ Clientid: clientID, Name: "Test Client Batch Cancel", }) require.NoError(t, err) // Create test batch batchID, err := cfg.GetDBQueries().CreateBatchUpload(t.Context(), &repository.CreateBatchUploadParams{ ClientID: clientID, OriginalFilename: "test.zip", TotalDocuments: 10, }) require.NoError(t, err) // Test canceling batch ctx, rec := createContext(t) err = cons.CancelDocumentBatch(ctx, clientID, queryapi.BatchID(batchID)) require.NoError(t, err) assert.Equal(t, http.StatusNoContent, rec.Code) // Verify batch was canceled batch, err := cfg.GetDBQueries().GetBatchUpload(t.Context(), &repository.GetBatchUploadParams{ ID: batchID, ClientID: clientID, }) 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()) }