query-orchestration/internal/document/get_test.go

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)
	})
}