From db469e74911d7bf09c4e005ec14901cf6c24e696 Mon Sep 17 00:00:00 2001 From: Michael McGuinness Date: Wed, 28 May 2025 12:36:47 +0000 Subject: [PATCH] Merged in feature/assertions (pull request #158) Assertions + Gen Text * assert --- cmd/textGenerator_test/text/textdetection.go | 104 ++++++++++++ .../textract/textdetection.go | 148 ++++++++++++++++++ internal/document/text/documentText.go | 2 +- internal/document/text/documentText_test.go | 24 +-- internal/document/text/execute.go | 2 +- internal/document/upload/get_test.go | 22 ++- internal/test/database.go | 5 +- 7 files changed, 290 insertions(+), 17 deletions(-) create mode 100644 cmd/textGenerator_test/text/textdetection.go create mode 100644 cmd/textGenerator_test/textract/textdetection.go diff --git a/cmd/textGenerator_test/text/textdetection.go b/cmd/textGenerator_test/text/textdetection.go new file mode 100644 index 00000000..c6184d1e --- /dev/null +++ b/cmd/textGenerator_test/text/textdetection.go @@ -0,0 +1,104 @@ +//go:build aws + +//go:generate go run -tags=aws ./textdetection.go +package main + +import ( + "context" + "fmt" + "io" + "log/slog" + "os" + documenttext "queryorchestration/internal/document/text" + documenttypes "queryorchestration/internal/document/types" + "queryorchestration/internal/serviceconfig" + "queryorchestration/internal/serviceconfig/aws" + "queryorchestration/internal/serviceconfig/objectstore" + "queryorchestration/internal/serviceconfig/queue/querysync" + "queryorchestration/internal/serviceconfig/textract" + "time" +) + +type Config struct { + serviceconfig.BaseConfig + textract.TextractConfig + querysync.QuerySyncConfig + objectstore.ObjectStoreConfig +} + +func main() { + ctx := context.Background() + + if len(os.Args) != 2 { + fmt.Println("Error: expected one argument") + fmt.Println("Usage:", os.Args[0], "") + os.Exit(1) + } + + name := os.Args[1] + + slog.Info("arguments", "file_base", name) + + cfg := &Config{} + profile := aws.Profile(os.Getenv("AWS_PROFILE")) + if profile == "" { + slog.Error("profile not available") + os.Exit(1) + } + + err := cfg.SetTextractClientWithProfile(ctx, profile) + if err != nil { + slog.Error("error creating client", "error", err) + os.Exit(1) + } + + root, err := os.OpenRoot("assets") + if err != nil { + slog.Error("error opening root", "error", err) + os.Exit(1) + } + + pdfFile, err := root.Open(fmt.Sprintf("original/%s.pdf", name)) + if err != nil { + slog.Error("error opening pdf file", "error", err) + os.Exit(1) + } + + pdf, err := documenttypes.NewPDFFromReader(pdfFile) + if err != nil { + slog.Error("error creating pdf file", "error", err) + os.Exit(1) + } + + textSvc := documenttext.New(cfg) + + startTime := time.Now() + + text, err := textSvc.GetDocumentText(ctx, pdf) + if err != nil { + slog.Error("error getting page count", "error", err) + os.Exit(1) + } + + outputFilename := fmt.Sprintf("text/%s.txt", name) + file, err := root.Create(outputFilename) + if err != nil { + slog.Error("error creating file", "error", err) + os.Exit(1) + } + defer file.Close() + + s, err := io.ReadAll(text) + if err != nil { + slog.Error("error reading text", "error", err) + os.Exit(1) + } + + _, err = file.Write(s) + if err != nil { + slog.Error("error writing file", "error", err) + os.Exit(1) + } + + slog.Info("extracted text", "output", outputFilename, "time", time.Since(startTime)) +} diff --git a/cmd/textGenerator_test/textract/textdetection.go b/cmd/textGenerator_test/textract/textdetection.go new file mode 100644 index 00000000..399c19f9 --- /dev/null +++ b/cmd/textGenerator_test/textract/textdetection.go @@ -0,0 +1,148 @@ +//go:build aws + +//go:generate go run -tags=aws ./textdetection.go +package main + +import ( + "context" + "encoding/json" + "fmt" + "log/slog" + "os" + documenttext "queryorchestration/internal/document/text" + documenttypes "queryorchestration/internal/document/types" + "queryorchestration/internal/serviceconfig" + "queryorchestration/internal/serviceconfig/aws" + "queryorchestration/internal/serviceconfig/objectstore" + "queryorchestration/internal/serviceconfig/queue/querysync" + "queryorchestration/internal/serviceconfig/textract" + "sync" + "time" + + awstextract "github.com/aws/aws-sdk-go-v2/service/textract" +) + +type Config struct { + serviceconfig.BaseConfig + textract.TextractConfig + querysync.QuerySyncConfig + objectstore.ObjectStoreConfig +} + +func main() { + ctx := context.Background() + + if len(os.Args) != 2 { + fmt.Println("Error: expected one argument") + fmt.Println("Usage:", os.Args[0], "") + os.Exit(1) + } + + name := os.Args[1] + + slog.Info("arguments", "file_base", name) + + cfg := &Config{} + profile := aws.Profile(os.Getenv("AWS_PROFILE")) + if profile == "" { + slog.Error("profile not available") + os.Exit(1) + } + + err := cfg.SetTextractClientWithProfile(ctx, profile) + if err != nil { + slog.Error("error creating client", "error", err) + os.Exit(1) + } + + root, err := os.OpenRoot("assets") + if err != nil { + slog.Error("error opening root", "error", err) + os.Exit(1) + } + + pdfFile, err := root.Open(fmt.Sprintf("original/%s.pdf", name)) + if err != nil { + slog.Error("error opening pdf file", "error", err) + os.Exit(1) + } + + pdf, err := documenttypes.NewPDFFromReader(pdfFile) + if err != nil { + slog.Error("error creating pdf file", "error", err) + os.Exit(1) + } + + textSvc := documenttext.New(cfg) + + count, err := pdf.GetPageCount(ctx) + if err != nil { + slog.Error("error getting page count", "error", err) + os.Exit(1) + } + + slog.Info("detecting text", "page_count", count) + + startTime := time.Now() + allResults := []map[string]*awstextract.AnalyzeDocumentOutput{} + allResults = make([]map[string]*awstextract.AnalyzeDocumentOutput, count) + var wg sync.WaitGroup + for index := range count { + i := index + wg.Add(1) + go func() { + slog.Info("getting page", "index", i) + + baseResult, err := textSvc.GetBasePage(ctx, pdf, i) + if err != nil { + slog.Error("error detecting text", "error", err) + os.Exit(1) + } + + allResults[i] = map[string]*awstextract.AnalyzeDocumentOutput{ + "base": &baseResult, + } + + features, err := textSvc.ListRequiredFeatures(ctx, baseResult.Blocks, i) + if err != nil { + slog.Error("error listing features", "error", err) + os.Exit(1) + } + + if len(features) > 0 { + fullFeaturesResult, err := textSvc.GetPageWithFeatures(ctx, pdf, i, features) + if err != nil { + slog.Error("error getting features", "error", err) + os.Exit(1) + } + + allResults[i]["full_features"] = &fullFeaturesResult + } + + wg.Done() + }() + } + wg.Wait() + + jsonContent, err := json.MarshalIndent(allResults, "", "\t") + if err != nil { + slog.Error("error marshaling json", "error", err) + os.Exit(1) + } + + outputFilename := fmt.Sprintf("textract/%s.gen", name) + file, err := root.Create(outputFilename) + if err != nil { + slog.Error("error creating file", "error", err) + os.Exit(1) + } + defer file.Close() + + _, err = file.Write(jsonContent) + if err != nil { + slog.Error("error writing file", "error", err) + os.Exit(1) + } + + slog.Info("extracted text", "output", outputFilename, "time", time.Since(startTime)) +} diff --git a/internal/document/text/documentText.go b/internal/document/text/documentText.go index 8acf8d99..c33e3ab4 100644 --- a/internal/document/text/documentText.go +++ b/internal/document/text/documentText.go @@ -10,7 +10,7 @@ import ( documenttypes "queryorchestration/internal/document/types" ) -func (s *Service) getDocumentText(ctx context.Context, doc documenttypes.File) (io.Reader, error) { +func (s *Service) GetDocumentText(ctx context.Context, doc documenttypes.File) (io.Reader, error) { pagesNum, err := doc.GetPageCount(ctx) if err != nil { return nil, err diff --git a/internal/document/text/documentText_test.go b/internal/document/text/documentText_test.go index 1807a430..73cc8677 100644 --- a/internal/document/text/documentText_test.go +++ b/internal/document/text/documentText_test.go @@ -37,7 +37,7 @@ func TestGetDocumentText(t *testing.T) { textExpectations(t, ctx, mockTextract, pdf, textractBaseOut) - text, err := svc.getDocumentText(ctx, pdf) + text, err := svc.GetDocumentText(ctx, pdf) require.NoError(t, err) assertReaders(t, txtFile, text) }) @@ -63,7 +63,7 @@ func TestGetDocumentText(t *testing.T) { ). Return(nil, errors.New("no textract response")) - _, err := svc.getDocumentText(ctx, pdf) + _, err := svc.GetDocumentText(ctx, pdf) require.EqualError(t, err, "no textract response") }) t.Run("merged_cell_table", func(t *testing.T) { @@ -80,7 +80,7 @@ func TestGetDocumentText(t *testing.T) { textExpectations(t, ctx, mockTextract, pdf, textractBaseOut) - text, err := svc.getDocumentText(ctx, pdf) + text, err := svc.GetDocumentText(ctx, pdf) require.NoError(t, err) assertReaders(t, txtFile, text) @@ -99,7 +99,7 @@ func TestGetDocumentText(t *testing.T) { textExpectations(t, ctx, mockTextract, pdf, textractBaseOut) - text, err := svc.getDocumentText(ctx, pdf) + text, err := svc.GetDocumentText(ctx, pdf) require.NoError(t, err) assertReaders(t, txtFile, text) @@ -118,7 +118,7 @@ func TestGetDocumentText(t *testing.T) { textExpectations(t, ctx, mockTextract, pdf, textractBaseOut) - text, err := svc.getDocumentText(ctx, pdf) + text, err := svc.GetDocumentText(ctx, pdf) require.NoError(t, err) assertReaders(t, txtFile, text) @@ -137,7 +137,7 @@ func TestGetDocumentText(t *testing.T) { textExpectations(t, ctx, mockTextract, pdf, textractBaseOut) - text, err := svc.getDocumentText(ctx, pdf) + text, err := svc.GetDocumentText(ctx, pdf) require.NoError(t, err) assertReaders(t, txtFile, text) @@ -156,7 +156,7 @@ func TestGetDocumentText(t *testing.T) { textExpectations(t, ctx, mockTextract, pdf, textractBaseOut) - text, err := svc.getDocumentText(ctx, pdf) + text, err := svc.GetDocumentText(ctx, pdf) require.NoError(t, err) assertReaders(t, txtFile, text) @@ -175,7 +175,7 @@ func TestGetDocumentText(t *testing.T) { textExpectations(t, ctx, mockTextract, pdf, textractBaseOut) - text, err := svc.getDocumentText(ctx, pdf) + text, err := svc.GetDocumentText(ctx, pdf) require.NoError(t, err) assertReaders(t, txtFile, text) @@ -195,7 +195,7 @@ func TestGetDocumentText(t *testing.T) { textExpectations(t, ctx, mockTextract, pdf, textractBaseOut) - text, err := svc.getDocumentText(ctx, pdf) + text, err := svc.GetDocumentText(ctx, pdf) require.NoError(t, err) assertReaders(t, txtFile, text) @@ -214,7 +214,7 @@ func TestGetDocumentText(t *testing.T) { textExpectations(t, ctx, mockTextract, pdf, textractBaseOut) - text, err := svc.getDocumentText(ctx, pdf) + text, err := svc.GetDocumentText(ctx, pdf) require.NoError(t, err) assertReaders(t, txtFile, text) @@ -233,7 +233,7 @@ func TestGetDocumentText(t *testing.T) { textExpectations(t, ctx, mockTextract, pdf, textractBaseOut) - text, err := svc.getDocumentText(ctx, pdf) + text, err := svc.GetDocumentText(ctx, pdf) require.NoError(t, err) assertReaders(t, txtFile, text) @@ -252,7 +252,7 @@ func TestGetDocumentText(t *testing.T) { textExpectations(t, ctx, mockTextract, pdf, textractBaseOut) - text, err := svc.getDocumentText(ctx, pdf) + text, err := svc.GetDocumentText(ctx, pdf) require.NoError(t, err) assertReaders(t, txtFile, text) diff --git a/internal/document/text/execute.go b/internal/document/text/execute.go index fb3ae4e4..e5622c02 100644 --- a/internal/document/text/execute.go +++ b/internal/document/text/execute.go @@ -34,7 +34,7 @@ func (s *Service) executeExtraction(ctx context.Context, documentId uuid.UUID) e return err } - text, err := s.getDocumentText(ctx, doc) + text, err := s.GetDocumentText(ctx, doc) if err != nil { return err } diff --git a/internal/document/upload/get_test.go b/internal/document/upload/get_test.go index 660c75b6..04f8927b 100644 --- a/internal/document/upload/get_test.go +++ b/internal/document/upload/get_test.go @@ -1,6 +1,7 @@ package documentupload_test import ( + "io" "strings" "testing" @@ -11,6 +12,8 @@ import ( documentupload "queryorchestration/internal/document/upload" + "github.com/aws/aws-sdk-go-v2/service/s3" + "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" ) @@ -37,9 +40,26 @@ func TestUpload(t *testing.T) { file := documentupload.File{ ClientID: "client_id", - Content: strings.NewReader("aaaaaaaaaaaa"), + Content: strings.NewReader("abc"), } err = svc.Upload(t.Context(), file) require.NoError(t, err) + + os, err := cfg.StoreClient.ListObjects(t.Context(), &s3.ListObjectsInput{ + Bucket: &cfg.Bucket, + }) + require.NoError(t, err) + + assert.Len(t, os.Contents, 1) + key := *os.Contents[0].Key + + o, err := cfg.StoreClient.GetObject(t.Context(), &s3.GetObjectInput{ + Bucket: &cfg.Bucket, + Key: &key, + }) + require.NoError(t, err) + body, err := io.ReadAll(o.Body) + require.NoError(t, err) + assert.Equal(t, []byte("abc"), body) } diff --git a/internal/test/database.go b/internal/test/database.go index 313c4c62..046eb969 100644 --- a/internal/test/database.go +++ b/internal/test/database.go @@ -5,9 +5,10 @@ import ( "strconv" "testing" - db "queryorchestration/internal/database" "queryorchestration/internal/serviceconfig" + db "queryorchestration/internal/database" + "github.com/docker/go-connections/nat" "github.com/stretchr/testify/require" "github.com/testcontainers/testcontainers-go" @@ -33,7 +34,7 @@ func CreateDBWithParams(t testing.TB, cfg serviceconfig.ConfigProvider, dcfg *Cr cfg.SetDBUser("pass") cfg.SetDBSecret("postgres") - cfg.SetDBName(GetAlias(t, "postgres")) + cfg.SetDBName(GetAlias(t, dbAlias)) cfg.SetDBNoSSL(true) network := GetNetwork(t)