package document_test import ( "testing" "queryorchestration/internal/database/repository" "queryorchestration/internal/document" "queryorchestration/internal/fieldextraction" "queryorchestration/internal/serviceconfig" "queryorchestration/internal/test" "github.com/stretchr/testify/require" "github.com/google/uuid" "github.com/pashagolub/pgxmock/v3" "github.com/stretchr/testify/assert" ) type TestConfig struct { serviceconfig.BaseConfig } func TestGetSummary(t *testing.T) { ctx := t.Context() pool, err := pgxmock.NewPool() require.NoError(t, err) cfg := &serviceconfig.BaseConfig{} cfg.DBPool = pool cfg.DBQueries = repository.New(pool) svc := document.New(cfg) doc := document.DocumentSummary{ ID: uuid.New(), ClientID: "example", Hash: "example_hash", } pool.ExpectQuery("name: GetDocumentSummary :one").WithArgs(doc.ID). WillReturnRows( pgxmock.NewRows([]string{"id", "clientId", "hash"}). AddRow(doc.ID, doc.ClientID, doc.Hash), ) adoc, err := svc.GetSummary(ctx, doc.ID) require.NoError(t, err) assert.Equal(t, &doc, adoc) } func TestGetExternal(t *testing.T) { ctx := t.Context() pool, err := pgxmock.NewPool() require.NoError(t, err) cfg := &serviceconfig.BaseConfig{} cfg.DBPool = pool cfg.DBQueries = repository.New(pool) svc := document.New(cfg) doc := document.DocumentExternal{ Id: uuid.New(), ClientID: "externalID", Hash: "example", Fields: map[string]interface{}{ "json": "hello", }, } pool.ExpectQuery("name: GetDocumentExternal :one").WithArgs(doc.Id). WillReturnRows( pgxmock.NewRows([]string{"id", "clientId", "hash", "fields"}). AddRow(doc.Id, "externalID", "example", []byte(`{"json": "hello"}`)), ) adoc, err := svc.GetExternal(ctx, doc.Id) require.NoError(t, err) assert.Equal(t, &doc, adoc) } func TestGetEnriched(t *testing.T) { ctx := t.Context() cfg := &TestConfig{} test.CreateDB(t, cfg) svc := document.New(cfg) // Create test client clientID := "test-client-enriched" err := cfg.GetDBQueries().CreateClient(ctx, &repository.CreateClientParams{ Clientid: clientID, Name: "Test Client Enriched", }) require.NoError(t, err) t.Run("get enriched document without text record", func(t *testing.T) { filename := "test-enriched.pdf" originalPath := "/documents/test-enriched.pdf" documentID, err := cfg.GetDBQueries().CreateDocument(ctx, &repository.CreateDocumentParams{ Clientid: clientID, Hash: "enrichedhash1", Filename: &filename, Originalpath: &originalPath, }) require.NoError(t, err) // Get enriched document without text record result, err := svc.GetEnriched(ctx, documentID, false) require.NoError(t, err) assert.NotNil(t, result) assert.Equal(t, documentID, result.ID) assert.Equal(t, clientID, result.ClientID) assert.Equal(t, "enrichedhash1", result.Hash) assert.False(t, result.HasTextRecord) assert.Equal(t, filename, *result.Filename) assert.Equal(t, originalPath, *result.OriginalPath) assert.Nil(t, result.TextRecord) assert.Empty(t, result.Labels) }) t.Run("get enriched document with text record requested but not present", func(t *testing.T) { filename := "test-enriched-no-text.pdf" documentID, err := cfg.GetDBQueries().CreateDocument(ctx, &repository.CreateDocumentParams{ Clientid: clientID, Hash: "enrichedhash2", Filename: &filename, }) require.NoError(t, err) // Get enriched document with text record requested result, err := svc.GetEnriched(ctx, documentID, true) require.NoError(t, err) assert.NotNil(t, result) assert.Equal(t, documentID, result.ID) assert.False(t, result.HasTextRecord) assert.Nil(t, result.TextRecord) // No text record exists }) t.Run("get enriched document with text record", func(t *testing.T) { filename := "test-enriched-with-text.pdf" documentID, err := cfg.GetDBQueries().CreateDocument(ctx, &repository.CreateDocumentParams{ Clientid: clientID, Hash: "enrichedhash3", Filename: &filename, }) require.NoError(t, err) // Create field extraction for this document fieldSvc := fieldextraction.New(cfg) contractTitle := "Test Contract" input := &fieldextraction.CreateFieldExtractionInput{ DocumentID: documentID, SingleFields: &repository.AddFieldExtractionParams{ Documentid: documentID, Filename: &filename, Contracttitle: &contractTitle, Createdby: "user123", }, ArrayFields: []*repository.AddFieldExtractionArrayFieldParams{}, } _, err = fieldSvc.CreateFieldExtraction(ctx, input) require.NoError(t, err) // Get enriched document with text record result, err := svc.GetEnriched(ctx, documentID, true) require.NoError(t, err) assert.NotNil(t, result) assert.Equal(t, documentID, result.ID) assert.True(t, result.HasTextRecord) assert.NotNil(t, result.TextRecord) assert.Equal(t, documentID, result.TextRecord.DocumentID) assert.Equal(t, int32(1), result.TextRecord.Version) assert.Equal(t, "user123", result.TextRecord.CreatedBy) }) t.Run("get enriched document with text record but not requesting it", func(t *testing.T) { filename := "test-enriched-skip-text.pdf" documentID, err := cfg.GetDBQueries().CreateDocument(ctx, &repository.CreateDocumentParams{ Clientid: clientID, Hash: "enrichedhash4", Filename: &filename, }) require.NoError(t, err) // Create field extraction fieldSvc := fieldextraction.New(cfg) contractTitle := "Test Contract Skip" input := &fieldextraction.CreateFieldExtractionInput{ DocumentID: documentID, SingleFields: &repository.AddFieldExtractionParams{ Documentid: documentID, Filename: &filename, Contracttitle: &contractTitle, Createdby: "user456", }, ArrayFields: []*repository.AddFieldExtractionArrayFieldParams{}, } _, err = fieldSvc.CreateFieldExtraction(ctx, input) require.NoError(t, err) // Get enriched document without requesting text record result, err := svc.GetEnriched(ctx, documentID, false) require.NoError(t, err) assert.NotNil(t, result) assert.True(t, result.HasTextRecord) // Has text record assert.Nil(t, result.TextRecord) // But not fetched }) t.Run("return error for non-existent document", func(t *testing.T) { nonExistentID := uuid.New() _, err := svc.GetEnriched(ctx, nonExistentID, false) require.Error(t, err) }) }