jacob-mathison/query-orchestration
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Activity

Merged in feature/assertions (pull request #158)

Browse Source 
Assertions + Gen Text

* assert
This commit is contained in:
Michael McGuinness
2025-05-28 12:36:47 +00:00
parent cc2278086f
commit db469e7491
7 changed files with 290 additions and 17 deletions
Download Patch File Download Diff File Expand all files Collapse all files
cmd/textGenerator_test/text/textdetection.go
+104
View File
@@ -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], "<file_base>")
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))
}
cmd/textGenerator_test/textract/textdetection.go
+148
View File
@@ -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], "<file_base>")
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))
}
internal/document/text/documentText.go
+1 -1
View File
@@ -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
internal/document/text/documentText_test.go
+12 -12
View File
@@ -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)
internal/document/text/execute.go
+1 -1
View File
@@ -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
}
internal/document/upload/get_test.go
+21 -1
View File
@@ -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)
}
internal/test/database.go
+3 -2
View File
@@ -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)