Merged in feature/textExtractionsPart3 (pull request #195)
enrich doc responses * enrich doc responses
This commit is contained in:
Jay Brown
@@ -21,6 +21,7 @@ import (
|
||||
|
||||
queryapi "queryorchestration/api/queryAPI"
|
||||
"queryorchestration/internal/database/repository"
|
||||
"queryorchestration/internal/fieldextraction"
|
||||
"queryorchestration/internal/test"
|
||||
)
|
||||
|
||||
@@ -109,17 +110,202 @@ func TestGetDocument(t *testing.T) {
|
||||
|
||||
ctx, rec := createContext(t)
|
||||
|
||||
err = cons.GetDocument(ctx, docId)
|
||||
err = cons.GetDocument(ctx, docId, queryapi.GetDocumentParams{})
|
||||
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{}{},
|
||||
assertBody(t, rec, queryapi.DocumentEnriched{
|
||||
Id: docId,
|
||||
ClientId: "client_id",
|
||||
Hash: "hash",
|
||||
Fields: map[string]interface{}{},
|
||||
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)
|
||||
|
||||
Reference in New Issue
Block a user