Jay Brown
917 lines
28 KiB
Go
917 lines
28 KiB
Go
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/fieldextraction"
|
|
"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)
|
|
|
|
// Use test helper to create client with root folder
|
|
test.CreateTestClient(t, cfg, "client_id", "client_name")
|
|
|
|
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)
|
|
|
|
// Use test helper to create client with root folder
|
|
test.CreateTestClient(t, cfg, "client_id", "client_name")
|
|
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)
|
|
|
|
// Use test helper to create client with root folder
|
|
test.CreateTestClient(t, cfg, "client_id", "client_name")
|
|
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, queryapi.GetDocumentParams{})
|
|
require.NoError(t, err)
|
|
assert.Equal(t, http.StatusOK, rec.Code)
|
|
assertBody(t, rec, queryapi.DocumentEnriched{
|
|
Id: docId,
|
|
ClientId: "client_id",
|
|
Hash: "hash",
|
|
HasTextRecord: false,
|
|
Labels: []queryapi.LabelRecord{},
|
|
})
|
|
}
|
|
|
|
func TestGetDocumentWithEnrichedFields(t *testing.T) {
|
|
cfg := &ControllerConfig{}
|
|
test.CreateDB(t, cfg)
|
|
initializeTestConfig(t, cfg)
|
|
|
|
svc := createControllerServices(cfg)
|
|
cons := queryapi.NewControllers(svc, cfg)
|
|
|
|
// Create test client with root folder
|
|
test.CreateTestClient(t, cfg, "enriched_client", "Enriched Client")
|
|
|
|
// Create folder
|
|
folder, err := cfg.GetDBQueries().CreateFolder(t.Context(), &repository.CreateFolderParams{
|
|
Path: "/documents",
|
|
Clientid: "enriched_client",
|
|
Createdby: "testuser",
|
|
})
|
|
require.NoError(t, err)
|
|
|
|
// Create document with folder, filename, originalPath
|
|
filename := "test-document.pdf"
|
|
originalPath := "/uploads/test-document.pdf"
|
|
docId, err := cfg.GetDBQueries().CreateDocument(t.Context(), &repository.CreateDocumentParams{
|
|
Clientid: "enriched_client",
|
|
Hash: "enriched_hash",
|
|
Filename: &filename,
|
|
Originalpath: &originalPath,
|
|
Folderid: &folder.ID,
|
|
})
|
|
require.NoError(t, err)
|
|
|
|
// Add labels to document - use a label that exists in the lookup table
|
|
_, err = cfg.GetDBQueries().ApplyLabel(t.Context(), &repository.ApplyLabelParams{
|
|
Documentid: docId,
|
|
Label: "Ingested", // This label is seeded in migration
|
|
Appliedby: "testuser@example.com",
|
|
})
|
|
require.NoError(t, err)
|
|
|
|
ctx, rec := createContext(t)
|
|
|
|
// Test without textRecord parameter
|
|
err = cons.GetDocument(ctx, docId, queryapi.GetDocumentParams{})
|
|
require.NoError(t, err)
|
|
assert.Equal(t, http.StatusOK, rec.Code)
|
|
|
|
var result queryapi.DocumentEnriched
|
|
err = json.Unmarshal(rec.Body.Bytes(), &result)
|
|
require.NoError(t, err)
|
|
|
|
// Verify enriched fields
|
|
assert.Equal(t, docId, result.Id)
|
|
assert.Equal(t, queryapi.ClientID("enriched_client"), result.ClientId)
|
|
assert.Equal(t, queryapi.Hash("enriched_hash"), result.Hash)
|
|
assert.NotNil(t, result.Filename)
|
|
assert.Equal(t, filename, *result.Filename)
|
|
assert.NotNil(t, result.OriginalPath)
|
|
assert.Equal(t, originalPath, *result.OriginalPath)
|
|
assert.NotNil(t, result.FolderId)
|
|
assert.Equal(t, folder.ID.String(), result.FolderId.String())
|
|
assert.False(t, result.HasTextRecord)
|
|
assert.Nil(t, result.TextRecord)
|
|
|
|
// Verify labels
|
|
assert.Len(t, result.Labels, 1)
|
|
assert.Equal(t, "Ingested", result.Labels[0].Label)
|
|
assert.Equal(t, "testuser@example.com", string(result.Labels[0].AppliedBy))
|
|
}
|
|
|
|
func TestGetDocumentWithTextRecord(t *testing.T) {
|
|
cfg := &ControllerConfig{}
|
|
test.CreateDB(t, cfg)
|
|
initializeTestConfig(t, cfg)
|
|
|
|
svc := createControllerServices(cfg)
|
|
cons := queryapi.NewControllers(svc, cfg)
|
|
|
|
// Create test client with root folder
|
|
test.CreateTestClient(t, cfg, "textrecord_client", "TextRecord Client")
|
|
|
|
// Create document
|
|
filename := "contract.pdf"
|
|
docId, err := cfg.GetDBQueries().CreateDocument(t.Context(), &repository.CreateDocumentParams{
|
|
Clientid: "textrecord_client",
|
|
Hash: "textrecord_hash",
|
|
Filename: &filename,
|
|
})
|
|
require.NoError(t, err)
|
|
|
|
// Create field extraction using the fieldextraction service (which properly inserts into versions table)
|
|
fieldSvc := fieldextraction.New(cfg)
|
|
contractTitle := "Test Contract Title"
|
|
extractionUser := "extraction@example.com"
|
|
input := &fieldextraction.CreateFieldExtractionInput{
|
|
DocumentID: docId,
|
|
SingleFields: &repository.AddFieldExtractionParams{
|
|
Documentid: docId,
|
|
Filename: &filename,
|
|
Contracttitle: &contractTitle,
|
|
Createdby: extractionUser,
|
|
},
|
|
ArrayFields: []*repository.AddFieldExtractionArrayFieldParams{},
|
|
}
|
|
_, err = fieldSvc.CreateFieldExtraction(t.Context(), input)
|
|
require.NoError(t, err)
|
|
|
|
// Test with textRecord=true parameter
|
|
includeTextRecord := true
|
|
ctx, rec := createContext(t)
|
|
|
|
err = cons.GetDocument(ctx, docId, queryapi.GetDocumentParams{
|
|
TextRecord: &includeTextRecord,
|
|
})
|
|
require.NoError(t, err)
|
|
assert.Equal(t, http.StatusOK, rec.Code)
|
|
|
|
var result queryapi.DocumentEnriched
|
|
err = json.Unmarshal(rec.Body.Bytes(), &result)
|
|
require.NoError(t, err)
|
|
|
|
// Verify document has text record
|
|
assert.Equal(t, docId, result.Id)
|
|
assert.True(t, result.HasTextRecord)
|
|
assert.NotNil(t, result.TextRecord)
|
|
assert.Equal(t, int32(1), result.TextRecord.Version)
|
|
assert.Equal(t, extractionUser, string(result.TextRecord.CreatedBy))
|
|
assert.NotNil(t, result.TextRecord.SingleFields.ContractTitle)
|
|
assert.Equal(t, contractTitle, *result.TextRecord.SingleFields.ContractTitle)
|
|
}
|
|
|
|
func TestGetDocumentWithTextRecordNotRequested(t *testing.T) {
|
|
cfg := &ControllerConfig{}
|
|
test.CreateDB(t, cfg)
|
|
initializeTestConfig(t, cfg)
|
|
|
|
svc := createControllerServices(cfg)
|
|
cons := queryapi.NewControllers(svc, cfg)
|
|
|
|
// Create test client with root folder
|
|
test.CreateTestClient(t, cfg, "notext_client", "NoText Client")
|
|
|
|
// Create document
|
|
filename := "contract2.pdf"
|
|
docId, err := cfg.GetDBQueries().CreateDocument(t.Context(), &repository.CreateDocumentParams{
|
|
Clientid: "notext_client",
|
|
Hash: "notext_hash",
|
|
Filename: &filename,
|
|
})
|
|
require.NoError(t, err)
|
|
|
|
// Create field extraction using the fieldextraction service
|
|
fieldSvc := fieldextraction.New(cfg)
|
|
contractTitle := "Another Contract"
|
|
input := &fieldextraction.CreateFieldExtractionInput{
|
|
DocumentID: docId,
|
|
SingleFields: &repository.AddFieldExtractionParams{
|
|
Documentid: docId,
|
|
Filename: &filename,
|
|
Contracttitle: &contractTitle,
|
|
Createdby: "extraction2@example.com",
|
|
},
|
|
ArrayFields: []*repository.AddFieldExtractionArrayFieldParams{},
|
|
}
|
|
_, err = fieldSvc.CreateFieldExtraction(t.Context(), input)
|
|
require.NoError(t, err)
|
|
|
|
// Test without textRecord parameter (should still have hasTextRecord=true but no textRecord body)
|
|
ctx, rec := createContext(t)
|
|
|
|
err = cons.GetDocument(ctx, docId, queryapi.GetDocumentParams{})
|
|
require.NoError(t, err)
|
|
assert.Equal(t, http.StatusOK, rec.Code)
|
|
|
|
var result queryapi.DocumentEnriched
|
|
err = json.Unmarshal(rec.Body.Bytes(), &result)
|
|
require.NoError(t, err)
|
|
|
|
// Verify document indicates it has text record but doesn't include the full body
|
|
assert.Equal(t, docId, result.Id)
|
|
assert.True(t, result.HasTextRecord)
|
|
assert.Nil(t, result.TextRecord) // Not requested, so should be nil
|
|
}
|
|
|
|
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 with root folder
|
|
clientID := "test_client_list"
|
|
test.CreateTestClient(t, cfg, clientID, "Test Client Batch List")
|
|
|
|
// 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 with root folder
|
|
clientID := "test_client_upload"
|
|
test.CreateTestClient(t, cfg, clientID, "Test Client Batch Upload")
|
|
|
|
// 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 with root folder
|
|
clientID := "test_client_get"
|
|
test.CreateTestClient(t, cfg, clientID, "Test Client Batch Get")
|
|
|
|
// 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)
|
|
}
|
|
|
|
// TestGetDocumentBatch_IncludesOutcomes verifies that GetDocumentBatch returns
|
|
// the document_outcomes array populated with per-file outcome rows.
|
|
func TestGetDocumentBatch_IncludesOutcomes(t *testing.T) {
|
|
cfg := &ControllerConfig{}
|
|
test.CreateDB(t, cfg)
|
|
initializeTestConfig(t, cfg)
|
|
|
|
svc := createControllerServices(cfg)
|
|
cons := queryapi.NewControllers(svc, cfg)
|
|
|
|
clientID := fmt.Sprintf("outcome_client_%s", uuid.New().String()[:8])
|
|
test.CreateTestClient(t, cfg, clientID, "Test Outcomes Client")
|
|
|
|
// Create batch
|
|
batchID, err := cfg.GetDBQueries().CreateBatchUpload(t.Context(), &repository.CreateBatchUploadParams{
|
|
ClientID: clientID,
|
|
OriginalFilename: "outcomes.zip",
|
|
TotalDocuments: 3,
|
|
})
|
|
require.NoError(t, err)
|
|
|
|
// Create a document so we can reference it in outcome rows
|
|
docID, err := cfg.GetDBQueries().CreateDocument(t.Context(), &repository.CreateDocumentParams{
|
|
Clientid: clientID,
|
|
Hash: "outcome_hash_" + uuid.New().String()[:8],
|
|
BatchID: &batchID,
|
|
})
|
|
require.NoError(t, err)
|
|
|
|
// Insert outcomes with different statuses
|
|
err = cfg.GetDBQueries().InsertBatchDocumentOutcome(t.Context(), &repository.InsertBatchDocumentOutcomeParams{
|
|
BatchID: batchID,
|
|
Filename: "invoice.pdf",
|
|
Column3: repository.BatchOutcomeStatusCleanPassed,
|
|
DocumentID: &docID,
|
|
})
|
|
require.NoError(t, err)
|
|
|
|
invalidErr := "not a PDF"
|
|
err = cfg.GetDBQueries().InsertBatchDocumentOutcome(t.Context(), &repository.InsertBatchDocumentOutcomeParams{
|
|
BatchID: batchID,
|
|
Filename: "spreadsheet.xlsx",
|
|
Column3: repository.BatchOutcomeStatusInvalidType,
|
|
ErrorDetail: &invalidErr,
|
|
})
|
|
require.NoError(t, err)
|
|
|
|
err = cfg.GetDBQueries().InsertBatchDocumentOutcome(t.Context(), &repository.InsertBatchDocumentOutcomeParams{
|
|
BatchID: batchID,
|
|
Filename: "receipt.pdf",
|
|
Column3: repository.BatchOutcomeStatusSubmitted,
|
|
DocumentID: &docID,
|
|
})
|
|
require.NoError(t, err)
|
|
|
|
// Call GetDocumentBatch
|
|
ctx, rec := createContext(t)
|
|
err = cons.GetDocumentBatch(ctx, clientID, queryapi.BatchID(batchID))
|
|
require.NoError(t, err)
|
|
assert.Equal(t, http.StatusOK, rec.Code)
|
|
|
|
var response queryapi.BatchUploadDetails
|
|
err = json.Unmarshal(rec.Body.Bytes(), &response)
|
|
require.NoError(t, err)
|
|
|
|
// Verify document_outcomes array is populated
|
|
require.NotNil(t, response.DocumentOutcomes, "document_outcomes should be present")
|
|
require.Len(t, *response.DocumentOutcomes, 3, "should have 3 outcome entries")
|
|
|
|
// Build a map by filename for easier assertion
|
|
outcomeMap := make(map[string]queryapi.BatchDocumentOutcome)
|
|
for _, o := range *response.DocumentOutcomes {
|
|
outcomeMap[o.Filename] = o
|
|
}
|
|
|
|
// Verify each outcome
|
|
invoiceOutcome := outcomeMap["invoice.pdf"]
|
|
assert.Equal(t, queryapi.BatchDocumentOutcomeOutcome("clean_passed"), invoiceOutcome.Outcome)
|
|
assert.True(t, invoiceOutcome.DocumentId.IsSpecified(), "clean_passed outcome should have document_id")
|
|
|
|
xlsxOutcome := outcomeMap["spreadsheet.xlsx"]
|
|
assert.Equal(t, queryapi.BatchDocumentOutcomeOutcome("invalid_type"), xlsxOutcome.Outcome)
|
|
assert.True(t, xlsxOutcome.ErrorDetail.IsSpecified(), "invalid_type outcome should have error_detail")
|
|
|
|
receiptOutcome := outcomeMap["receipt.pdf"]
|
|
assert.Equal(t, queryapi.BatchDocumentOutcomeOutcome("submitted"), receiptOutcome.Outcome)
|
|
}
|
|
|
|
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 and root folder
|
|
clientID := fmt.Sprintf("test_client_part1_%d", time.Now().UnixNano())
|
|
test.CreateTestClient(t, cfg, clientID, "Test Client Part1")
|
|
|
|
// 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 and root folder
|
|
clientID := fmt.Sprintf("test_client_invalid_%d", time.Now().UnixNano())
|
|
test.CreateTestClient(t, cfg, clientID, "Test Client Invalid")
|
|
|
|
// 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 and root folder
|
|
clientID := fmt.Sprintf("test_client_toolarge_%d", time.Now().UnixNano())
|
|
test.CreateTestClient(t, cfg, clientID, "Test Client TooLarge")
|
|
|
|
// 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")
|
|
}
|
|
|
|
// TestGetDocument_IncludesFileSize verifies that file sizes are returned
|
|
// in the document response when documents have file_size_bytes set.
|
|
func TestGetDocument_IncludesFileSize(t *testing.T) {
|
|
cfg := &ControllerConfig{}
|
|
test.CreateDB(t, cfg)
|
|
initializeTestConfig(t, cfg)
|
|
|
|
svc := createControllerServices(cfg)
|
|
cons := queryapi.NewControllers(svc, cfg)
|
|
|
|
// Create test client with root folder (unique ID to avoid conflicts)
|
|
clientID := fmt.Sprintf("filesize_client_%s", uuid.New().String()[:8])
|
|
test.CreateTestClient(t, cfg, clientID, "FileSize Client")
|
|
|
|
// Create document with file size
|
|
filename := "measured-doc.pdf"
|
|
fileSize := int64(1048576) // 1MB
|
|
docId, err := cfg.GetDBQueries().CreateDocument(t.Context(), &repository.CreateDocumentParams{
|
|
Clientid: clientID,
|
|
Hash: fmt.Sprintf("filesizehash-%s", uuid.New().String()[:8]),
|
|
Filename: &filename,
|
|
FileSizeBytes: &fileSize,
|
|
})
|
|
require.NoError(t, err)
|
|
|
|
ctx, rec := createContext(t)
|
|
|
|
err = cons.GetDocument(ctx, docId, queryapi.GetDocumentParams{})
|
|
require.NoError(t, err)
|
|
assert.Equal(t, http.StatusOK, rec.Code)
|
|
|
|
var result queryapi.DocumentEnriched
|
|
err = json.Unmarshal(rec.Body.Bytes(), &result)
|
|
require.NoError(t, err)
|
|
|
|
// Verify file size is present and correct
|
|
assert.Equal(t, docId, result.Id)
|
|
require.NotNil(t, result.FileSizeBytes, "document should have file size")
|
|
assert.Equal(t, fileSize, *result.FileSizeBytes, "file size should be 1048576 (1MB)")
|
|
}
|
|
|
|
// TestGetDocument_NilFileSize verifies that documents without file sizes
|
|
// (legacy documents) return nil for fileSizeBytes.
|
|
func TestGetDocument_NilFileSize(t *testing.T) {
|
|
cfg := &ControllerConfig{}
|
|
test.CreateDB(t, cfg)
|
|
initializeTestConfig(t, cfg)
|
|
|
|
svc := createControllerServices(cfg)
|
|
cons := queryapi.NewControllers(svc, cfg)
|
|
|
|
// Create test client with root folder (unique ID to avoid conflicts)
|
|
clientID := fmt.Sprintf("legacy_client_%s", uuid.New().String()[:8])
|
|
test.CreateTestClient(t, cfg, clientID, "Legacy Client")
|
|
|
|
// Create document WITHOUT file size (simulating legacy document)
|
|
filename := "legacy-document.pdf"
|
|
docId, err := cfg.GetDBQueries().CreateDocument(t.Context(), &repository.CreateDocumentParams{
|
|
Clientid: clientID,
|
|
Hash: fmt.Sprintf("legacyhash-%s", uuid.New().String()[:8]),
|
|
Filename: &filename,
|
|
FileSizeBytes: nil, // No file size
|
|
})
|
|
require.NoError(t, err)
|
|
|
|
ctx, rec := createContext(t)
|
|
|
|
err = cons.GetDocument(ctx, docId, queryapi.GetDocumentParams{})
|
|
require.NoError(t, err)
|
|
assert.Equal(t, http.StatusOK, rec.Code)
|
|
|
|
var result queryapi.DocumentEnriched
|
|
err = json.Unmarshal(rec.Body.Bytes(), &result)
|
|
require.NoError(t, err)
|
|
|
|
// Verify file size is nil for the legacy document
|
|
assert.Equal(t, docId, result.Id)
|
|
assert.Nil(t, result.FileSizeBytes, "legacy document should have nil file size")
|
|
}
|
|
|
|
func TestDownloadDocument_Handler(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 with root folder
|
|
test.CreateTestClient(t, cfg, "download_client", "Download Client")
|
|
|
|
// Create document with filename
|
|
filename := "test-report.pdf"
|
|
fileContent := []byte("test file content for download")
|
|
docId, err := cfg.GetDBQueries().CreateDocument(t.Context(), &repository.CreateDocumentParams{
|
|
Clientid: "download_client",
|
|
Hash: "download_hash",
|
|
Filename: &filename,
|
|
})
|
|
require.NoError(t, err)
|
|
|
|
// Upload content to S3
|
|
bucket := cfg.GetBucket()
|
|
s3Key := "test/download/" + docId.String()
|
|
_, err = cfg.GetStoreClient().PutObject(t.Context(), &s3.PutObjectInput{
|
|
Bucket: &bucket,
|
|
Key: &s3Key,
|
|
Body: bytes.NewReader(fileContent),
|
|
})
|
|
require.NoError(t, err)
|
|
|
|
// Create document entry linking document to S3 object
|
|
err = cfg.GetDBQueries().AddDocumentEntry(t.Context(), &repository.AddDocumentEntryParams{
|
|
Documentid: docId,
|
|
Bucket: bucket,
|
|
Key: s3Key,
|
|
})
|
|
require.NoError(t, err)
|
|
|
|
// Call handler
|
|
ctx, rec := createContext(t)
|
|
err = cons.DownloadDocument(ctx, docId)
|
|
require.NoError(t, err)
|
|
assert.Equal(t, http.StatusOK, rec.Code)
|
|
assert.Equal(t, "application/pdf", rec.Header().Get("Content-Type"))
|
|
assert.Contains(t, rec.Header().Get("Content-Disposition"), "test-report.pdf")
|
|
assert.Equal(t, fileContent, rec.Body.Bytes())
|
|
}
|
|
|
|
func TestDownloadDocument_Handler_NotFound(t *testing.T) {
|
|
cfg := &ControllerConfig{}
|
|
test.CreateDB(t, cfg)
|
|
initializeTestConfig(t, cfg)
|
|
test.CreateAWSResources(t, cfg)
|
|
|
|
svc := createControllerServices(cfg)
|
|
cons := queryapi.NewControllers(svc, cfg)
|
|
|
|
// Call handler with non-existent document ID
|
|
ctx, _ := createContext(t)
|
|
err := cons.DownloadDocument(ctx, uuid.New())
|
|
require.Error(t, err)
|
|
httpErr, ok := err.(*echo.HTTPError)
|
|
require.True(t, ok)
|
|
assert.Equal(t, http.StatusNotFound, httpErr.Code)
|
|
}
|
|
|
|
func TestDownloadDocument_Handler_ContentDisposition(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 client
|
|
test.CreateTestClient(t, cfg, "dl_disposition_client", "Disposition Client")
|
|
|
|
// Create document with specific filename
|
|
filename := "quarterly-financial-report.pdf"
|
|
fileContent := []byte("financial report content")
|
|
docId, err := cfg.GetDBQueries().CreateDocument(t.Context(), &repository.CreateDocumentParams{
|
|
Clientid: "dl_disposition_client",
|
|
Hash: "disposition_hash",
|
|
Filename: &filename,
|
|
})
|
|
require.NoError(t, err)
|
|
|
|
bucket := cfg.GetBucket()
|
|
s3Key := "test/download/" + docId.String()
|
|
_, err = cfg.GetStoreClient().PutObject(t.Context(), &s3.PutObjectInput{
|
|
Bucket: &bucket,
|
|
Key: &s3Key,
|
|
Body: bytes.NewReader(fileContent),
|
|
})
|
|
require.NoError(t, err)
|
|
|
|
err = cfg.GetDBQueries().AddDocumentEntry(t.Context(), &repository.AddDocumentEntryParams{
|
|
Documentid: docId,
|
|
Bucket: bucket,
|
|
Key: s3Key,
|
|
})
|
|
require.NoError(t, err)
|
|
|
|
ctx, rec := createContext(t)
|
|
err = cons.DownloadDocument(ctx, docId)
|
|
require.NoError(t, err)
|
|
assert.Equal(t, http.StatusOK, rec.Code)
|
|
|
|
// Verify Content-Disposition header contains the original filename
|
|
contentDisposition := rec.Header().Get("Content-Disposition")
|
|
assert.Contains(t, contentDisposition, "attachment")
|
|
assert.Contains(t, contentDisposition, "quarterly-financial-report.pdf")
|
|
}
|
|
|
|
// 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())
|
|
}
|