Merged in feature/assertions (pull request #158)

Assertions + Gen Text

* assert

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], "<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

@@ -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

@@ -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

@@ -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

@@ -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

@@ -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

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