Merged in feature/checkbox (pull request #145)

Checkbox Primary Edge Cases

* hascheckboxlogic

* gen

* preppinggen

* exploringcheckbox

* removeunselected

* tests

* failingtest

* failintests

* nomockqueryapi

* db

* name

* nocontainer

* deleterow

* cnc

* extradocsandupdatetextracts

* moretables

* appndedtables

* baseversion

api/clientSyncRunner/runner_test.go

@@ -1,7 +1,6 @@
 package clientsyncrunner_test
 
 import (
-	"context"
 	"fmt"
 	"testing"
 
@@ -10,11 +9,10 @@ import (
 	"queryorchestration/internal/database/repository"
 	"queryorchestration/internal/server/runner"
 	"queryorchestration/internal/serviceconfig/queue/documentsync"
+	"queryorchestration/internal/test"
 	queuemock "queryorchestration/mocks/queue"
 
 	"github.com/aws/aws-sdk-go-v2/service/sqs"
-	"github.com/google/uuid"
-	"github.com/pashagolub/pgxmock/v3"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/mock"
 	"github.com/stretchr/testify/require"
@@ -26,14 +24,11 @@ type ClientSyncConfig struct {
 }
 
 func TestQueryRunner(t *testing.T) {
-	ctx := context.Background()
-
-	pool, err := pgxmock.NewPool()
-	require.NoError(t, err)
-
+	t.Parallel()
 	cfg := &ClientSyncConfig{}
-	cfg.DBPool = pool
-	cfg.DBQueries = repository.New(pool)
+	net := test.GetNetwork(t)
+	test.CreateDB(t, cfg, net, &test.CreateDatabaseConfig{})
+
 	mockSQS := queuemock.NewMockSQSClient(t)
 	cfg.QueueClient = mockSQS
 	cfg.DocumentSyncURL = "/i/am/here"
@@ -44,29 +39,30 @@ func TestQueryRunner(t *testing.T) {
 		ClientSync: svc,
 	})
 
-	t.Run("valid", func(t *testing.T) {
-		bod := clientsyncrunner.Body{
-			ClientID: "hello",
-		}
-
-		docId := uuid.New()
-
-		total := int64(1)
-		pool.ExpectQuery("name: ListDocumentIDsBatch :many").WithArgs(bod.ClientID, int32(100), int32(0)).
-			WillReturnRows(
-				pgxmock.NewRows([]string{"id", "totalCount"}).
-					AddRow(&docId, &total),
-			)
-		mockSQS.EXPECT().
-			SendMessage(
-				mock.Anything,
-				mock.MatchedBy(func(in *sqs.SendMessageInput) bool {
-					return *in.QueueUrl == cfg.DocumentSyncURL && *in.MessageBody == fmt.Sprintf("{\"id\":\"%s\"}", docId.String())
-				}),
-				mock.Anything,
-			).
-			Return(&sqs.SendMessageOutput{}, nil)
-
-		assert.True(t, runner.Process(ctx, bod))
+	err := cfg.GetDBQueries().CreateClient(t.Context(), &repository.CreateClientParams{
+		Clientid: "client_id",
+		Name:     "client_name",
 	})
+	require.NoError(t, err)
+	docId, err := cfg.GetDBQueries().CreateDocument(t.Context(), &repository.CreateDocumentParams{
+		Clientid: "client_id",
+		Hash:     "hash",
+	})
+	require.NoError(t, err)
+
+	bod := clientsyncrunner.Body{
+		ClientID: "client_id",
+	}
+
+	mockSQS.EXPECT().
+		SendMessage(
+			mock.Anything,
+			mock.MatchedBy(func(in *sqs.SendMessageInput) bool {
+				return *in.QueueUrl == cfg.DocumentSyncURL && *in.MessageBody == fmt.Sprintf("{\"id\":\"%s\"}", docId.String())
+			}),
+			mock.Anything,
+		).
+		Return(&sqs.SendMessageOutput{}, nil)
+
+	assert.True(t, runner.Process(t.Context(), bod))
 }

api/docCleanRunner/runner_test.go

@@ -1,7 +1,6 @@
 package doccleanrunner_test
 
 import (
-	"context"
 	"fmt"
 	"io"
 	"strings"
@@ -10,18 +9,17 @@ import (
 
 	doccleanrunner "queryorchestration/api/docCleanRunner"
 	"queryorchestration/internal/database/repository"
-	"queryorchestration/internal/document"
 	documentclean "queryorchestration/internal/document/clean"
 	"queryorchestration/internal/server/runner"
 	"queryorchestration/internal/serviceconfig/objectstore"
 	"queryorchestration/internal/serviceconfig/queue/documenttexttrigger"
+	"queryorchestration/internal/test"
 	objectstoremock "queryorchestration/mocks/objectstore"
 	queuemock "queryorchestration/mocks/queue"
 
 	"github.com/aws/aws-sdk-go-v2/service/s3"
 	"github.com/aws/aws-sdk-go-v2/service/sqs"
 	"github.com/google/uuid"
-	"github.com/pashagolub/pgxmock/v3"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/mock"
 	"github.com/stretchr/testify/require"
@@ -34,14 +32,11 @@ type DocCleanConfig struct {
 }
 
 func TestDocCleanRunner(t *testing.T) {
-	ctx := context.Background()
-
-	pool, err := pgxmock.NewPool()
-	require.NoError(t, err)
-
+	t.Parallel()
 	cfg := &DocCleanConfig{}
-	cfg.DBPool = pool
-	cfg.DBQueries = repository.New(pool)
+	net := test.GetNetwork(t)
+	test.CreateDB(t, cfg, net, &test.CreateDatabaseConfig{})
+
 	mockStore := objectstoremock.NewMockS3Client(t)
 	cfg.StoreClient = mockStore
 	mockSQS := queuemock.NewMockSQSClient(t)
@@ -52,96 +47,73 @@ func TestDocCleanRunner(t *testing.T) {
 		Clean: documentclean.New(cfg),
 	})
 
-	t.Run("valid", func(t *testing.T) {
-		bod := doccleanrunner.Body{
-			DocumentID: uuid.New(),
-		}
-
-		doc := document.DocumentSummary{
-			ID:       bod.DocumentID,
-			ClientID: "hello",
-			Hash:     "example_hash",
-		}
-		bucket := "hello"
-		key := objectstore.BucketKey{
-			CreatedAt: time.Now().UTC(),
-			ClientID:  "hi",
-			EntityID:  uuid.New(),
-			Location:  objectstore.Import,
-		}
-		keyStr := key.String()
-		dbid := doc.ID
-
-		pool.ExpectQuery("name: HasDocumentCleanEntry :one").WithArgs(dbid).WillReturnRows(
-			pgxmock.NewRows([]string{"isclean"}).
-				AddRow(false),
-		)
-		pool.ExpectQuery("name: GetDocumentSummary :one").WithArgs(dbid).
-			WillReturnRows(
-				pgxmock.NewRows([]string{"id", "clientId", "hash"}).
-					AddRow(dbid, doc.ClientID, doc.Hash),
-			)
-		pool.ExpectQuery("name: GetDocumentEntry :one").WithArgs(dbid).WillReturnRows(
-			pgxmock.NewRows([]string{"documentId", "bucket", "key"}).
-				AddRow(dbid, bucket, keyStr),
-		)
-		mimeType := "application/pdf"
-		dbmimetype := repository.NullCleanmimetype{
-			Valid:         true,
-			Cleanmimetype: repository.Cleanmimetype(mimeType),
-		}
-		pool.ExpectBegin()
-		cleanid := uuid.New()
-		pool.ExpectQuery("name: GetMostRecentDocumentCleanEntry :one").WithArgs(dbid).
-			WillReturnRows(
-				pgxmock.NewRows([]string{"id", "documentId", "bucket", "key", "version", "hash", "mimetype", "fail"}),
-			)
-		pool.ExpectQuery("name: AddDocumentClean :one").WithArgs(dbid, &bucket, &keyStr, &doc.Hash, dbmimetype, repository.NullCleanfailtype{}).
-			WillReturnRows(
-				pgxmock.NewRows([]string{"id"}).
-					AddRow(cleanid),
-			)
-		pool.ExpectExec("name: AddDocumentCleanEntry :exec").WithArgs(cleanid, pgxmock.AnyArg()).
-			WillReturnResult(pgxmock.NewResult("", 1))
-		pool.ExpectCommit()
-
-		mockSQS.EXPECT().
-			SendMessage(
-				mock.Anything,
-				mock.MatchedBy(func(in *sqs.SendMessageInput) bool {
-					return *in.QueueUrl == cfg.DocumentTextTriggerURL && *in.MessageBody == fmt.Sprintf("{\"id\":\"%s\"}", doc.ID.String())
-				}),
-				mock.Anything,
-			).
-			Return(&sqs.SendMessageOutput{}, nil)
-
-		mockStore.EXPECT().
-			HeadObject(
-				mock.Anything,
-				mock.MatchedBy(func(in *s3.HeadObjectInput) bool {
-					return *in.Bucket == bucket && *in.Key == keyStr
-				}),
-				mock.Anything,
-			).
-			Return(&s3.HeadObjectOutput{
-				ContentType: &mimeType,
-			}, nil)
-
-		bodyMsg := io.NopCloser(strings.NewReader(pdfHelloWorld))
-		mockStore.EXPECT().
-			GetObject(
-				mock.Anything,
-				mock.MatchedBy(func(in *s3.GetObjectInput) bool {
-					return *in.Bucket == bucket && *in.Key == keyStr
-				}),
-				mock.Anything,
-			).
-			Return(&s3.GetObjectOutput{
-				Body: bodyMsg,
-			}, nil)
-
-		assert.True(t, runner.Process(ctx, bod))
+	err := cfg.GetDBQueries().CreateClient(t.Context(), &repository.CreateClientParams{
+		Clientid: "clientid",
+		Name:     "client_name",
 	})
+	require.NoError(t, err)
+	docId, err := cfg.GetDBQueries().CreateDocument(t.Context(), &repository.CreateDocumentParams{
+		Clientid: "clientid",
+		Hash:     "hash",
+	})
+	require.NoError(t, err)
+	part := uint16(1)
+	filetype := "pdf"
+	key := objectstore.BucketKey{
+		CreatedAt: time.Now().UTC(),
+		ClientID:  "clientid",
+		EntityID:  uuid.New(),
+		Location:  objectstore.Import,
+		Part:      &part,
+		FileType:  &filetype,
+	}
+	err = cfg.GetDBQueries().AddDocumentEntry(t.Context(), &repository.AddDocumentEntryParams{
+		Documentid: docId,
+		Bucket:     "bucket",
+		Key:        key.String(),
+	})
+	require.NoError(t, err)
+
+	mockSQS.EXPECT().
+		SendMessage(
+			mock.Anything,
+			mock.MatchedBy(func(in *sqs.SendMessageInput) bool {
+				return *in.QueueUrl == cfg.DocumentTextTriggerURL && *in.MessageBody == fmt.Sprintf("{\"id\":\"%s\"}", docId.String())
+			}),
+			mock.Anything,
+		).
+		Return(&sqs.SendMessageOutput{}, nil)
+
+	mimeType := "application/pdf"
+	mockStore.EXPECT().
+		HeadObject(
+			mock.Anything,
+			mock.MatchedBy(func(in *s3.HeadObjectInput) bool {
+				return *in.Bucket == "bucket" && *in.Key == key.String()
+			}),
+			mock.Anything,
+		).
+		Return(&s3.HeadObjectOutput{
+			ContentType: &mimeType,
+		}, nil)
+
+	bodyMsg := io.NopCloser(strings.NewReader(pdfHelloWorld))
+	mockStore.EXPECT().
+		GetObject(
+			mock.Anything,
+			mock.MatchedBy(func(in *s3.GetObjectInput) bool {
+				return *in.Bucket == "bucket" && *in.Key == key.String()
+			}),
+			mock.Anything,
+		).
+		Return(&s3.GetObjectOutput{
+			Body: bodyMsg,
+		}, nil)
+
+	bod := doccleanrunner.Body{
+		DocumentID: docId,
+	}
+	assert.True(t, runner.Process(t.Context(), bod))
 }
 
 const pdfHelloWorld = `%PDF-1.4

api/docInitRunner/runner_test.go

@@ -1,23 +1,19 @@
 package docinitrunner_test
 
 import (
-	"context"
-	"fmt"
 	"testing"
 	"time"
 
 	docinitrunner "queryorchestration/api/docInitRunner"
 	"queryorchestration/internal/database/repository"
-	"queryorchestration/internal/document"
 	documentinit "queryorchestration/internal/document/init"
 	"queryorchestration/internal/server/runner"
 	"queryorchestration/internal/serviceconfig/objectstore"
 	"queryorchestration/internal/serviceconfig/queue/documentsync"
+	"queryorchestration/internal/test"
 	queuemock "queryorchestration/mocks/queue"
 
 	"github.com/aws/aws-sdk-go-v2/service/sqs"
-	"github.com/google/uuid"
-	"github.com/pashagolub/pgxmock/v3"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/mock"
 	"github.com/stretchr/testify/require"
@@ -29,14 +25,11 @@ type DocInitConfig struct {
 }
 
 func TestDocInitRunner(t *testing.T) {
-	ctx := context.Background()
-
-	pool, err := pgxmock.NewPool()
-	require.NoError(t, err)
-
+	t.Parallel()
 	cfg := &DocInitConfig{}
-	cfg.DBPool = pool
-	cfg.DBQueries = repository.New(pool)
+	net := test.GetNetwork(t)
+	test.CreateDB(t, cfg, net, &test.CreateDatabaseConfig{})
+
 	mockSQS := queuemock.NewMockSQSClient(t)
 	cfg.QueueClient = mockSQS
 
@@ -44,53 +37,32 @@ func TestDocInitRunner(t *testing.T) {
 		Document: documentinit.New(cfg),
 	})
 
-	t.Run("valid", func(t *testing.T) {
-		clientId := "clientid"
-		bucketName := "bucketName"
-		location := objectstore.BucketKey{
-			Location:  objectstore.Import,
-			ClientID:  clientId,
-			CreatedAt: time.Now().UTC(),
-		}
-		docinfo := document.DocumentSummary{
-			ID:       uuid.New(),
-			ClientID: "hello",
-			Hash:     "example_hash",
-		}
-		doc := docinitrunner.Body{
-			Bucket: bucketName,
-			Key:    location.String(),
-			Hash:   docinfo.Hash,
-		}
-
-		pool.ExpectQuery("name: GetDocumentIDByHash :one").WithArgs(docinfo.Hash, clientId).WillReturnRows(
-			pgxmock.NewRows([]string{"id"}),
-		)
-		pool.ExpectBegin()
-		pool.ExpectQuery("name: CreateDocument :one").WithArgs(clientId, docinfo.Hash).
-			WillReturnRows(
-				pgxmock.NewRows([]string{"id"}).
-					AddRow(docinfo.ID),
-			)
-		pool.ExpectExec("name: AddDocumentEntry :exec").WithArgs(docinfo.ID, bucketName, location.String()).
-			WillReturnResult(pgxmock.NewResult("", 1))
-		pool.ExpectCommit()
-		pool.ExpectQuery("name: GetDocumentSummary :one").WithArgs(docinfo.ID).
-			WillReturnRows(
-				pgxmock.NewRows([]string{"id", "clientId", "hash"}).
-					AddRow(docinfo.ID, docinfo.ClientID, "example"),
-			)
-
-		mockSQS.EXPECT().
-			SendMessage(
-				mock.Anything,
-				mock.MatchedBy(func(in *sqs.SendMessageInput) bool {
-					return *in.QueueUrl == cfg.DocumentSyncURL && *in.MessageBody == fmt.Sprintf("{\"id\":\"%s\"}", docinfo.ID.String())
-				}),
-				mock.Anything,
-			).
-			Return(&sqs.SendMessageOutput{}, nil)
-
-		assert.True(t, runner.Process(ctx, doc))
+	err := cfg.GetDBQueries().CreateClient(t.Context(), &repository.CreateClientParams{
+		Clientid: "clientid",
+		Name:     "client_name",
 	})
+	require.NoError(t, err)
+
+	location := objectstore.BucketKey{
+		Location:  objectstore.Import,
+		ClientID:  "clientid",
+		CreatedAt: time.Now().UTC(),
+	}
+	doc := docinitrunner.Body{
+		Bucket: "bucket",
+		Key:    location.String(),
+		Hash:   "hash",
+	}
+
+	mockSQS.EXPECT().
+		SendMessage(
+			mock.Anything,
+			mock.MatchedBy(func(in *sqs.SendMessageInput) bool {
+				return *in.QueueUrl == cfg.DocumentSyncURL
+			}),
+			mock.Anything,
+		).
+		Return(&sqs.SendMessageOutput{}, nil)
+
+	assert.True(t, runner.Process(t.Context(), doc))
 }

api/docSyncRunner/runner_test.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"fmt"
 	"testing"
+	"time"
 
 	docsyncrunner "queryorchestration/api/docSyncRunner"
 	"queryorchestration/internal/client"
@@ -11,12 +12,13 @@ import (
 	"queryorchestration/internal/document"
 	documentsync "queryorchestration/internal/document/sync"
 	"queryorchestration/internal/server/runner"
+	"queryorchestration/internal/serviceconfig/objectstore"
 	"queryorchestration/internal/serviceconfig/queue/documentclean"
+	"queryorchestration/internal/test"
 	queuemock "queryorchestration/mocks/queue"
 
 	"github.com/aws/aws-sdk-go-v2/service/sqs"
 	"github.com/google/uuid"
-	"github.com/pashagolub/pgxmock/v3"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/mock"
 	"github.com/stretchr/testify/require"
@@ -27,15 +29,13 @@ type DocSyncConfig struct {
 	documentclean.DocCleanConfig
 }
 
-func TestDocInitRunner(t *testing.T) {
+func TestDocSyncRunner(t *testing.T) {
+	t.Parallel()
+	cfg := &DocSyncConfig{}
+	net := test.GetNetwork(t)
+	test.CreateDB(t, cfg, net, &test.CreateDatabaseConfig{})
 	ctx := context.Background()
 
-	pool, err := pgxmock.NewPool()
-	require.NoError(t, err)
-
-	cfg := &DocSyncConfig{}
-	cfg.DBPool = pool
-	cfg.DBQueries = repository.New(pool)
 	mockSQS := queuemock.NewMockSQSClient(t)
 	cfg.QueueClient = mockSQS
 
@@ -46,39 +46,51 @@ func TestDocInitRunner(t *testing.T) {
 		}),
 	})
 
-	t.Run("valid", func(t *testing.T) {
-		j := client.Client{
-			ID: "hello",
-		}
-		docinfo := document.DocumentSummary{
-			ID:       uuid.New(),
-			ClientID: j.ID,
-		}
-		doc := docsyncrunner.Body{
-			DocumentID: docinfo.ID,
-		}
-
-		pool.ExpectQuery("name: GetDocumentSummary :one").WithArgs(docinfo.ID).
-			WillReturnRows(
-				pgxmock.NewRows([]string{"id", "clientId", "hash"}).
-					AddRow(docinfo.ID, docinfo.ClientID, "example"),
-			)
-		pool.ExpectQuery("name: GetClient :one").WithArgs(j.ID).
-			WillReturnRows(
-				pgxmock.NewRows([]string{"id", "name", "canSync"}).
-					AddRow(j.ID, "client_name", true),
-			)
-
-		mockSQS.EXPECT().
-			SendMessage(
-				mock.Anything,
-				mock.MatchedBy(func(in *sqs.SendMessageInput) bool {
-					return *in.QueueUrl == cfg.DocumentCleanURL && *in.MessageBody == fmt.Sprintf("{\"id\":\"%s\"}", doc.DocumentID)
-				}),
-				mock.Anything,
-			).
-			Return(&sqs.SendMessageOutput{}, nil)
-
-		assert.True(t, runner.Process(ctx, doc))
+	err := cfg.GetDBQueries().CreateClient(t.Context(), &repository.CreateClientParams{
+		Clientid: "clientid",
+		Name:     "client_name",
 	})
+	require.NoError(t, err)
+	err = cfg.GetDBQueries().AddClientCanSync(t.Context(), &repository.AddClientCanSyncParams{
+		Clientid: "clientid",
+		Cansync:  true,
+	})
+	require.NoError(t, err)
+	docId, err := cfg.GetDBQueries().CreateDocument(t.Context(), &repository.CreateDocumentParams{
+		Clientid: "clientid",
+		Hash:     "hash",
+	})
+	require.NoError(t, err)
+	part := uint16(1)
+	filetype := "pdf"
+	key := objectstore.BucketKey{
+		CreatedAt: time.Now().UTC(),
+		ClientID:  "clientid",
+		EntityID:  uuid.New(),
+		Location:  objectstore.Import,
+		Part:      &part,
+		FileType:  &filetype,
+	}
+	err = cfg.GetDBQueries().AddDocumentEntry(t.Context(), &repository.AddDocumentEntryParams{
+		Documentid: docId,
+		Bucket:     "bucket",
+		Key:        key.String(),
+	})
+	require.NoError(t, err)
+
+	mockSQS.EXPECT().
+		SendMessage(
+			mock.Anything,
+			mock.MatchedBy(func(in *sqs.SendMessageInput) bool {
+				return *in.QueueUrl == cfg.DocumentCleanURL && *in.MessageBody == fmt.Sprintf("{\"id\":\"%s\"}", docId)
+			}),
+			mock.Anything,
+		).
+		Return(&sqs.SendMessageOutput{}, nil)
+
+	doc := docsyncrunner.Body{
+		DocumentID: docId,
+	}
+
+	assert.True(t, runner.Process(ctx, doc))
 }

api/docTextRunner/runner_test.go

@@ -1,7 +1,6 @@
 package doctextrunner_test
 
 import (
-	"context"
 	"fmt"
 	"io"
 	"strings"
@@ -25,7 +24,6 @@ import (
 	"github.com/aws/aws-sdk-go-v2/service/sqs"
 	awstextract "github.com/aws/aws-sdk-go-v2/service/textract"
 	"github.com/google/uuid"
-	"github.com/pashagolub/pgxmock/v3"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/mock"
 	"github.com/stretchr/testify/require"
@@ -39,14 +37,12 @@ type DocTextConfig struct {
 	objectstore.ObjectStoreConfig
 }
 
-func TestDocCleanRunner(t *testing.T) {
-	ctx := context.Background()
-	pool, err := pgxmock.NewPool()
-	require.NoError(t, err)
-
+func TestDocTextRunner(t *testing.T) {
+	t.Parallel()
 	cfg := &DocTextConfig{}
-	cfg.DBPool = pool
-	cfg.DBQueries = repository.New(pool)
+	net := test.GetNetwork(t)
+	test.CreateDB(t, cfg, net, &test.CreateDatabaseConfig{})
+
 	mockStore := objectstoremock.NewMockS3Client(t)
 	cfg.StoreClient = mockStore
 	mockTextract := textractmock.NewMockTextractClient(t)
@@ -58,84 +54,99 @@ func TestDocCleanRunner(t *testing.T) {
 	runner := doctextrunner.New(&doctextrunner.Services{
 		Text: documenttext.New(cfg),
 	})
-	doc := doctextrunner.Body{
-		DocumentID: uuid.New(),
-	}
 
-	cleanId := uuid.New()
-	bucket := "bucket"
-	hash := `"a06d5e0e795adeb6d0b3732330dbcc15"`
-	textId := uuid.New()
+	err := cfg.GetDBQueries().CreateClient(t.Context(), &repository.CreateClientParams{
+		Clientid: "clientid",
+		Name:     "client_name",
+	})
+	require.NoError(t, err)
+	docId, err := cfg.GetDBQueries().CreateDocument(t.Context(), &repository.CreateDocumentParams{
+		Clientid: "clientid",
+		Hash:     "hash",
+	})
+	require.NoError(t, err)
+	part := uint16(1)
+	filetype := "pdf"
 	key := objectstore.BucketKey{
-		ClientID:  "aa",
-		Location:  objectstore.Import,
 		CreatedAt: time.Now().UTC(),
-		EntityID:  cleanId,
+		ClientID:  "clientid",
+		EntityID:  uuid.New(),
+		Location:  objectstore.Import,
+		Part:      &part,
+		FileType:  &filetype,
 	}
+	err = cfg.GetDBQueries().AddDocumentEntry(t.Context(), &repository.AddDocumentEntryParams{
+		Documentid: docId,
+		Bucket:     "bucket",
+		Key:        key.String(),
+	})
+	require.NoError(t, err)
+	bucket := "bucket"
+	hash := "hash"
 	keyStr := key.String()
-	filename := "../../assets/textract/helloWorld.gen"
-	textractBaseOut := test.GetTextractFileResponse(t, filename)
+	cleanId, err := cfg.GetDBQueries().AddDocumentClean(t.Context(), &repository.AddDocumentCleanParams{
+		Documentid: docId,
+		Bucket:     &bucket,
+		Hash:       &hash,
+		Key:        &keyStr,
+		Mimetype: repository.NullCleanmimetype{
+			Valid:         true,
+			Cleanmimetype: repository.CleanmimetypeApplicationPdf,
+		},
+	})
+	require.NoError(t, err)
+	err = cfg.GetDBQueries().AddDocumentCleanEntry(t.Context(), &repository.AddDocumentCleanEntryParams{
+		Cleanid: cleanId,
+		Version: 1,
+	})
+	require.NoError(t, err)
 
-	mimetype := repository.NullCleanmimetype{
-		Valid:         true,
-		Cleanmimetype: repository.CleanmimetypeApplicationPdf,
-	}
-	pool.ExpectQuery("name: IsDocumentTextExtracted :one").WithArgs(doc.DocumentID).WillReturnRows(
-		pgxmock.NewRows([]string{"isextracted"}).
-			AddRow(false),
-	)
-	pool.ExpectQuery("name: GetCleanEntryByDocId :one").WithArgs(doc.DocumentID).
-		WillReturnRows(
-			pgxmock.NewRows([]string{"id", "documentId", "bucket", "key", "version", "hash", "mimetype", "fail", "clientId"}).
-				AddRow(cleanId, doc.DocumentID, &bucket, &keyStr, int64(123), &hash, mimetype, nil, key.ClientID),
-		)
 	bodyMsg := io.NopCloser(strings.NewReader(pdfHelloWorld))
 	mockStore.EXPECT().
 		GetObject(
 			mock.Anything,
 			mock.MatchedBy(func(in *s3.GetObjectInput) bool {
-				return *in.Bucket == bucket && *in.Key == keyStr
+				return *in.Bucket == "bucket" && *in.Key == key.String()
 			}),
 			mock.Anything,
 		).
 		Return(&s3.GetObjectOutput{
 			Body: bodyMsg,
 		}, nil)
-	sent := 0
+	filename := "../../assets/textract/helloWorld.gen"
+	textractBaseOut := test.GetTextractFileResponse(t, filename)
 	mockTextract.EXPECT().
 		AnalyzeDocument(
 			mock.Anything,
 			mock.MatchedBy(func(in *awstextract.AnalyzeDocumentInput) bool {
-				if sent == 0 {
-					sent++
-					return true
-				}
-				return false
+				return true
 			}),
 			mock.Anything,
 		).
 		Return(textractBaseOut[0]["base"], nil).Once()
-
-	pool.ExpectBegin()
-	pool.ExpectQuery("name: GetDocumentTextExtractionByHash :one").WithArgs(cleanId, hash).WillReturnRows(
-		pgxmock.NewRows([]string{"id"}).
-			AddRow(textId),
-	)
-	pool.ExpectExec("name: AddDocumentTextEntry :exec").WithArgs(textId, pgxmock.AnyArg()).
-		WillReturnResult(pgxmock.NewResult("", 1))
-	pool.ExpectCommit()
-
+	mockStore.EXPECT().
+		PutObject(
+			mock.Anything,
+			mock.MatchedBy(func(in *s3.PutObjectInput) bool {
+				return *in.Bucket == "bucket"
+			}),
+			mock.Anything,
+		).
+		Return(&s3.PutObjectOutput{}, nil)
 	mockQueue.EXPECT().
 		SendMessage(
 			mock.Anything,
 			mock.MatchedBy(func(in *sqs.SendMessageInput) bool {
-				return *in.QueueUrl == cfg.QuerySyncURL && *in.MessageBody == fmt.Sprintf("{\"id\":\"%s\"}", doc.DocumentID.String())
+				return *in.QueueUrl == cfg.QuerySyncURL && *in.MessageBody == fmt.Sprintf("{\"id\":\"%s\"}", docId.String())
 			}),
 			mock.Anything,
 		).
 		Return(&sqs.SendMessageOutput{}, nil)
 
-	assert.True(t, runner.Process(ctx, doc))
+	doc := doctextrunner.Body{
+		DocumentID: docId,
+	}
+	assert.True(t, runner.Process(t.Context(), doc))
 }
 
 const pdfHelloWorld = `%PDF-1.4

api/queryAPI/client_test.go

@@ -2,7 +2,6 @@ package queryapi_test
 
 import (
 	"encoding/json"
-	"fmt"
 	"net/http"
 	"net/http/httptest"
 	"strings"
@@ -14,25 +13,23 @@ import (
 	"queryorchestration/internal/database/repository"
 	"queryorchestration/internal/serviceconfig"
 	"queryorchestration/internal/serviceconfig/queue/clientsync"
+	"queryorchestration/internal/test"
 
 	"github.com/labstack/echo/v4"
-	"github.com/pashagolub/pgxmock/v3"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 )
 
 type ClientConfig struct {
 	serviceconfig.BaseConfig
-	clientsync.ConfigProvider
+	clientsync.ClientSyncConfig
 }
 
 func TestCreateClient(t *testing.T) {
-	pool, err := pgxmock.NewPool()
-	require.NoError(t, err)
-
+	t.Parallel()
 	cfg := &serviceconfig.BaseConfig{}
-	cfg.DBPool = pool
-	cfg.DBQueries = repository.New(pool)
+	net := test.GetNetwork(t)
+	test.CreateDB(t, cfg, net, &test.CreateDatabaseConfig{})
 
 	cons := queryapi.NewControllers(&queryapi.Services{
 		Client: client.New(cfg),
@@ -42,71 +39,62 @@ func TestCreateClient(t *testing.T) {
 		Name: "example_name",
 		Id:   "external_id",
 	}
-	bodyBytes, err := json.Marshal(body)
-	require.NoError(t, err)
 
-	e := echo.New()
-	req := httptest.NewRequest(http.MethodPost, "/", strings.NewReader(string(bodyBytes)))
-	req.Header.Set(echo.HeaderContentType, echo.MIMEApplicationJSON)
-	rec := httptest.NewRecorder()
-	ctx := e.NewContext(req, rec)
+	ctx, rec := createContextWithBody(t, body)
 
-	pool.ExpectExec("name: CreateClient :exec").WithArgs(body.Id, body.Name).
-		WillReturnResult(pgxmock.NewResult("", 1))
-
-	err = cons.CreateClient(ctx)
+	err := cons.CreateClient(ctx)
 	require.NoError(t, err)
 	assert.Equal(t, http.StatusCreated, rec.Code)
-	assert.Equal(t, fmt.Sprintf("{\"id\":\"%s\"}\n", body.Id), rec.Body.String())
+	assertBody(t, rec, queryapi.ClientIDBody{
+		Id: body.Id,
+	})
+
+	client, err := cfg.GetDBQueries().GetClient(t.Context(), "external_id")
+	require.NoError(t, err)
+	assert.Equal(t, "example_name", client.Name)
+	assert.Equal(t, "external_id", client.Clientid)
 }
 
 func TestGetClient(t *testing.T) {
-	pool, err := pgxmock.NewPool()
-	require.NoError(t, err)
-
+	t.Parallel()
 	cfg := &serviceconfig.BaseConfig{}
-	cfg.DBPool = pool
-	cfg.DBQueries = repository.New(pool)
+	net := test.GetNetwork(t)
+	test.CreateDB(t, cfg, net, &test.CreateDatabaseConfig{})
 
 	cons := queryapi.NewControllers(&queryapi.Services{
 		Client: client.New(cfg),
 	})
 
-	e := echo.New()
-	req := httptest.NewRequest(http.MethodPost, "/", strings.NewReader(""))
-	rec := httptest.NewRecorder()
-	ctx := e.NewContext(req, rec)
-
 	id := "client_id"
 
-	ctx.Set("id", id)
+	err := cfg.GetDBQueries().CreateClient(t.Context(), &repository.CreateClientParams{
+		Clientid: id,
+		Name:     "client_name",
+	})
+	require.NoError(t, err)
+	err = cfg.GetDBQueries().AddClientCanSync(t.Context(), &repository.AddClientCanSyncParams{
+		Clientid: id,
+		Cansync:  true,
+	})
+	require.NoError(t, err)
 
-	pool.ExpectQuery("name: GetClient :one").WithArgs(id).WillReturnRows(
-		pgxmock.NewRows([]string{"id", "name", "canSync"}).
-			AddRow("client_id", "client_name", true),
-	)
+	ctx, rec := createContext(t)
 
 	err = cons.GetClient(ctx, id)
 	require.NoError(t, err)
 	assert.Equal(t, http.StatusOK, rec.Code)
-
-	var res queryapi.DocClient
-	err = json.Unmarshal(rec.Body.Bytes(), &res)
-	require.NoError(t, err)
-	assert.EqualExportedValues(t, queryapi.DocClient{
+	assertBody(t, rec, queryapi.DocClient{
 		Id:      id,
 		Name:    "client_name",
 		CanSync: true,
-	}, res)
+	})
 }
 
 func TestUpdateClient(t *testing.T) {
-	pool, err := pgxmock.NewPool()
-	require.NoError(t, err)
-
+	t.Parallel()
 	cfg := &ClientConfig{}
-	cfg.DBPool = pool
-	cfg.DBQueries = repository.New(pool)
+	net := test.GetNetwork(t)
+	test.CreateDB(t, cfg, net, &test.CreateDatabaseConfig{})
 
 	cons := queryapi.NewControllers(&queryapi.Services{
 		ClientUpdate: clientupdate.New(cfg, &clientupdate.Services{
@@ -114,10 +102,36 @@ func TestUpdateClient(t *testing.T) {
 		}),
 	})
 
-	cs := false
+	id := "client_id"
+
+	err := cfg.GetDBQueries().CreateClient(t.Context(), &repository.CreateClientParams{
+		Clientid: id,
+		Name:     "client_name",
+	})
+	require.NoError(t, err)
+
+	newClientName := "new_name"
 	body := queryapi.ClientUpdate{
-		CanSync: &cs,
+		Name: &newClientName,
 	}
+
+	ctx, rec := createContextWithBody(t, body)
+
+	err = cons.UpdateClient(ctx, id)
+	require.NoError(t, err)
+	assert.Equal(t, http.StatusOK, rec.Code)
+	assert.Empty(t, rec.Body.String())
+
+	client, err := cfg.GetDBQueries().GetClient(t.Context(), "client_id")
+	require.NoError(t, err)
+	assert.Equal(t, "new_name", client.Name)
+}
+
+func createContext(t testing.TB) (echo.Context, *httptest.ResponseRecorder) {
+	return createContextWithBody(t, struct{}{})
+}
+
+func createContextWithBody(t testing.TB, body interface{}) (echo.Context, *httptest.ResponseRecorder) {
 	bodyBytes, err := json.Marshal(body)
 	require.NoError(t, err)
 
@@ -125,23 +139,18 @@ func TestUpdateClient(t *testing.T) {
 	req := httptest.NewRequest(http.MethodPost, "/", strings.NewReader(string(bodyBytes)))
 	req.Header.Set(echo.HeaderContentType, echo.MIMEApplicationJSON)
 	rec := httptest.NewRecorder()
-	ctx := e.NewContext(req, rec)
-
-	id := "clientid"
-
-	ctx.Set("id", id)
-
-	pool.ExpectQuery("name: GetClient :one").WithArgs(id).WillReturnRows(
-		pgxmock.NewRows([]string{"id", "name", "canSync"}).
-			AddRow("clientid", "client_name", true),
-	)
-	pool.ExpectBegin()
-	pool.ExpectExec("name: AddClientCanSync :exec").WithArgs(*body.CanSync, id).
-		WillReturnResult(pgxmock.NewResult("", 1))
-	pool.ExpectCommit()
-
-	err = cons.UpdateClient(ctx, id)
-	require.NoError(t, err)
-	assert.Equal(t, http.StatusOK, rec.Code)
-	assert.Empty(t, rec.Body.String())
+	return e.NewContext(req, rec), rec
+}
+
+func assertBody[T any](t testing.TB, rec *httptest.ResponseRecorder, expected T) {
+	res := getBody[T](t, rec)
+	assert.EqualExportedValues(t, expected, res)
+}
+
+func getBody[T any](t testing.TB, rec *httptest.ResponseRecorder) T {
+	var res T
+	err := json.Unmarshal(rec.Body.Bytes(), &res)
+	require.NoError(t, err)
+
+	return res
 }

api/queryAPI/collector_test.go

@@ -1,11 +1,8 @@
 package queryapi_test
 
 import (
-	"encoding/json"
 	"fmt"
 	"net/http"
-	"net/http/httptest"
-	"strings"
 	"testing"
 
 	queryapi "queryorchestration/api/queryAPI"
@@ -13,39 +10,36 @@ import (
 	collectorset "queryorchestration/internal/collector/set"
 	"queryorchestration/internal/database/repository"
 	"queryorchestration/internal/serviceconfig"
-	"queryorchestration/internal/serviceconfig/queue/clientsync"
+	"queryorchestration/internal/test"
 	queuemock "queryorchestration/mocks/queue"
 
 	"github.com/aws/aws-sdk-go-v2/service/sqs"
-	"github.com/google/uuid"
-	"github.com/jackc/pgx/v5"
-	"github.com/labstack/echo/v4"
-	"github.com/pashagolub/pgxmock/v3"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/mock"
 	"github.com/stretchr/testify/require"
 )
 
 func TestSetCollector(t *testing.T) {
-	pool, err := pgxmock.NewPool()
-	require.NoError(t, err)
-	cfg := &struct {
-		serviceconfig.BaseConfig
-		clientsync.ClientSyncConfig
-	}{}
-	cfg.DBPool = pool
-	cfg.DBQueries = repository.New(pool)
+	t.Parallel()
+	cfg := &ClientConfig{}
+	cfg.ClientSyncURL = "example"
+	net := test.GetNetwork(t)
+	test.CreateDB(t, cfg, net, &test.CreateDatabaseConfig{})
+
 	mockSQS := queuemock.NewMockSQSClient(t)
 	cfg.QueueClient = mockSQS
 
 	id := "clientid"
-	current := collector.Collector{
-		ClientID:      id,
-		ActiveVersion: 1,
-		LatestVersion: 4,
-	}
 
-	av := int32(2)
+	err := cfg.GetDBQueries().CreateClient(t.Context(), &repository.CreateClientParams{
+		Clientid: id,
+		Name:     "client_name",
+	})
+	require.NoError(t, err)
+	queryId, err := cfg.GetDBQueries().CreateQuery(t.Context(), repository.QuerytypeContextFull)
+	require.NoError(t, err)
+
+	av := int32(1)
 	cv := int64(1)
 	body := queryapi.CollectorSet{
 		ActiveVersion:         &av,
@@ -53,18 +47,11 @@ func TestSetCollector(t *testing.T) {
 		Fields: &[]queryapi.CollectorField{
 			{
 				Name:    "a",
-				QueryId: uuid.New(),
+				QueryId: queryId,
 			},
 		},
 	}
-	bodyBytes, err := json.Marshal(body)
-	require.NoError(t, err)
-
-	e := echo.New()
-	req := httptest.NewRequest(http.MethodPost, "/", strings.NewReader(string(bodyBytes)))
-	req.Header.Set(echo.HeaderContentType, echo.MIMEApplicationJSON)
-	rec := httptest.NewRecorder()
-	ctx := e.NewContext(req, rec)
+	ctx, rec := createContextWithBody(t, body)
 
 	cons := queryapi.NewControllers(&queryapi.Services{
 		CollectorSet: collectorset.New(cfg, &collectorset.Services{
@@ -72,36 +59,11 @@ func TestSetCollector(t *testing.T) {
 		}),
 	})
 
-	ctx.Set("id", id)
-
-	pool.ExpectQuery("name: GetCollectorByClientID :one").WithArgs(current.ClientID).
-		WillReturnRows(
-			pgxmock.NewRows([]string{"clientId", "minCleanVersion", "minTextVersion", "activeVersion", "latestVersion", "fields"}).
-				AddRow(current.ClientID, current.MinCleanVersion, current.MinTextVersion, current.ActiveVersion, current.LatestVersion, []byte("")),
-		)
-	pool.ExpectQuery("name: AllQueriesExist :one").WithArgs([]uuid.UUID{(*body.Fields)[0].QueryId}).
-		WillReturnRows(
-			pgxmock.NewRows([]string{"exist"}).
-				AddRow(true),
-		)
-	pool.ExpectBeginTx(pgx.TxOptions{})
-	pool.ExpectQuery("name: AddLatestCollectorVersion :one").WithArgs(current.ClientID).WillReturnRows(
-		pgxmock.NewRows([]string{"version"}).
-			AddRow(int32(5)),
-	)
-	pool.ExpectExec("name: SetActiveCollectorVersion :exec").WithArgs(current.ClientID, int32(2)).
-		WillReturnResult(pgxmock.NewResult("", 1))
-	pool.ExpectExec("name: SetCollectorCleanVersion :exec").WithArgs(current.ClientID, int32(5), *body.MinimumCleanerVersion).
-		WillReturnResult(pgxmock.NewResult("", 1))
-	pool.ExpectExec("name: AddCollectorQuery :exec").WithArgs(current.ClientID, "a", (*body.Fields)[0].QueryId, int32(5)).
-		WillReturnResult(pgxmock.NewResult("", 1))
-	pool.ExpectCommit()
-
 	mockSQS.EXPECT().
 		SendMessage(
 			mock.Anything,
 			mock.MatchedBy(func(in *sqs.SendMessageInput) bool {
-				return *in.QueueUrl == cfg.ClientSyncURL && *in.MessageBody == fmt.Sprintf("{\"id\":\"%s\"}", current.ClientID)
+				return *in.QueueUrl == cfg.GetClientSyncURL() && *in.MessageBody == fmt.Sprintf("{\"id\":\"%s\"}", id)
 			}),
 			mock.Anything,
 		).
@@ -113,50 +75,35 @@ func TestSetCollector(t *testing.T) {
 	assert.Empty(t, rec.Body.String())
 }
 
-func TestGetCollectorByclientId(t *testing.T) {
-	pool, err := pgxmock.NewPool()
-	require.NoError(t, err)
+func TestGetCollectorByClientId(t *testing.T) {
+	t.Parallel()
 	cfg := &serviceconfig.BaseConfig{}
-	cfg.DBPool = pool
-	cfg.DBQueries = repository.New(pool)
+	net := test.GetNetwork(t)
+	test.CreateDB(t, cfg, net, &test.CreateDatabaseConfig{})
 
-	e := echo.New()
-	req := httptest.NewRequest(http.MethodPost, "/", strings.NewReader(""))
-	rec := httptest.NewRecorder()
-	ctx := e.NewContext(req, rec)
-
-	svc := collector.New(cfg)
 	cons := queryapi.NewControllers(&queryapi.Services{
-		Collector: svc,
+		Collector: collector.New(cfg),
 	})
 
-	coll := collector.Collector{
-		ClientID:        "hello",
-		MinCleanVersion: 1,
-		MinTextVersion:  2,
-	}
-
 	id := "clientid"
 
-	pool.ExpectQuery("name: GetCollectorByClientID :one").WithArgs(id).
-		WillReturnRows(
-			pgxmock.NewRows([]string{"clientId", "minCleanVersion", "minTextVersion", "activeVersion", "latestVersion", "fields"}).
-				AddRow(coll.ClientID, coll.MinCleanVersion, coll.MinTextVersion, int32(1), int32(2), []byte("")),
-		)
+	err := cfg.GetDBQueries().CreateClient(t.Context(), &repository.CreateClientParams{
+		Clientid: id,
+		Name:     "client_name",
+	})
+	require.NoError(t, err)
+
+	ctx, rec := createContext(t)
 
 	err = cons.GetCollectorByClientId(ctx, id)
 	require.NoError(t, err)
 	assert.Equal(t, http.StatusOK, rec.Code)
-
-	var res queryapi.Collector
-	err = json.Unmarshal(rec.Body.Bytes(), &res)
-	require.NoError(t, err)
-	assert.EqualExportedValues(t, queryapi.Collector{
+	assertBody(t, rec, queryapi.Collector{
 		ClientId:              id,
-		MinimumCleanerVersion: coll.MinCleanVersion,
-		MinimumTextVersion:    coll.MinTextVersion,
-		ActiveVersion:         int32(1),
-		LatestVersion:         int32(2),
+		MinimumCleanerVersion: 0,
+		MinimumTextVersion:    0,
+		ActiveVersion:         0,
+		LatestVersion:         0,
 		Fields:                []queryapi.CollectorField{},
-	}, res)
+	})
 }

api/queryAPI/documents_test.go

@@ -1,10 +1,7 @@
 package queryapi_test
 
 import (
-	"encoding/json"
 	"net/http"
-	"net/http/httptest"
-	"strings"
 	"testing"
 
 	queryapi "queryorchestration/api/queryAPI"
@@ -12,109 +9,82 @@ import (
 	"queryorchestration/internal/document"
 	"queryorchestration/internal/serviceconfig"
 	"queryorchestration/internal/serviceconfig/queue/clientsync"
-	queuemock "queryorchestration/mocks/queue"
+	"queryorchestration/internal/test"
 
-	"github.com/google/uuid"
-	"github.com/labstack/echo/v4"
-	"github.com/pashagolub/pgxmock/v3"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 )
 
 func TestListDocumentsByClientId(t *testing.T) {
-	pool, err := pgxmock.NewPool()
-	require.NoError(t, err)
+	t.Parallel()
 	cfg := &struct {
 		serviceconfig.BaseConfig
 		clientsync.ClientSyncConfig
 	}{}
-	cfg.DBPool = pool
-	cfg.DBQueries = repository.New(pool)
-	mockSQS := queuemock.NewMockSQSClient(t)
-	cfg.QueueClient = mockSQS
-
-	clientId := "clientid"
-
-	e := echo.New()
-	req := httptest.NewRequest(http.MethodPost, "/", strings.NewReader(""))
-	req.Header.Set(echo.HeaderContentType, echo.MIMEApplicationJSON)
-	rec := httptest.NewRecorder()
-	ctx := e.NewContext(req, rec)
+	net := test.GetNetwork(t)
+	test.CreateDB(t, cfg, net, &test.CreateDatabaseConfig{})
 
 	cons := queryapi.NewControllers(&queryapi.Services{
 		Document: document.New(cfg),
 	})
 
-	ctx.Set("id", clientId)
-	doc := queryapi.ListDocuments{
-		{
-			Id:   uuid.New(),
-			Hash: "hash",
-		},
-	}
+	err := cfg.GetDBQueries().CreateClient(t.Context(), &repository.CreateClientParams{
+		Clientid: "client_id",
+		Name:     "client_name",
+	})
+	require.NoError(t, err)
+	docId, err := cfg.GetDBQueries().CreateDocument(t.Context(), &repository.CreateDocumentParams{
+		Clientid: "client_id",
+		Hash:     "hash",
+	})
+	require.NoError(t, err)
 
-	pool.ExpectQuery("name: ListDocumentsByClient :many").WithArgs(clientId).
-		WillReturnRows(
-			pgxmock.NewRows([]string{"id", "hash"}).
-				AddRow(doc[0].Id, "hash"),
-		)
+	ctx, rec := createContext(t)
 
-	err = cons.ListDocumentsByClientId(ctx, clientId)
+	err = cons.ListDocumentsByClientId(ctx, "client_id")
 	require.NoError(t, err)
 	assert.Equal(t, http.StatusOK, rec.Code)
-
-	var res queryapi.ListDocuments
-	err = json.Unmarshal(rec.Body.Bytes(), &res)
-	require.NoError(t, err)
-	assert.EqualExportedValues(t, doc, res)
+	assertBody(t, rec, queryapi.ListDocuments{
+		{
+			Id:   docId,
+			Hash: "hash",
+		},
+	})
 }
 
 func TestGetDocument(t *testing.T) {
-	pool, err := pgxmock.NewPool()
-	require.NoError(t, err)
+	t.Parallel()
 	cfg := &struct {
 		serviceconfig.BaseConfig
 		clientsync.ClientSyncConfig
 	}{}
-	cfg.DBPool = pool
-	cfg.DBQueries = repository.New(pool)
-	mockSQS := queuemock.NewMockSQSClient(t)
-	cfg.QueueClient = mockSQS
-
-	docId := uuid.New()
-
-	e := echo.New()
-	req := httptest.NewRequest(http.MethodPost, "/", strings.NewReader(""))
-	req.Header.Set(echo.HeaderContentType, echo.MIMEApplicationJSON)
-	rec := httptest.NewRecorder()
-	ctx := e.NewContext(req, rec)
+	net := test.GetNetwork(t)
+	test.CreateDB(t, cfg, net, &test.CreateDatabaseConfig{})
 
 	cons := queryapi.NewControllers(&queryapi.Services{
 		Document: document.New(cfg),
 	})
 
-	ctx.Set("id", docId)
-	doc := queryapi.Document{
-		Id:       docId,
-		ClientId: "externalID",
-		Hash:     "example",
-		Fields: map[string]interface{}{
-			"json": "hello",
-		},
-	}
+	err := cfg.GetDBQueries().CreateClient(t.Context(), &repository.CreateClientParams{
+		Clientid: "client_id",
+		Name:     "client_name",
+	})
+	require.NoError(t, err)
+	docId, err := cfg.GetDBQueries().CreateDocument(t.Context(), &repository.CreateDocumentParams{
+		Clientid: "client_id",
+		Hash:     "hash",
+	})
+	require.NoError(t, err)
 
-	pool.ExpectQuery("name: GetDocumentExternal :one").WithArgs(docId).
-		WillReturnRows(
-			pgxmock.NewRows([]string{"id", "clientId", "hash", "fields"}).
-				AddRow(doc.Id, "externalID", "example", []byte(`{"json": "hello"}`)),
-		)
+	ctx, rec := createContext(t)
 
 	err = cons.GetDocument(ctx, docId)
 	require.NoError(t, err)
 	assert.Equal(t, http.StatusOK, rec.Code)
-
-	var res queryapi.Document
-	err = json.Unmarshal(rec.Body.Bytes(), &res)
-	require.NoError(t, err)
-	assert.EqualExportedValues(t, doc, res)
+	assertBody(t, rec, queryapi.Document{
+		Id:       docId,
+		ClientId: "client_id",
+		Hash:     "hash",
+		Fields:   map[string]interface{}{},
+	})
 }

api/queryAPI/export_test.go

@@ -2,23 +2,17 @@ package queryapi_test
 
 import (
 	"net/http"
-	"net/http/httptest"
-	"strings"
 	"testing"
 
 	queryapi "queryorchestration/api/queryAPI"
 
 	"github.com/google/uuid"
-	"github.com/labstack/echo/v4"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 )
 
 func TestTriggerExport(t *testing.T) {
-	e := echo.New()
-	req := httptest.NewRequest(http.MethodPost, "/", strings.NewReader(""))
-	rec := httptest.NewRecorder()
-	ctx := e.NewContext(req, rec)
+	ctx, rec := createContext(t)
 
 	cons := queryapi.NewControllers(&queryapi.Services{})
 
@@ -29,10 +23,7 @@ func TestTriggerExport(t *testing.T) {
 }
 
 func TestExportState(t *testing.T) {
-	e := echo.New()
-	req := httptest.NewRequest(http.MethodPost, "/", strings.NewReader(""))
-	rec := httptest.NewRecorder()
-	ctx := e.NewContext(req, rec)
+	ctx, rec := createContext(t)
 
 	cons := queryapi.NewControllers(&queryapi.Services{})
 

api/queryAPI/query_test.go

@@ -1,12 +1,10 @@
 package queryapi_test
 
 import (
-	"encoding/json"
 	"fmt"
 	"net/http"
-	"net/http/httptest"
-	"strings"
 	"testing"
+	"time"
 
 	queryapi "queryorchestration/api/queryAPI"
 	"queryorchestration/internal/collector"
@@ -14,29 +12,28 @@ import (
 	"queryorchestration/internal/document"
 	"queryorchestration/internal/query"
 	"queryorchestration/internal/query/result"
+	resultprocessor "queryorchestration/internal/query/result/processor"
 	querytest "queryorchestration/internal/query/test"
 	queryupdate "queryorchestration/internal/query/update"
 	"queryorchestration/internal/serviceconfig"
 	"queryorchestration/internal/serviceconfig/queue/queryversionsync"
+	"queryorchestration/internal/test"
 	queuemock "queryorchestration/mocks/queue"
 
+	"github.com/jackc/pgx/v5/pgtype"
 	"github.com/stretchr/testify/require"
 
 	"github.com/aws/aws-sdk-go-v2/service/sqs"
 	"github.com/google/uuid"
-	"github.com/jackc/pgx/v5"
-	"github.com/labstack/echo/v4"
-	"github.com/pashagolub/pgxmock/v3"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/mock"
 )
 
 func TestCreateQuery(t *testing.T) {
-	pool, err := pgxmock.NewPool()
-	require.NoError(t, err)
+	t.Parallel()
 	cfg := &serviceconfig.BaseConfig{}
-	cfg.DBPool = pool
-	cfg.DBQueries = repository.New(pool)
+	net := test.GetNetwork(t)
+	test.CreateDB(t, cfg, net, &test.CreateDatabaseConfig{})
 
 	cons := queryapi.NewControllers(&queryapi.Services{
 		Query: query.New(cfg),
@@ -45,125 +42,79 @@ func TestCreateQuery(t *testing.T) {
 	body := queryapi.QueryCreate{
 		Type: queryapi.CONTEXTFULL,
 	}
-	bodyBytes, err := json.Marshal(body)
-	require.NoError(t, err)
+	ctx, rec := createContextWithBody(t, body)
 
-	e := echo.New()
-	req := httptest.NewRequest(http.MethodPost, "/", strings.NewReader(string(bodyBytes)))
-	req.Header.Set(echo.HeaderContentType, echo.MIMEApplicationJSON)
-	rec := httptest.NewRecorder()
-	ctx := e.NewContext(req, rec)
-
-	id := uuid.New()
-
-	pool.ExpectBeginTx(pgx.TxOptions{})
-	pool.ExpectQuery("name: CreateQuery :one").WithArgs(repository.QuerytypeContextFull).WillReturnRows(
-		pgxmock.NewRows([]string{"id"}).
-			AddRow(id),
-	)
-	pool.ExpectQuery("name: AddLatestQueryVersion :one").WithArgs(id).WillReturnRows(
-		pgxmock.NewRows([]string{"version"}).
-			AddRow(int32(1)),
-	)
-	pool.ExpectExec("name: AddActiveQueryVersion :exec").WithArgs(id, int32(1)).
-		WillReturnResult(pgxmock.NewResult("", 1))
-	pool.ExpectCommit()
-
-	err = cons.CreateQuery(ctx)
+	err := cons.CreateQuery(ctx)
 	require.NoError(t, err)
 	assert.Equal(t, http.StatusCreated, rec.Code)
-	assert.Equal(t, fmt.Sprintf("{\"id\":\"%s\"}\n", id), rec.Body.String())
+	res := getBody[queryapi.IdMessage](t, rec)
+	assert.NotEqual(t, uuid.Nil, res.Id)
 }
 
 func TestListQueries(t *testing.T) {
-	pool, err := pgxmock.NewPool()
-	require.NoError(t, err)
+	t.Parallel()
 	cfg := &serviceconfig.BaseConfig{}
-	cfg.DBPool = pool
-	cfg.DBQueries = repository.New(pool)
+	net := test.GetNetwork(t)
+	test.CreateDB(t, cfg, net, &test.CreateDatabaseConfig{})
 
 	cons := queryapi.NewControllers(&queryapi.Services{
 		Query: query.New(cfg),
 	})
 
-	e := echo.New()
-	req := httptest.NewRequest(http.MethodPost, "/", strings.NewReader(""))
-	rec := httptest.NewRecorder()
-	ctx := e.NewContext(req, rec)
+	ctx, rec := createContext(t)
 
-	id := uuid.New()
-
-	pool.ExpectQuery("name: ListQueries :many").WithArgs().WillReturnRows(
-		pgxmock.NewRows([]string{"id", "type", "activeVersion", "latestVersion", "config", "requiredIds"}).
-			AddRow(id, repository.QuerytypeContextFull, int32(1), int32(2), []byte(""), []uuid.UUID{}),
-	)
+	id, err := cfg.GetDBQueries().CreateQuery(t.Context(), repository.QuerytypeContextFull)
+	require.NoError(t, err)
 
 	err = cons.ListQueries(ctx)
 	require.NoError(t, err)
 	assert.Equal(t, http.StatusOK, rec.Code)
-
-	var res queryapi.ListQueries
-	err = json.Unmarshal(rec.Body.Bytes(), &res)
-	require.NoError(t, err)
-	assert.ElementsMatch(t, res.Queries, []queryapi.Query{
-		{
-			Id:            id,
-			Type:          queryapi.CONTEXTFULL,
-			ActiveVersion: 1,
-			LatestVersion: 2,
+	assertBody(t, rec, queryapi.ListQueries{
+		Queries: []queryapi.Query{
+			{
+				Id:            id,
+				Type:          queryapi.CONTEXTFULL,
+				ActiveVersion: 0,
+				LatestVersion: 0,
+			},
 		},
 	})
 }
 
 func TestGetQuery(t *testing.T) {
-	pool, err := pgxmock.NewPool()
-	require.NoError(t, err)
+	t.Parallel()
 	cfg := &serviceconfig.BaseConfig{}
-	cfg.DBPool = pool
-	cfg.DBQueries = repository.New(pool)
+	net := test.GetNetwork(t)
+	test.CreateDB(t, cfg, net, &test.CreateDatabaseConfig{})
 
 	cons := queryapi.NewControllers(&queryapi.Services{
 		Query: query.New(cfg),
 	})
 
-	e := echo.New()
-	req := httptest.NewRequest(http.MethodPost, "/", strings.NewReader(""))
-	rec := httptest.NewRecorder()
-	ctx := e.NewContext(req, rec)
+	ctx, rec := createContext(t)
 
-	id := uuid.New()
-
-	ctx.Set("id", id)
-
-	pool.ExpectQuery("name: GetQuery :one").WithArgs(id).WillReturnRows(
-		pgxmock.NewRows([]string{"id", "type", "activeVersion", "latestVersion", "config", "requiredIds"}).
-			AddRow(id, repository.QuerytypeContextFull, int32(1), int32(2), nil, []uuid.UUID{}),
-	)
+	id, err := cfg.GetDBQueries().CreateQuery(t.Context(), repository.QuerytypeContextFull)
+	require.NoError(t, err)
 
 	err = cons.GetQuery(ctx, id)
 	require.NoError(t, err)
 	assert.Equal(t, http.StatusOK, rec.Code)
-
-	var res queryapi.Query
-	err = json.Unmarshal(rec.Body.Bytes(), &res)
-	require.NoError(t, err)
-	assert.EqualExportedValues(t, queryapi.Query{
+	assertBody(t, rec, queryapi.Query{
 		Id:            id,
 		Type:          queryapi.CONTEXTFULL,
-		ActiveVersion: 1,
-		LatestVersion: 2,
-	}, res)
+		ActiveVersion: 0,
+		LatestVersion: 0,
+	})
 }
 
 func TestUpdateQuery(t *testing.T) {
-	pool, err := pgxmock.NewPool()
-	require.NoError(t, err)
+	t.Parallel()
 	cfg := &struct {
 		serviceconfig.BaseConfig
 		queryversionsync.QueryVersionSyncConfig
 	}{}
-	cfg.DBPool = pool
-	cfg.DBQueries = repository.New(pool)
+	net := test.GetNetwork(t)
+	test.CreateDB(t, cfg, net, &test.CreateDatabaseConfig{})
 	mockSQS := queuemock.NewMockSQSClient(t)
 	cfg.QueueClient = mockSQS
 	cfg.QueryVersionSyncURL = "here"
@@ -174,32 +125,17 @@ func TestUpdateQuery(t *testing.T) {
 		}),
 	})
 
-	av := int32(2)
-	body := queryapi.QueryUpdate{
-		ActiveVersion: &av,
-	}
-	bodyBytes, err := json.Marshal(body)
+	id, err := cfg.GetDBQueries().CreateQuery(t.Context(), repository.QuerytypeJsonExtractor)
 	require.NoError(t, err)
 
-	e := echo.New()
-	req := httptest.NewRequest(http.MethodPost, "/", strings.NewReader(string(bodyBytes)))
-	req.Header.Set(echo.HeaderContentType, echo.MIMEApplicationJSON)
-	rec := httptest.NewRecorder()
-	ctx := e.NewContext(req, rec)
+	av := int32(1)
+	c := `{"path": "val"}`
+	body := queryapi.QueryUpdate{
+		ActiveVersion: &av,
+		Config:        &c,
+	}
+	ctx, rec := createContextWithBody(t, body)
 
-	id := uuid.New()
-
-	ctx.Set("id", id)
-
-	pool.ExpectQuery("name: GetQuery :one").WithArgs(id).WillReturnRows(
-		pgxmock.NewRows([]string{"id", "type", "activeVersion", "latestVersion", "config", "requiredIds"}).
-			AddRow(id, repository.QuerytypeContextFull, int32(1), int32(2), []byte(""), []uuid.UUID{}),
-	)
-
-	pool.ExpectBeginTx(pgx.TxOptions{})
-	pool.ExpectExec("name: AddActiveQueryVersion :exec").WithArgs(id, av).
-		WillReturnResult(pgxmock.NewResult("", 1))
-	pool.ExpectCommit()
 	mockSQS.EXPECT().
 		SendMessage(
 			mock.Anything,
@@ -217,11 +153,10 @@ func TestUpdateQuery(t *testing.T) {
 }
 
 func TestTestQuery(t *testing.T) {
-	pool, err := pgxmock.NewPool()
-	require.NoError(t, err)
+	t.Parallel()
 	cfg := &serviceconfig.BaseConfig{}
-	cfg.DBPool = pool
-	cfg.DBQueries = repository.New(pool)
+	net := test.GetNetwork(t)
+	test.CreateDB(t, cfg, net, &test.CreateDatabaseConfig{})
 
 	docsvc := document.New(cfg)
 	col := collector.New(cfg)
@@ -238,51 +173,83 @@ func TestTestQuery(t *testing.T) {
 		}),
 	})
 
-	doc := document.DocumentSummary{
-		ID: uuid.New(),
-	}
-	params := &result.Process{
-		QueryID:      uuid.New(),
-		DocumentID:   doc.ID,
-		QueryVersion: int32(3),
-	}
-	body := queryapi.QueryTestRequest{
-		DocumentId:   params.DocumentID,
-		QueryVersion: params.QueryVersion,
-	}
-	bodyBytes, err := json.Marshal(body)
+	contextQueryId, err := que.Create(t.Context(), &resultprocessor.Create{
+		Type: resultprocessor.TypeContextFull,
+	})
 	require.NoError(t, err)
 
-	e := echo.New()
-	req := httptest.NewRequest(http.MethodPost, "/", strings.NewReader(string(bodyBytes)))
-	req.Header.Set(echo.HeaderContentType, echo.MIMEApplicationJSON)
-	rec := httptest.NewRecorder()
-	ctx := e.NewContext(req, rec)
+	err = cfg.GetDBQueries().CreateClient(t.Context(), &repository.CreateClientParams{
+		Clientid: "client_id",
+		Name:     "client_name",
+	})
+	require.NoError(t, err)
+	docId, err := cfg.GetDBQueries().CreateDocument(t.Context(), &repository.CreateDocumentParams{
+		Clientid: "client_id",
+		Hash:     "hash",
+	})
+	require.NoError(t, err)
+	fill := "fill"
+	cleanId, err := cfg.GetDBQueries().AddDocumentClean(t.Context(), &repository.AddDocumentCleanParams{
+		Documentid: docId,
+		Bucket:     &fill,
+		Key:        &fill,
+		Hash:       &fill,
+		Mimetype: repository.NullCleanmimetype{
+			Valid:         true,
+			Cleanmimetype: repository.CleanmimetypeApplicationPdf,
+		},
+	})
+	require.NoError(t, err)
+	err = cfg.GetDBQueries().AddDocumentCleanEntry(t.Context(), &repository.AddDocumentCleanEntryParams{
+		Cleanid: cleanId,
+		Version: 1,
+	})
+	require.NoError(t, err)
+	textId, err := cfg.GetDBQueries().AddDocumentText(t.Context(), &repository.AddDocumentTextParams{
+		Cleanid: cleanId,
+		Bucket:  fill,
+		Key:     fill,
+		Hash:    fill,
+		Createdat: pgtype.Timestamp{
+			Time:  time.Now(),
+			Valid: true,
+		},
+	})
+	require.NoError(t, err)
+	err = cfg.GetDBQueries().AddDocumentTextEntry(t.Context(), &repository.AddDocumentTextEntryParams{
+		Textid:  textId,
+		Version: 1,
+	})
+	require.NoError(t, err)
+	strVal := `{"mykey": "example_value", "oldkey": "old_value"}`
+	_, err = cfg.GetDBQueries().AddResult(t.Context(), &repository.AddResultParams{
+		Queryid:      contextQueryId,
+		Value:        strVal,
+		Textentryid:  textId,
+		Queryversion: 1,
+	})
+	require.NoError(t, err)
 
-	reqID := uuid.New()
-	pool.ExpectQuery("name: GetQueryWithVersion :one").WithArgs(&params.QueryID, &params.QueryVersion).WillReturnRows(
-		pgxmock.NewRows([]string{"id", "type", "activeVersion", "latestVersion", "config", "requiredIds"}).
-			AddRow(params.QueryID, repository.QuerytypeJsonExtractor, int32(1), params.QueryVersion+1, []byte("{\"path\":\"oldkey\"}"), []uuid.UUID{reqID}),
-	)
-	strVal := "{\"mykey\":\"example_value\",\"oldkey\":\"old_value\"}"
-	pool.ExpectQuery("name: ListQueryRequirementValues :many").WithArgs(&params.QueryID, &params.QueryVersion, &params.DocumentID).
-		WillReturnRows(
-			pgxmock.NewRows([]string{"id", "queryId", "type", "value"}).
-				AddRow(&uuid.UUID{}, reqID, repository.QuerytypeContextFull, &strVal),
-		)
-	pool.ExpectQuery("name: GetActiveQueryConfig :one").WithArgs(params.QueryID).WillReturnRows(
-		pgxmock.NewRows([]string{"config"}).
-			AddRow([]byte("{\"path\":\"oldkey\"}")),
-	)
+	c := `{"path": "oldkey"}`
+	queryId, err := que.Create(t.Context(), &resultprocessor.Create{
+		Type:   resultprocessor.TypeJsonExtractor,
+		Config: &c,
+		RequiredQueryIDs: &[]uuid.UUID{
+			contextQueryId,
+		},
+	})
+	require.NoError(t, err)
 
-	err = cons.TestQuery(ctx, params.QueryID)
+	body := queryapi.QueryTestRequest{
+		DocumentId:   docId,
+		QueryVersion: 1,
+	}
+	ctx, rec := createContextWithBody(t, body)
+
+	err = cons.TestQuery(ctx, queryId)
 	require.NoError(t, err)
 	assert.Equal(t, http.StatusOK, rec.Code)
-
-	var res queryapi.QueryTestResponse
-	err = json.Unmarshal(rec.Body.Bytes(), &res)
-	require.NoError(t, err)
-	assert.EqualExportedValues(t, queryapi.QueryTestResponse{
+	assertBody(t, rec, queryapi.QueryTestResponse{
 		Value: "old_value",
-	}, res)
+	})
 }

api/queryAPI/status_test.go

@@ -1,60 +1,81 @@
 package queryapi_test
 
 import (
-	"encoding/json"
 	"net/http"
-	"net/http/httptest"
-	"strings"
 	"testing"
 
 	queryapi "queryorchestration/api/queryAPI"
 	"queryorchestration/internal/client"
 	"queryorchestration/internal/database/repository"
 	"queryorchestration/internal/export"
-	"queryorchestration/internal/serviceconfig"
+	"queryorchestration/internal/test"
 
-	"github.com/labstack/echo/v4"
-	"github.com/pashagolub/pgxmock/v3"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 )
 
 func TestGetClientStatus(t *testing.T) {
-	pool, err := pgxmock.NewPool()
-	require.NoError(t, err)
-	cfg := &serviceconfig.BaseConfig{}
-	cfg.DBPool = pool
-	cfg.DBQueries = repository.New(pool)
+	t.Parallel()
+	cfg := &ClientConfig{}
+	net := test.GetNetwork(t)
+	test.CreateDB(t, cfg, net, &test.CreateDatabaseConfig{})
 
 	cons := queryapi.NewControllers(&queryapi.Services{
 		Client: client.New(cfg),
 		Export: export.New(),
 	})
 
-	e := echo.New()
-	req := httptest.NewRequest(http.MethodPost, "/", strings.NewReader(""))
-	rec := httptest.NewRecorder()
-	ctx := e.NewContext(req, rec)
+	ctx, rec := createContext(t)
 
 	clientId := "clientid"
 
-	pool.ExpectQuery("name: GetClient :one").WithArgs(clientId).WillReturnRows(
-		pgxmock.NewRows([]string{"id", "name", "can_sync"}).
-			AddRow(clientId, "name", true),
-	)
-	pool.ExpectQuery("name: IsClientSynced :one").WithArgs(&clientId).WillReturnRows(
-		pgxmock.NewRows([]string{"issynced"}).
-			AddRow(true),
-	)
+	err := cfg.GetDBQueries().CreateClient(t.Context(), &repository.CreateClientParams{
+		Clientid: clientId,
+		Name:     "client_name",
+	})
+	require.NoError(t, err)
+	err = cfg.GetDBQueries().AddClientCanSync(t.Context(), &repository.AddClientCanSyncParams{
+		Clientid: clientId,
+		Cansync:  true,
+	})
+	require.NoError(t, err)
 
 	err = cons.GetStatusByClientId(ctx, clientId)
 	require.NoError(t, err)
 	assert.Equal(t, http.StatusOK, rec.Code)
-
-	var res queryapi.ClientStatusBody
-	err = json.Unmarshal(rec.Body.Bytes(), &res)
-	require.NoError(t, err)
-	assert.EqualExportedValues(t, queryapi.ClientStatusBody{
+	assertBody(t, rec, queryapi.ClientStatusBody{
 		Status: queryapi.INSYNC,
-	}, res)
+	})
+}
+
+func BenchmarkGetClientStatus(b *testing.B) {
+	cfg := &ClientConfig{}
+	net := test.GetNetwork(b)
+	test.CreateDB(b, cfg, net, &test.CreateDatabaseConfig{})
+
+	cons := queryapi.NewControllers(&queryapi.Services{
+		Client: client.New(cfg),
+		Export: export.New(),
+	})
+
+	ctx, _ := createContext(b)
+
+	clientId := "clientid"
+
+	err := cfg.GetDBQueries().CreateClient(b.Context(), &repository.CreateClientParams{
+		Clientid: clientId,
+		Name:     "client_name",
+	})
+	require.NoError(b, err)
+	err = cfg.GetDBQueries().AddClientCanSync(b.Context(), &repository.AddClientCanSyncParams{
+		Clientid: clientId,
+		Cansync:  true,
+	})
+	require.NoError(b, err)
+
+	b.ResetTimer()
+
+	for b.Loop() {
+		_ = cons.GetStatusByClientId(ctx, clientId)
+	}
 }

api/queryRunner/runner_test.go

@@ -53,80 +53,78 @@ func TestQueryRunner(t *testing.T) {
 		}),
 	})
 
-	t.Run("valid", func(t *testing.T) {
-		doc := queryrunner.Body{
-			DocumentID: uuid.New(),
-			QueryID:    uuid.New(),
-		}
+	doc := queryrunner.Body{
+		DocumentID: uuid.New(),
+		QueryID:    uuid.New(),
+	}
 
-		qcfg := "{\"path\":\"examplekey\"}"
-		reqQuery := uuid.New()
-		query := &resultprocessor.Query{
-			ID:               doc.QueryID,
-			Version:          2,
-			RequiredQueryIDs: &[]uuid.UUID{reqQuery},
-			Config:           &qcfg,
-		}
-		params := &resultset.Set{
-			DocumentID: doc.DocumentID,
-		}
+	qcfg := "{\"path\":\"examplekey\"}"
+	reqQuery := uuid.New()
+	query := &resultprocessor.Query{
+		ID:               doc.QueryID,
+		Version:          2,
+		RequiredQueryIDs: &[]uuid.UUID{reqQuery},
+		Config:           &qcfg,
+	}
+	params := &resultset.Set{
+		DocumentID: doc.DocumentID,
+	}
 
-		textEntryId := uuid.New()
-		pool.ExpectQuery("name: GetTextEntryByDocId :one").WithArgs(params.DocumentID).
-			WillReturnRows(
-				pgxmock.NewRows([]string{"id", "documentId", "bucket", "key", "hash", "cleanId", "extractionVersion"}).
-					AddRow(textEntryId, params.DocumentID, "buket", "/i/am/here", "example", uuid.New(), int64(543)),
-			)
-		pool.ExpectQuery("name: GetQuery :one").WithArgs(query.ID).WillReturnRows(
-			pgxmock.NewRows([]string{"id", "type", "activeVersion", "latestVersion", "config", "requiredIds"}).
-				AddRow(query.ID, repository.QuerytypeJsonExtractor, query.Version, query.Version, []byte(*query.Config), *query.RequiredQueryIDs),
+	textEntryId := uuid.New()
+	pool.ExpectQuery("name: GetTextEntryByDocId :one").WithArgs(params.DocumentID).
+		WillReturnRows(
+			pgxmock.NewRows([]string{"id", "documentId", "bucket", "key", "hash", "cleanId", "extractionVersion"}).
+				AddRow(textEntryId, params.DocumentID, "buket", "/i/am/here", "example", uuid.New(), int64(543)),
 		)
-		pool.ExpectQuery("name: GetQueryWithVersion :one").WithArgs(&query.ID, &query.Version).WillReturnRows(
-			pgxmock.NewRows([]string{"id", "type", "activeVersion", "latestVersion", "config", "requiredIds"}).
-				AddRow(query.ID, repository.QuerytypeJsonExtractor, query.Version, query.Version, []byte(*query.Config), *query.RequiredQueryIDs),
+	pool.ExpectQuery("name: GetQuery :one").WithArgs(query.ID).WillReturnRows(
+		pgxmock.NewRows([]string{"id", "type", "activeVersion", "latestVersion", "config", "requiredIds"}).
+			AddRow(query.ID, repository.QuerytypeJsonExtractor, query.Version, query.Version, []byte(*query.Config), *query.RequiredQueryIDs),
+	)
+	pool.ExpectQuery("name: GetQueryWithVersion :one").WithArgs(&query.ID, &query.Version).WillReturnRows(
+		pgxmock.NewRows([]string{"id", "type", "activeVersion", "latestVersion", "config", "requiredIds"}).
+			AddRow(query.ID, repository.QuerytypeJsonExtractor, query.Version, query.Version, []byte(*query.Config), *query.RequiredQueryIDs),
+	)
+	requiredResultId := uuid.New()
+	strVal := "{\"examplekey\":\"example_value\"}"
+	pool.ExpectQuery("name: ListQueryRequirementValues :many").WithArgs(&query.ID, &query.Version, &params.DocumentID).
+		WillReturnRows(
+			pgxmock.NewRows([]string{"id", "queryId", "type", "value"}).
+				AddRow(&requiredResultId, (*query.RequiredQueryIDs)[0], repository.QuerytypeContextFull, &strVal),
 		)
-		requiredResultId := uuid.New()
-		strVal := "{\"examplekey\":\"example_value\"}"
-		pool.ExpectQuery("name: ListQueryRequirementValues :many").WithArgs(&query.ID, &query.Version, &params.DocumentID).
-			WillReturnRows(
-				pgxmock.NewRows([]string{"id", "queryId", "type", "value"}).
-					AddRow(&requiredResultId, (*query.RequiredQueryIDs)[0], repository.QuerytypeContextFull, &strVal),
-			)
-		pool.ExpectQuery("name: GetActiveQueryConfig :one").WithArgs(query.ID).WillReturnRows(
-			pgxmock.NewRows([]string{"config"}).
-				AddRow([]byte(qcfg)),
+	pool.ExpectQuery("name: GetActiveQueryConfig :one").WithArgs(query.ID).WillReturnRows(
+		pgxmock.NewRows([]string{"config"}).
+			AddRow([]byte(qcfg)),
+	)
+	pool.ExpectBegin()
+	resultId := uuid.New()
+	pool.ExpectQuery("name: AddResult :one").WithArgs(query.ID, pgxmock.AnyArg(), textEntryId, query.Version).
+		WillReturnRows(
+			pgxmock.NewRows([]string{"id"}).
+				AddRow(resultId),
+		)
+	pool.ExpectQuery("name: ListQueryRequirementValues :many").WithArgs(&query.ID, &query.Version, &params.DocumentID).
+		WillReturnRows(
+			pgxmock.NewRows([]string{"id", "queryId", "type", "value"}).
+				AddRow(&requiredResultId, (*query.RequiredQueryIDs)[0], repository.QuerytypeContextFull, &strVal),
+		)
+	pool.ExpectExec("name: AddResultDependency :exec").WithArgs(resultId, requiredResultId).
+		WillReturnResult(pgxmock.NewResult("", 1))
+	pool.ExpectCommit()
+	pool.ExpectQuery("name: ListQueryDirectDependentsByDocumentID :many").WithArgs(&query.ID, &doc.DocumentID).
+		WillReturnRows(
+			pgxmock.NewRows([]string{"queryId"}).
+				AddRow(&reqQuery),
 		)
-		pool.ExpectBegin()
-		resultId := uuid.New()
-		pool.ExpectQuery("name: AddResult :one").WithArgs(query.ID, pgxmock.AnyArg(), textEntryId, query.Version).
-			WillReturnRows(
-				pgxmock.NewRows([]string{"id"}).
-					AddRow(resultId),
-			)
-		pool.ExpectQuery("name: ListQueryRequirementValues :many").WithArgs(&query.ID, &query.Version, &params.DocumentID).
-			WillReturnRows(
-				pgxmock.NewRows([]string{"id", "queryId", "type", "value"}).
-					AddRow(&requiredResultId, (*query.RequiredQueryIDs)[0], repository.QuerytypeContextFull, &strVal),
-			)
-		pool.ExpectExec("name: AddResultDependency :exec").WithArgs(resultId, requiredResultId).
-			WillReturnResult(pgxmock.NewResult("", 1))
-		pool.ExpectCommit()
-		pool.ExpectQuery("name: ListQueryDirectDependentsByDocumentID :many").WithArgs(&query.ID, &doc.DocumentID).
-			WillReturnRows(
-				pgxmock.NewRows([]string{"queryId"}).
-					AddRow(&reqQuery),
-			)
 
-		mockSQS.EXPECT().
-			SendMessage(
-				mock.Anything,
-				mock.MatchedBy(func(in *sqs.SendMessageInput) bool {
-					return *in.QueueUrl == cfg.QueryURL && *in.MessageBody == fmt.Sprintf("{\"document_id\":\"%s\",\"query_id\":\"%s\"}", params.DocumentID.String(), (*query.RequiredQueryIDs)[0].String())
-				}),
-				mock.Anything,
-			).
-			Return(&sqs.SendMessageOutput{}, nil)
+	mockSQS.EXPECT().
+		SendMessage(
+			mock.Anything,
+			mock.MatchedBy(func(in *sqs.SendMessageInput) bool {
+				return *in.QueueUrl == cfg.QueryURL && *in.MessageBody == fmt.Sprintf("{\"document_id\":\"%s\",\"query_id\":\"%s\"}", params.DocumentID.String(), (*query.RequiredQueryIDs)[0].String())
+			}),
+			mock.Anything,
+		).
+		Return(&sqs.SendMessageOutput{}, nil)
 
-		assert.True(t, runner.Process(ctx, doc))
-	})
+	assert.True(t, runner.Process(ctx, doc))
 }

api/querySyncRunner/runner_test.go

@@ -47,37 +47,35 @@ func TestQueryRunner(t *testing.T) {
 		QuerySync: svc,
 	})
 
-	t.Run("valid", func(t *testing.T) {
-		doc := querysyncrunner.Body{
-			DocumentID: uuid.New(),
-		}
+	doc := querysyncrunner.Body{
+		DocumentID: uuid.New(),
+	}
 
-		reqId := uuid.New()
-		qs := []uuid.UUID{
-			reqId,
-		}
+	reqId := uuid.New()
+	qs := []uuid.UUID{
+		reqId,
+	}
 
-		pool.ExpectQuery("name: IsDocumentTextExtracted :one").WithArgs(doc.DocumentID).
-			WillReturnRows(
-				pgxmock.NewRows([]string{"isextracted"}).
-					AddRow(true),
-			)
-		pool.ExpectQuery("name: ListUnsyncedNoDepsQueriesByDocId :many").WithArgs(&doc.DocumentID).
-			WillReturnRows(
-				pgxmock.NewRows([]string{"id"}).
-					AddRow(&reqId),
-			)
+	pool.ExpectQuery("name: IsDocumentTextExtracted :one").WithArgs(doc.DocumentID).
+		WillReturnRows(
+			pgxmock.NewRows([]string{"isextracted"}).
+				AddRow(true),
+		)
+	pool.ExpectQuery("name: ListUnsyncedNoDepsQueriesByDocId :many").WithArgs(&doc.DocumentID).
+		WillReturnRows(
+			pgxmock.NewRows([]string{"id"}).
+				AddRow(&reqId),
+		)
 
-		mockSQS.EXPECT().
-			SendMessage(
-				mock.Anything,
-				mock.MatchedBy(func(in *sqs.SendMessageInput) bool {
-					return *in.QueueUrl == cfg.QueryURL && *in.MessageBody == fmt.Sprintf("{\"document_id\":\"%s\",\"query_id\":\"%s\"}", doc.DocumentID.String(), qs[0].String())
-				}),
-				mock.Anything,
-			).
-			Return(&sqs.SendMessageOutput{}, nil)
+	mockSQS.EXPECT().
+		SendMessage(
+			mock.Anything,
+			mock.MatchedBy(func(in *sqs.SendMessageInput) bool {
+				return *in.QueueUrl == cfg.QueryURL && *in.MessageBody == fmt.Sprintf("{\"document_id\":\"%s\",\"query_id\":\"%s\"}", doc.DocumentID.String(), qs[0].String())
+			}),
+			mock.Anything,
+		).
+		Return(&sqs.SendMessageOutput{}, nil)
 
-		assert.True(t, runner.Process(ctx, doc))
-	})
+	assert.True(t, runner.Process(ctx, doc))
 }

api/queryVersionSyncRunner/runner_test.go

@@ -44,44 +44,42 @@ func TestQueryRunner(t *testing.T) {
 		Sync: svc,
 	})
 
-	t.Run("valid", func(t *testing.T) {
-		doc := queryversionsyncrunner.Body{
-			QueryID: uuid.New(),
-		}
+	doc := queryversionsyncrunner.Body{
+		QueryID: uuid.New(),
+	}
 
-		clientOne := "hello"
-		clientTwo := "bye"
-		clientIds := []string{
-			clientOne,
-			clientTwo,
-		}
+	clientOne := "hello"
+	clientTwo := "bye"
+	clientIds := []string{
+		clientOne,
+		clientTwo,
+	}
 
-		pool.ExpectQuery("name: ListQueryClientIDs :many").WithArgs(&doc.QueryID).
-			WillReturnRows(
-				pgxmock.NewRows([]string{"clientId"}).
-					AddRow(clientOne).
-					AddRow(clientTwo),
-			)
+	pool.ExpectQuery("name: ListQueryClientIDs :many").WithArgs(&doc.QueryID).
+		WillReturnRows(
+			pgxmock.NewRows([]string{"clientId"}).
+				AddRow(clientOne).
+				AddRow(clientTwo),
+		)
 
-		mockSQS.EXPECT().
-			SendMessage(
-				mock.Anything,
-				mock.MatchedBy(func(in *sqs.SendMessageInput) bool {
-					return *in.QueueUrl == cfg.ClientSyncURL && *in.MessageBody == fmt.Sprintf("{\"id\":\"%s\"}", clientIds[0])
-				}),
-				mock.Anything,
-			).
-			Return(&sqs.SendMessageOutput{}, nil)
-		mockSQS.EXPECT().
-			SendMessage(
-				mock.Anything,
-				mock.MatchedBy(func(in *sqs.SendMessageInput) bool {
-					return *in.QueueUrl == cfg.ClientSyncURL && *in.MessageBody == fmt.Sprintf("{\"id\":\"%s\"}", clientIds[1])
-				}),
-				mock.Anything,
-			).
-			Return(&sqs.SendMessageOutput{}, nil)
+	mockSQS.EXPECT().
+		SendMessage(
+			mock.Anything,
+			mock.MatchedBy(func(in *sqs.SendMessageInput) bool {
+				return *in.QueueUrl == cfg.ClientSyncURL && *in.MessageBody == fmt.Sprintf("{\"id\":\"%s\"}", clientIds[0])
+			}),
+			mock.Anything,
+		).
+		Return(&sqs.SendMessageOutput{}, nil)
+	mockSQS.EXPECT().
+		SendMessage(
+			mock.Anything,
+			mock.MatchedBy(func(in *sqs.SendMessageInput) bool {
+				return *in.QueueUrl == cfg.ClientSyncURL && *in.MessageBody == fmt.Sprintf("{\"id\":\"%s\"}", clientIds[1])
+			}),
+			mock.Anything,
+		).
+		Return(&sqs.SendMessageOutput{}, nil)
 
-		assert.True(t, runner.Process(ctx, doc))
-	})
+	assert.True(t, runner.Process(ctx, doc))
 }

api/storeEventRunner/runner_test.go

@@ -43,64 +43,62 @@ func TestDocInitRunner(t *testing.T) {
 		documentstore.New(cfg),
 	})
 
-	t.Run("valid", func(t *testing.T) {
-		clientId := "hi"
-		bucketName := "bucketName"
-		location := objectstore.BucketKey{
-			Location:  objectstore.Import,
-			ClientID:  clientId,
-			CreatedAt: time.Now().UTC(),
-		}
-		docinfo := document.DocumentSummary{
-			ID:       uuid.New(),
-			ClientID: clientId,
-			Hash:     "example_hash",
-		}
-		doc := S3EventNotification{
-			Records: []S3EventRecord{
-				{
-					EventName: EventS3ObjectCreatedPut,
-					S3: S3EventRecordDetails{
-						Bucket: S3Bucket{
-							Name: bucketName,
-						},
-						Object: S3Object{
-							Key:  location.String(),
-							ETag: docinfo.Hash,
-						},
+	clientId := "hi"
+	bucketName := "bucketName"
+	location := objectstore.BucketKey{
+		Location:  objectstore.Import,
+		ClientID:  clientId,
+		CreatedAt: time.Now().UTC(),
+	}
+	docinfo := document.DocumentSummary{
+		ID:       uuid.New(),
+		ClientID: clientId,
+		Hash:     "example_hash",
+	}
+	doc := S3EventNotification{
+		Records: []S3EventRecord{
+			{
+				EventName: EventS3ObjectCreatedPut,
+				S3: S3EventRecordDetails{
+					Bucket: S3Bucket{
+						Name: bucketName,
+					},
+					Object: S3Object{
+						Key:  location.String(),
+						ETag: docinfo.Hash,
 					},
 				},
 			},
-		}
+		},
+	}
 
-		pool.ExpectQuery("name: GetDocumentIDByHash :one").WithArgs(docinfo.Hash, clientId).WillReturnRows(
-			pgxmock.NewRows([]string{"id"}),
+	pool.ExpectQuery("name: GetDocumentIDByHash :one").WithArgs(docinfo.Hash, clientId).WillReturnRows(
+		pgxmock.NewRows([]string{"id"}),
+	)
+	pool.ExpectBegin()
+	pool.ExpectQuery("name: CreateDocument :one").WithArgs(clientId, docinfo.Hash).
+		WillReturnRows(
+			pgxmock.NewRows([]string{"id"}).
+				AddRow(docinfo.ID),
+		)
+	pool.ExpectExec("name: AddDocumentEntry :exec").WithArgs(docinfo.ID, bucketName, location.String()).
+		WillReturnResult(pgxmock.NewResult("", 1))
+	pool.ExpectCommit()
+	pool.ExpectQuery("name: GetDocumentSummary :one").WithArgs(docinfo.ID).
+		WillReturnRows(
+			pgxmock.NewRows([]string{"id", "clientId", "hash"}).
+				AddRow(docinfo.ID, docinfo.ClientID, "example"),
 		)
-		pool.ExpectBegin()
-		pool.ExpectQuery("name: CreateDocument :one").WithArgs(clientId, docinfo.Hash).
-			WillReturnRows(
-				pgxmock.NewRows([]string{"id"}).
-					AddRow(docinfo.ID),
-			)
-		pool.ExpectExec("name: AddDocumentEntry :exec").WithArgs(docinfo.ID, bucketName, location.String()).
-			WillReturnResult(pgxmock.NewResult("", 1))
-		pool.ExpectCommit()
-		pool.ExpectQuery("name: GetDocumentSummary :one").WithArgs(docinfo.ID).
-			WillReturnRows(
-				pgxmock.NewRows([]string{"id", "clientId", "hash"}).
-					AddRow(docinfo.ID, docinfo.ClientID, "example"),
-			)
 
-		mockSQS.EXPECT().
-			SendMessage(
-				mock.Anything,
-				mock.MatchedBy(func(in *sqs.SendMessageInput) bool {
-					return *in.QueueUrl == cfg.DocInitURL
-				}),
-				mock.Anything,
-			).
-			Return(&sqs.SendMessageOutput{}, nil)
+	mockSQS.EXPECT().
+		SendMessage(
+			mock.Anything,
+			mock.MatchedBy(func(in *sqs.SendMessageInput) bool {
+				return *in.QueueUrl == cfg.DocInitURL
+			}),
+			mock.Anything,
+		).
+		Return(&sqs.SendMessageOutput{}, nil)
 
-		assert.True(t, runner.Process(ctx, doc))
-	})
+	assert.True(t, runner.Process(ctx, doc))
 }

assets/text/hn_23-70.txt

@@ -181,7 +181,7 @@ Addendum C.2. DIVISION OF FINANCIAL RESPONSIBILITY MATRIX OF HMO AND PPG
 CAPITATED SERVICES MEDICARE BENEFIT PROGRAM of the Agreement shall be amended only for the
 following Service categories as indicated below:
 -------Table Start--------
-[["","PPG CAPITATED SERVICES","HMO RISK SERVICES","SHARED RISK/HOSPITAL CAPITATED SERVICES"],["CHEMICAL DEPENDENCE","","",""],["(Rehabilitation Services)","","",""],["Inpatient Facility Component","","X",""],["Inpatient Professional Component","","",""],["- Outpatient Facility Component","","",""],["- Outpatient Professional Component","","X",""],["Inpatient Detox Facility Component","","","X"],["Inpatient Detox Professional Component","","",""],["MENTAL HEALTH - Inpatient","","",""],["Facility Component","","X",""],["Professional Component","","X",""],["MENTAL HEALTH - Outpatient","","",""],["Facility Component","","",""],["Professional Component","","X",""]]
+[["","PPG CAPITATED SERVICES","HMO RISK SERVICES","SHARED RISK/HOSPITAL CAPITATED SERVICES"],["CHEMICAL DEPENDENCE","","",""],["(Rehabilitation Services)","","",""],["MENTAL HEALTH - Inpatient","","",""],["MENTAL HEALTH - Outpatient","","",""]]
 -------Table End--------
 7.
 Addendum E is deleted in its entirety and is replaced with a new Addendum E attached

assets/text/multi_column.txt

@@ -0,0 +1,121 @@
+Start of Page No. = 1
+
+employee of Centene and such Assigned Worker enrolls
+in a medical plan that constitutes minimum essential
+coverage (as defined in Code Section 5000A) provided
+by Vendor for the period beginning on or after January 1.
+2015. Centene hereby agrees to pay an additional fee to
+Vendor with respect to such Assigned Worker ("Health
+Plan Enrollment Fee"). but only insofar as applicable
+regulations under Section 4980H of the Code require the
+imposition of such a fee in order to treat the medical plan
+coverage provided by Vendor as being provided on
+behalf of Centene The amount of any such Health Plan
+Enrollment Fee will be determined by Centene and
+Vendor at the time of any such re-characterization but
+will not exceed [$0.05] per hour for all hours of work
+performed by any such re-characterized Assigned Worker
+(including hours of work previously performed by the
+Assigned Worker and billed to Centene) during the period
+of time for which (i) the Assigned Worker was or is
+enrolled in Vendor's medical plan and (ii) the Assigned
+Worker was or is characterized as a common law
+employee of the Company If applicable the Health Plan
+Enrollment Fee will appear separately on each invoice
+provided by Vendor to Centene
+(b) Vendor represents and warrants to Centene that
+Vendor is solely responsible for any assessable payment
+under Section 4980H of the Code assessed against
+Vendor, even if any Assigned Worker is re-characterized
+by the Internal Revenue Service Department of Labor or
+a court of law as a common law employee of Centene
+Vendor further represents and warrants to Centene that
+should an assessable payment under Section 4980H of
+the Code be assessed against Centene due to Vendor's
+failure to offer any Assigned Worker Affordable Group
+Health Plan Coverage as required by this agreement
+upon notice from Centene Vendor shall fully indemnify
+and promptly reimburse Centene for such assessable
+payment
+11.17 Entire Agreement This Agreement its addenda
+all SOWs. Change Orders and all exhibits and addenda
+thereto are incorporated herein and constitute the entire
+agreement of the parties This Agreement supersedes all
+prior and contemporaneous negotiations representations
+promises, and agreements concerning the subject matter
+herein whether written or oral
+Non-Solicitation During the term of this
+Agreement and for one (1) year thereafter neither party
+shall without the prior written consent of the other party,
+which may be withheld at such other party's sole
+discretion solicit for hire any person or contractor
+employed by the other party then or within the preceding
+twelve (12) months For this purpose, solicitation does
+not include contact resulting from indirect means such as
+public advertisement placement firm searches or similar
+means not directed specifically at the employee to which
+the employee responds on his or her own initiative nor
+shall it include contacts initiated by the employee If a
+party breaches this Non-Solicitation provision the
+breaching party. as its sole liability and as the exclusive
+remedy to the non-breaching party, shall pay
+compensation to the non-breaching party in the form of
+liquidated damages equal to three (3) months of the
+solicited employee's starting base compensation with the
+non-breaching party
+11.19 Limitation of Liability NEITHER PARTY'S TOTAL
+LIABILITY RELATING TO THIS AGREEMENT SHALL
+EXCEED THE GREATER OF (a) TWO TIMES THE
+AMOUNT OF FEES PAYABLE UNDER THE SERVICES
+AGREEMENTS DURING THE TWELVE (12) MONTHS
+IMMEDIATELY PRECEDING THE EVENT THAT GAVE
+RISE TO A PARTY'S CLAIM AND (b) TEN MILLION
+DOLLARS ($10,000,000) NOR SHALL EITHER PARTY
+BE LIABLE TO THE OTHER FOR ANY SPECIAL
+CONSEQUENTIAL INCIDENTAL OR EXEMPLARY
+DAMAGES SUCH AS LOST PROFITS OR LOST
+SAVINGS The foregoing exclusions and limitations of
+liability do not apply to damages caused by a party's
+breach of Section 7 (Intellectual Property and
+Confidentiality) Section 9 (Indemnity) gross negligence
+or willful misconduct Without limiting either party's (a)
+responsibility for direct damages or (b) right to claim other
+direct damages, the following damages shall be
+considered as direct damages and shall not be excluded
+from liability under this Agreement (i) reasonable costs
+incurred by Centene to correct the Services/Deliverables
+or acquire substitute services as a result of any uncured
+breach of this Agreement by Vendor and (ii) amounts
+charged by Vendor to provide transition assistance
+pursuant to Section 1.2 (Transition Assistance) for a
+period of up to ninety (90) days in the event of an
+uncured material breach of this Agreement by Vendor
+Further, without limiting a party's right to claim (i) other
+expenses are both reasonable and subject to
+indemnification hereunder and/or (ii) additional amounts
+for identity-Related Services (defined below) are
+reasonable the parties acknowledge and agree that
+expenses of not more than $200 for Identity-Related
+Services per affected individual are reasonable subject
+to indemnity if such expenses are incurred in response to
+an Incident and not precluded by the foregoing limitation
+of liability
+"Identity-Related Services" means
+notification letters credit monitoring services identity
+theft insurance reimbursement for credit freezes fraud
+resolution services identity and credit restoration
+services toll free information services for affected
+individuals and any similar service which corporate
+entities which create or maintain Protected Health
+Information make available to impacted individuals in the
+event of a Breach or alleged Breach regarding such
+information
+The Services will be in support of one or more of the following (check all that apply):
+Exchange (ACA)
+Commercial
+Duals
+Medicare
+TRICARE
+Medicaid (please list states below AND attach the appropriate Medicaid/Regulatory Addendum)
+
+

assets/text/table_check.txt

@@ -0,0 +1,17 @@
+Start of Page No. = 1
+
+EXHIBIT B-2
+COMPENSATION
+Health Insurance Marketplace (HIM)
+Applicable
+Does not participate in Kelsey Marketplace
+Benefit Plan(s):
+Does not participate in ERS
+Provider Type:
+Specialist
+Services:
+Professional Services
+Physician/Provider agrees to participate in the Benefit Plan/Program described in this Exhibit and authorizes,
+through its signature below, the transfer of all payment/reimbursement terms and obligations under the
+
+

assets/text/table_check_row.txt

@@ -0,0 +1,7 @@
+Start of Page No. = 1
+
+-------Table Start--------
+[["Type","Service (Check all that apply)","MMP Reimbursement","DSNI Reimbursement"],["Therapy","50 Occupational Therapy 50 Physical Therapy to Speech Therapy","80%","80%"]]
+-------Table End--------
+
+

assets/textract/helloWorld.gen

@@ -35,14 +35,14 @@
 							}
 						]
 					},
-					"Id": "15d237c3-4f14-41e3-90c5-8e9393f95735",
+					"Id": "538b9a69-e418-4f18-8d05-e00b05158032",
 					"Page": null,
 					"Query": null,
 					"Relationships": [
 						{
 							"Ids": [
-								"b60df4b4-76d7-4b11-bef8-b2f605d19b7e",
-								"355b6ac6-d19f-40ec-954c-2432c773d198"
+								"7fda8da3-faf1-4826-9ef6-485937ed84c6",
+								"d2c76ce9-d421-46a7-9144-35829514042d"
 							],
 							"Type": "CHILD"
 						}
@@ -85,14 +85,14 @@
 							}
 						]
 					},
-					"Id": "b60df4b4-76d7-4b11-bef8-b2f605d19b7e",
+					"Id": "7fda8da3-faf1-4826-9ef6-485937ed84c6",
 					"Page": null,
 					"Query": null,
 					"Relationships": [
 						{
 							"Ids": [
-								"180bd18f-9d4a-4b0a-a986-7028d2555cf2",
-								"f54a60f6-64d4-4434-8dc0-4d14f3d110b4"
+								"0faae859-4c2e-4aaf-a779-ea47eb7d8bba",
+								"48191175-2f21-4992-9da6-df00599af5d8"
 							],
 							"Type": "CHILD"
 						}
@@ -135,7 +135,7 @@
 							}
 						]
 					},
-					"Id": "180bd18f-9d4a-4b0a-a986-7028d2555cf2",
+					"Id": "0faae859-4c2e-4aaf-a779-ea47eb7d8bba",
 					"Page": null,
 					"Query": null,
 					"Relationships": null,
@@ -177,7 +177,7 @@
 							}
 						]
 					},
-					"Id": "f54a60f6-64d4-4434-8dc0-4d14f3d110b4",
+					"Id": "48191175-2f21-4992-9da6-df00599af5d8",
 					"Page": null,
 					"Query": null,
 					"Relationships": null,
@@ -219,13 +219,13 @@
 							}
 						]
 					},
-					"Id": "355b6ac6-d19f-40ec-954c-2432c773d198",
+					"Id": "d2c76ce9-d421-46a7-9144-35829514042d",
 					"Page": null,
 					"Query": null,
 					"Relationships": [
 						{
 							"Ids": [
-								"b60df4b4-76d7-4b11-bef8-b2f605d19b7e"
+								"7fda8da3-faf1-4826-9ef6-485937ed84c6"
 							],
 							"Type": "CHILD"
 						}

cmd/textractGenerator_test/textdetection.go

@@ -20,7 +20,6 @@ import (
 	"time"
 
 	awstextract "github.com/aws/aws-sdk-go-v2/service/textract"
-	"github.com/aws/aws-sdk-go-v2/service/textract/types"
 )
 
 type Config struct {
@@ -88,7 +87,8 @@ func main() {
 	allResults := []map[string]*awstextract.AnalyzeDocumentOutput{}
 	allResults = make([]map[string]*awstextract.AnalyzeDocumentOutput, count)
 	var wg sync.WaitGroup
-	for i := range count {
+	for index := range count {
+		i := index
 		wg.Add(1)
 		go func() {
 			slog.Info("getting page", "index", i)
@@ -103,36 +103,22 @@ func main() {
 				"base": &baseResult,
 			}
 
-			hasTable := false
-			for _, block := range baseResult.Blocks {
-				if block.BlockType == types.BlockTypeLayoutTable {
-					hasTable = true
-				}
+			features, err := textSvc.ListRequiredFeatures(ctx, baseResult.Blocks, i)
+			if err != nil {
+				slog.Error("error listing features", "error", err)
+				os.Exit(1)
 			}
 
-			if hasTable {
-				slog.Info("getting table", "index", i)
-				pageBytes, err := pdf.GetPageAsPNG(ctx, i)
+			if len(features) > 0 {
+				fullFeaturesResult, err := textSvc.GetPageWithFeatures(ctx, pdf, i, features)
 				if err != nil {
-					slog.Error("error getting page", "error", err)
+					slog.Error("error getting features", "error", err)
 					os.Exit(1)
 				}
 
-				tableResult, err := cfg.GetTextractClient().AnalyzeDocument(ctx, &awstextract.AnalyzeDocumentInput{
-					Document: &types.Document{
-						Bytes: pageBytes,
-					},
-					FeatureTypes: []types.FeatureType{
-						types.FeatureTypeTables,
-					},
-				})
-				if err != nil {
-					slog.Error("error detecting table", "error", err)
-					os.Exit(1)
-				}
-
-				allResults[i]["table"] = tableResult
+				allResults[i]["full_features"] = &fullFeaturesResult
 			}
+
 			wg.Done()
 		}()
 	}

internal/document/clean/create.go

@@ -2,6 +2,7 @@ package documentclean
 
 import (
 	"context"
+	"fmt"
 
 	doctextrunner "queryorchestration/api/docTextRunner"
 	"queryorchestration/internal/serviceconfig/queue"
@@ -12,13 +13,13 @@ import (
 func (s *Service) Clean(ctx context.Context, documentId uuid.UUID) error {
 	isclean, err := s.cfg.GetDBQueries().HasDocumentCleanEntry(ctx, documentId)
 	if err != nil {
-		return err
+		return fmt.Errorf("unable to verify if document has been cleaned: %s", err)
 	}
 
 	if !isclean {
 		err = s.clean(ctx, documentId)
 		if err != nil {
-			return err
+			return fmt.Errorf("unable to clean document: %s", err)
 		}
 	}
 

internal/document/text/checkbox.go

@@ -0,0 +1,54 @@
+package documenttext
+
+import (
+	"strings"
+)
+
+func (s *Service) hasCheckboxKeywords(text string) bool {
+	conditions := []func(string) bool{
+		s.hasExhibit,
+		s.hasBenefit,
+		s.hasPlans,
+		s.hasServices,
+	}
+
+	elementsContained := 0
+	for _, hasCondition := range conditions {
+		if hasCondition(text) {
+			elementsContained++
+		}
+	}
+
+	return elementsContained >= 3
+}
+
+func (s *Service) hasExhibit(text string) bool {
+	return strings.Contains(text, "exhibit")
+}
+
+func (s *Service) hasBenefit(text string) bool {
+	keywords := []string{"benefit plan", "applicable benefit", "benefits"}
+
+	return orKeywords(text, keywords)
+}
+
+func (s *Service) hasPlans(text string) bool {
+	keywords := []string{"plan/program:", "plan(s)", "plans"}
+
+	return orKeywords(text, keywords)
+}
+
+func (s *Service) hasServices(text string) bool {
+	keywords := []string{"service:", "services", "service(s)"}
+	return orKeywords(text, keywords)
+}
+
+func orKeywords(text string, keywords []string) bool {
+	for _, keyword := range keywords {
+		if strings.Contains(text, keyword) {
+			return true
+		}
+	}
+
+	return false
+}

internal/document/text/checkbox_test.go

@@ -0,0 +1,57 @@
+package documenttext
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestHasCheckboxKeywords(t *testing.T) {
+	svc := Service{}
+	assert.False(t, svc.hasCheckboxKeywords(""))
+	assert.False(t, svc.hasCheckboxKeywords("aaaaaa"))
+	assert.False(t, svc.hasCheckboxKeywords("benefit plan"))
+	assert.False(t, svc.hasCheckboxKeywords("benefit plans"))
+	assert.True(t, svc.hasCheckboxKeywords("exhibit benefit plans"))
+}
+
+func TestHasExhibit(t *testing.T) {
+	svc := Service{}
+	assert.False(t, svc.hasExhibit(""))
+	assert.False(t, svc.hasExhibit("aaaaaa"))
+	assert.True(t, svc.hasExhibit("exhibit"))
+}
+
+func TestHasBenefit(t *testing.T) {
+	svc := Service{}
+	assert.False(t, svc.hasBenefit(""))
+	assert.False(t, svc.hasBenefit("aaaaaa"))
+	assert.False(t, svc.hasBenefit("benefit"))
+	assert.True(t, svc.hasBenefit("benefit plan"))
+	assert.True(t, svc.hasBenefit("benefit plan(s)"))
+	assert.True(t, svc.hasBenefit("applicable benefit"))
+	assert.True(t, svc.hasBenefit("benefits"))
+}
+
+func TestHasPlans(t *testing.T) {
+	svc := Service{}
+	assert.False(t, svc.hasPlans(""))
+	assert.False(t, svc.hasPlans("aaaaaa"))
+	assert.True(t, svc.hasPlans("plan(s):"))
+	assert.True(t, svc.hasPlans("plan(s)/program(s):"))
+	assert.True(t, svc.hasPlans("plans:"))
+	assert.True(t, svc.hasPlans("plan/program:"))
+	assert.True(t, svc.hasPlans("plan(s)"))
+	assert.True(t, svc.hasPlans("plans"))
+}
+
+func TestHasServices(t *testing.T) {
+	svc := Service{}
+	assert.False(t, svc.hasServices(""))
+	assert.False(t, svc.hasServices("aaaaaa"))
+	assert.True(t, svc.hasServices("services:"))
+	assert.True(t, svc.hasServices("service:"))
+	assert.True(t, svc.hasServices("service(s):"))
+	assert.True(t, svc.hasServices("services"))
+	assert.True(t, svc.hasServices("service(s)"))
+}

internal/document/text/childElements.go

@@ -7,11 +7,11 @@ import (
 	"github.com/google/uuid"
 )
 
-func (s *Service) getDirectChildren(block types.Block, blockMap map[uuid.UUID]types.Block) []uuid.UUID {
+func (s *Service) GetDirectChildren(id uuid.UUID, blockMap map[uuid.UUID]types.Block) []uuid.UUID {
 	lines := []uuid.UUID{}
 	found := map[uuid.UUID]*bool{}
 
-	for _, id := range s.getChildIds(block) {
+	for _, id := range s.getRelationshipIds(blockMap[id]) {
 		if found[id] != nil && !*found[id] {
 			continue
 		}
@@ -20,7 +20,7 @@ func (s *Service) getDirectChildren(block types.Block, blockMap map[uuid.UUID]ty
 		f := true
 		found[id] = &f
 
-		s.removeChildren(blockMap[id], blockMap, &found)
+		s.removeChildren(id, blockMap, &found)
 	}
 
 	output := []uuid.UUID{}
@@ -35,15 +35,15 @@ func (s *Service) getDirectChildren(block types.Block, blockMap map[uuid.UUID]ty
 	return output
 }
 
-func (s *Service) removeChildren(block types.Block, blockMap map[uuid.UUID]types.Block, found *map[uuid.UUID]*bool) {
-	for _, childId := range s.getChildIds(block) {
+func (s *Service) removeChildren(id uuid.UUID, blockMap map[uuid.UUID]types.Block, found *map[uuid.UUID]*bool) {
+	for _, childId := range s.getRelationshipIds(blockMap[id]) {
 		f := false
 		(*found)[childId] = &f
-		s.removeChildren(blockMap[childId], blockMap, found)
+		s.removeChildren(childId, blockMap, found)
 	}
 }
 
-func (s *Service) getBlockMap(blocksList []types.Block) map[uuid.UUID]types.Block {
+func (s *Service) GetBlockMap(blocksList []types.Block) map[uuid.UUID]types.Block {
 	blocks := make(map[uuid.UUID]types.Block, len(blocksList))
 	for _, block := range blocksList {
 		if block.Id == nil {
@@ -61,7 +61,7 @@ func (s *Service) getBlockMap(blocksList []types.Block) map[uuid.UUID]types.Bloc
 	return blocks
 }
 
-func (s *Service) getChildIds(block types.Block) []uuid.UUID {
+func (s *Service) getRelationshipIds(block types.Block) []uuid.UUID {
 	childIds := []uuid.UUID{}
 	for _, relationship := range block.Relationships {
 		for _, idStr := range relationship.Ids {
@@ -77,3 +77,23 @@ func (s *Service) getChildIds(block types.Block) []uuid.UUID {
 
 	return childIds
 }
+
+func (s *Service) getRelationshipIdsByRel(block types.Block, relType types.RelationshipType) []uuid.UUID {
+	childIds := []uuid.UUID{}
+	for _, relationship := range block.Relationships {
+		if relationship.Type != relType {
+			continue
+		}
+		for _, idStr := range relationship.Ids {
+			id, err := uuid.Parse(idStr)
+			if err != nil {
+				slog.Error("invalid child uuid", "id", id, "error", err)
+				continue
+			}
+
+			childIds = append(childIds, id)
+		}
+	}
+
+	return childIds
+}

internal/document/text/documentText_test.go

@@ -66,6 +66,82 @@ func TestGetDocumentText(t *testing.T) {
 		_, err := svc.getDocumentText(ctx, pdf)
 		require.EqualError(t, err, "no textract response")
 	})
+	t.Run("merged_cell_table", func(t *testing.T) {
+		t.Parallel()
+		cfg := DocTextConfig{}
+		svc := Service{
+			cfg: &cfg,
+		}
+		mockTextract := textractmock.NewMockTextractClient(t)
+		cfg.TextractClient = mockTextract
+
+		textractBaseOut, pdf, txtFile, close := getFiles(t, "merged_cell_table")
+		defer close()
+
+		textExpectations(t, ctx, mockTextract, pdf, textractBaseOut)
+
+		text, err := svc.getDocumentText(ctx, pdf)
+		require.NoError(t, err)
+
+		assertReaders(t, txtFile, text)
+	})
+	t.Run("table_check", func(t *testing.T) {
+		t.Parallel()
+		cfg := DocTextConfig{}
+		svc := Service{
+			cfg: &cfg,
+		}
+		mockTextract := textractmock.NewMockTextractClient(t)
+		cfg.TextractClient = mockTextract
+
+		textractBaseOut, pdf, txtFile, close := getFiles(t, "table_check")
+		defer close()
+
+		textExpectations(t, ctx, mockTextract, pdf, textractBaseOut)
+
+		text, err := svc.getDocumentText(ctx, pdf)
+		require.NoError(t, err)
+
+		assertReaders(t, txtFile, text)
+	})
+	t.Run("multi_column", func(t *testing.T) {
+		t.Parallel()
+		cfg := DocTextConfig{}
+		svc := Service{
+			cfg: &cfg,
+		}
+		mockTextract := textractmock.NewMockTextractClient(t)
+		cfg.TextractClient = mockTextract
+
+		textractBaseOut, pdf, txtFile, close := getFiles(t, "multi_column")
+		defer close()
+
+		textExpectations(t, ctx, mockTextract, pdf, textractBaseOut)
+
+		text, err := svc.getDocumentText(ctx, pdf)
+		require.NoError(t, err)
+
+		assertReaders(t, txtFile, text)
+	})
+	t.Run("table_check_row", func(t *testing.T) {
+		t.Parallel()
+		cfg := DocTextConfig{}
+		svc := Service{
+			cfg: &cfg,
+		}
+		mockTextract := textractmock.NewMockTextractClient(t)
+		cfg.TextractClient = mockTextract
+
+		textractBaseOut, pdf, txtFile, close := getFiles(t, "table_check_row")
+		defer close()
+
+		textExpectations(t, ctx, mockTextract, pdf, textractBaseOut)
+
+		text, err := svc.getDocumentText(ctx, pdf)
+		require.NoError(t, err)
+
+		assertReaders(t, txtFile, text)
+	})
 
 	t.Run("cnc_01-06", func(t *testing.T) {
 		t.Parallel()
@@ -141,25 +217,6 @@ func TestGetDocumentText(t *testing.T) {
 		text, err := svc.getDocumentText(ctx, pdf)
 		require.NoError(t, err)
 
-		assertReaders(t, txtFile, text)
-	})
-	t.Run("merged_cell_table", func(t *testing.T) {
-		t.Parallel()
-		cfg := DocTextConfig{}
-		svc := Service{
-			cfg: &cfg,
-		}
-		mockTextract := textractmock.NewMockTextractClient(t)
-		cfg.TextractClient = mockTextract
-
-		textractBaseOut, pdf, txtFile, close := getFiles(t, "merged_cell_table")
-		defer close()
-
-		textExpectations(t, ctx, mockTextract, pdf, textractBaseOut)
-
-		text, err := svc.getDocumentText(ctx, pdf)
-		require.NoError(t, err)
-
 		assertReaders(t, txtFile, text)
 	})
 }
@@ -208,13 +265,20 @@ func textExpectations(t testing.TB, ctx context.Context, mockTextract *textractm
 			).
 			Return(base, nil)
 
-		table := pageResult["table"]
+		table := pageResult["full_features"]
 		if table != nil {
 			mockTextract.EXPECT().
 				AnalyzeDocument(
 					mock.Anything,
 					mock.MatchedBy(func(in *textract.AnalyzeDocumentInput) bool {
-						return string(in.Document.Bytes) == string(page) && in.FeatureTypes[0] == types.FeatureTypeTables
+						found := false
+						for _, feature := range in.FeatureTypes {
+							if feature == types.FeatureTypeForms || feature == types.FeatureTypeTables {
+								found = true
+								break
+							}
+						}
+						return string(in.Document.Bytes) == string(page) && found
 					}),
 					mock.Anything,
 				).
@@ -223,17 +287,21 @@ func textExpectations(t testing.TB, ctx context.Context, mockTextract *textractm
 	}
 }
 
+func assertReaderAndString(t testing.TB, expected io.Reader, actual string) {
+	expectedBytes, err := io.ReadAll(expected)
+	require.NoError(t, err)
+	expectedStr := string(expectedBytes)
+
+	if !assert.Equal(t, expectedStr, actual) {
+		diff := cmp.Diff(expectedStr, actual)
+		t.Logf("Detailed diff (-expected +actual):\n%s", diff)
+	}
+}
+
 func assertReaders(t testing.TB, expected io.Reader, actual io.Reader) {
 	actualBytes, err := io.ReadAll(actual)
 	require.NoError(t, err)
 	actualStr := string(actualBytes)
 
-	expectedBytes, err := io.ReadAll(expected)
-	require.NoError(t, err)
-	expectedStr := string(expectedBytes)
-
-	if !assert.Equal(t, expectedStr, actualStr) {
-		diff := cmp.Diff(expectedStr, actualStr)
-		t.Logf("Detailed diff (-expected +actual):\n%s", diff)
-	}
+	assertReaderAndString(t, expected, actualStr)
 }

internal/document/text/keyvalue.go

@@ -0,0 +1,113 @@
+package documenttext
+
+import (
+	"fmt"
+	"log/slog"
+	"strings"
+
+	"github.com/aws/aws-sdk-go-v2/service/textract/types"
+	"github.com/google/uuid"
+)
+
+func (s *Service) getKeyValueSets(id uuid.UUID, elements PageElements) (map[uuid.UUID]uuid.UUID, error) {
+	block := elements.BlockMap[id]
+	keyValuePairs := map[uuid.UUID]bool{}
+	lines := map[uuid.UUID]bool{}
+	lineToKey := map[uuid.UUID]uuid.UUID{}
+
+	for _, lineId := range s.getRelationshipIds(block) {
+		for _, wordId := range s.getRelationshipIds(elements.BlockMap[lineId]) {
+			for _, b := range elements.BlockMap {
+				for _, childId := range s.getRelationshipIds(b) {
+					if wordId != childId || b.BlockType != types.BlockTypeKeyValueSet {
+						continue
+					}
+
+					id, err := uuid.Parse(*b.Id)
+					if err != nil {
+						return nil, err
+					}
+
+					lines[lineId] = true
+					keyValuePairs[id] = true
+					lineToKey[lineId] = id
+				}
+			}
+		}
+	}
+
+	for _, lineId := range s.getRelationshipIds(block) {
+		if lines[lineId] {
+			continue
+		}
+		for _, wordId := range s.getRelationshipIds(elements.BlockMap[lineId]) {
+			for _, b := range elements.BlockMap {
+				for _, childId := range s.getRelationshipIds(b) {
+					id, err := uuid.Parse(*b.Id)
+					if err != nil {
+						return nil, err
+					}
+
+					if wordId != childId || keyValuePairs[id] {
+						continue
+					}
+					keyValuePairs[id] = true
+					lineToKey[lineId] = id
+				}
+			}
+		}
+	}
+
+	return lineToKey, nil
+}
+
+func (s *Service) addKeyValue(id uuid.UUID, elements PageElements, builder *strings.Builder) error {
+	lineToKey, err := s.getKeyValueSets(id, elements)
+	if err != nil {
+		return err
+	}
+
+	for _, lineId := range s.getRelationshipIds(elements.BlockMap[id]) {
+		keysBlock := lineToKey[lineId]
+		needsSelection := false
+		isSelected := false
+
+		for _, id := range s.getRelationshipIdsByRel(elements.BlockMap[keysBlock], types.RelationshipTypeValue) {
+			for _, id := range s.getRelationshipIdsByRel(elements.BlockMap[id], types.RelationshipTypeChild) {
+				keyBlock := elements.BlockMap[id]
+				if keyBlock.BlockType != types.BlockTypeSelectionElement {
+					continue
+				}
+				needsSelection = true
+				if keyBlock.SelectionStatus == types.SelectionStatusSelected {
+					isSelected = true
+					break
+				}
+			}
+		}
+
+		if !needsSelection {
+			builder.WriteString(fmt.Sprintf("%s\n", *elements.BlockMap[lineId].Text))
+			continue
+		} else if !isSelected {
+			continue
+		}
+
+		var key strings.Builder
+		for _, id := range s.getRelationshipIdsByRel(elements.BlockMap[keysBlock], types.RelationshipTypeChild) {
+			keyBlock := elements.BlockMap[id]
+			if keyBlock.BlockType != types.BlockTypeWord {
+				slog.Info("uncaught element in key value", "type", keyBlock.BlockType)
+				continue
+			}
+			if key.String() != "" {
+				key.WriteRune(' ')
+			}
+			key.WriteString(*keyBlock.Text)
+		}
+
+		builder.WriteString(fmt.Sprintf("%s\n", key.String()))
+	}
+
+	return nil
+}

internal/document/text/pageText.go

@@ -2,6 +2,7 @@ package documenttext
 
 import (
 	"context"
+	"errors"
 	"fmt"
 	"log/slog"
 	"strings"
@@ -19,102 +20,126 @@ func (s *Service) getPageText(ctx context.Context, document documenttypes.File,
 		return "", err
 	}
 
-	var pageBlock *types.Block
-	for _, block := range res.Blocks {
-		if block.BlockType == types.BlockTypePage {
-			pageBlock = &block
-			break
-		}
-	}
-	if pageBlock == nil {
-		return "", fmt.Errorf("no page block found for index: %d", index)
-	}
-
-	blockMap := s.getBlockMap(res.Blocks)
-
-	lines, specialFeatures, err := s.getPageElements(ctx, document, index, *pageBlock, blockMap)
+	features, err := s.ListRequiredFeatures(ctx, res.Blocks, index)
 	if err != nil {
 		return "", err
 	}
 
-	return s.buildPageText(ctx, lines, blockMap, index, specialFeatures)
+	blocks := res.Blocks
+	if len(features) > 0 {
+		res, err := s.GetPageWithFeatures(ctx, document, index, features)
+		if err != nil {
+			return "", err
+		}
+
+		blocks = res.Blocks
+	}
+
+	elements, err := s.getPageElements(ctx, document, index, blocks)
+	if err != nil {
+		return "", err
+	}
+
+	return s.buildPageText(ctx, elements, index)
 }
 
-func (s *Service) getPageElements(ctx context.Context, document documenttypes.File, index int, pageBlock types.Block, blockMap map[uuid.UUID]types.Block) ([]uuid.UUID, specialisedFeatures, error) {
-	directChildren := s.getDirectChildren(pageBlock, blockMap)
+type PageElements struct {
+	Page           uuid.UUID
+	BlockMap       map[uuid.UUID]types.Block
+	DirectChildren []uuid.UUID
+	Tables         []uuid.UUID
+}
 
-	hasTable := false
-	for _, id := range directChildren {
-		if blockMap[id].BlockType == types.BlockTypeLayoutTable {
-			hasTable = true
+func (s *Service) getPageElements(ctx context.Context, document documenttypes.File, index int, blocks []types.Block) (PageElements, error) {
+	elements := PageElements{}
+
+	pageBlock, err := s.GetPageBlock(blocks)
+	if err != nil {
+		return PageElements{}, err
+	}
+
+	elements.Page = pageBlock
+	elements.BlockMap = s.GetBlockMap(blocks)
+	elements.DirectChildren = s.GetDirectChildren(pageBlock, elements.BlockMap)
+
+	for _, block := range blocks {
+		if block.BlockType != types.BlockTypeTable {
+			continue
+		}
+		id, err := uuid.Parse(*block.Id)
+		if err != nil {
+			return PageElements{}, err
+		}
+
+		elements.Tables = append(elements.Tables, id)
+	}
+
+	return elements, nil
+}
+
+func (s *Service) ListRequiredFeatures(ctx context.Context, blocks []types.Block, index int) ([]types.FeatureType, error) {
+	features := []types.FeatureType{}
+	for _, block := range blocks {
+		switch block.BlockType {
+		case types.BlockTypeLayoutTable:
+			features = append(features, types.FeatureTypeTables)
+		case types.BlockTypeLayoutKeyValue:
+			features = append(features, types.FeatureTypeForms)
 		}
 	}
 
-	features := []types.FeatureType{}
-	if hasTable {
-		features = append(features, types.FeatureTypeTables)
-		// features = append(features, types.FeatureTypeForms)
+	if len(features) > 0 {
+		features = append(features, types.FeatureTypeLayout)
+		features = append(features, types.FeatureTypeSignatures)
 	}
 
-	// Detecting Checkboxes -
-	// if has tables
-	// For a page with no tables, we look for 3 out of 4 of the following conditions
-	//     has_exhibit→ 'exhibit' is on the page
-	//     has_benefit→ one of ['benefit plan', 'benefit plan(s)', 'applicable benefit', 'benefits'] is on the page
-	//     has_plans→ one of ['plan(s):', 'plan(s)/program(s):', 'plans:', 'plan/program:', 'plan(s)', 'plans'] is on the page
-	//     has_services → one of ['services:', 'service:', 'service(s):', 'services', 'service(s)'] is on the page
-
-	if len(features) == 0 {
-		return directChildren, specialisedFeatures{}, nil
-	}
-
-	featuresOut, err := s.getSpecialisedFeatures(ctx, document, index, features)
-	if err != nil {
-		return nil, specialisedFeatures{}, err
-	}
-
-	return directChildren, featuresOut, nil
+	return features, nil
 }
 
 type buildParams struct {
-	numOfSignatures int
-	currentTable    int
+	tableIndex int
 }
 
-func (s *Service) buildPageText(ctx context.Context, lines []uuid.UUID, blockMap map[uuid.UUID]types.Block, index int, specialFeatures specialisedFeatures) (string, error) {
+func (s *Service) buildPageText(ctx context.Context, elements PageElements, index int) (string, error) {
 	var builder strings.Builder
-	params := buildParams{
-		numOfSignatures: 0,
-		currentTable:    0,
-	}
+	params := buildParams{}
 
-	for _, id := range lines {
-		err := s.addComponentText(ctx, blockMap[id], blockMap, specialFeatures, &params, &builder)
+	for _, id := range elements.DirectChildren {
+		err := s.addComponentText(ctx, id, elements, &params, &builder)
 		if err != nil {
 			return "", err
 		}
 	}
 
-	for params.currentTable < len(specialFeatures.Tables) {
-		id := specialFeatures.Tables[params.currentTable]
-		err := s.addComponentText(ctx, specialFeatures.BlockMap[id], blockMap, specialFeatures, &params, &builder)
+	for params.tableIndex < len(elements.Tables) {
+		err := s.addTable(ctx, elements.Tables[params.tableIndex], elements.BlockMap, &builder)
 		if err != nil {
 			return "", err
 		}
+
+		params.tableIndex++
 	}
 
-	if params.numOfSignatures > 0 {
-		builder.WriteString(fmt.Sprintf("\nThis page has %d signature.\n", params.numOfSignatures))
+	numOfSignatures := 0
+	for _, block := range elements.BlockMap {
+		if block.BlockType == types.BlockTypeSignature {
+			numOfSignatures++
+		}
+	}
+
+	if numOfSignatures > 0 {
+		builder.WriteString(fmt.Sprintf("\nThis page has %d signature.\n", numOfSignatures))
 	}
 
 	return builder.String(), nil
 }
 
-func (s *Service) addComponentText(ctx context.Context, block types.Block, blockMap map[uuid.UUID]types.Block, specialFeatures specialisedFeatures, params *buildParams, builder *strings.Builder) error {
+func (s *Service) addComponentText(ctx context.Context, id uuid.UUID, elements PageElements, params *buildParams, builder *strings.Builder) error {
+	block := elements.BlockMap[id]
 	switch block.BlockType {
-	case types.BlockTypeLayoutText, types.BlockTypeLayoutTitle, types.BlockTypeLayoutHeader, types.BlockTypeLayoutFooter, types.BlockTypeLayoutSectionHeader, types.BlockTypeLayoutPageNumber, types.BlockTypeLayoutList, types.BlockTypeLayoutFigure, types.BlockTypeLayoutKeyValue:
-		for _, id := range s.getChildIds(block) {
-			err := s.addComponentText(ctx, blockMap[id], blockMap, specialFeatures, params, builder)
+	case types.BlockTypeLayoutText, types.BlockTypeLayoutTitle, types.BlockTypeLayoutHeader, types.BlockTypeLayoutFooter, types.BlockTypeLayoutSectionHeader, types.BlockTypeLayoutPageNumber, types.BlockTypeLayoutList, types.BlockTypeLayoutFigure:
+		for _, id := range s.getRelationshipIds(block) {
+			err := s.addComponentText(ctx, id, elements, params, builder)
 			if err != nil {
 				return err
 			}
@@ -131,23 +156,24 @@ func (s *Service) addComponentText(ctx context.Context, block types.Block, block
 			slog.Error("unable to write string", "error", err)
 			return err
 		}
-	case types.BlockTypeSignature:
-		params.numOfSignatures++
-	case types.BlockTypeLayoutTable, types.BlockTypeTable:
-		if block.BlockType == types.BlockTypeLayoutTable && len(specialFeatures.Tables) <= params.currentTable {
-			for _, child := range s.getChildIds(block) {
-				err := s.addComponentText(ctx, blockMap[child], blockMap, specialFeatures, params, builder)
+	case types.BlockTypeLayoutTable:
+		if params.tableIndex < len(elements.Tables) {
+			err := s.addTable(ctx, elements.Tables[params.tableIndex], elements.BlockMap, builder)
+			if err != nil {
+				return err
+			}
+		} else {
+			for _, child := range s.getRelationshipIds(block) {
+				err := s.addComponentText(ctx, child, elements, params, builder)
 				if err != nil {
 					return err
 				}
 			}
-			params.currentTable++
-			break
 		}
-		tableId := specialFeatures.Tables[params.currentTable]
-		params.currentTable++
-		tableBlock := specialFeatures.BlockMap[tableId]
-		err := s.addTable(ctx, tableBlock, specialFeatures.BlockMap, builder)
+
+		params.tableIndex++
+	case types.BlockTypeLayoutKeyValue:
+		err := s.addKeyValue(id, elements, builder)
 		if err != nil {
 			return err
 		}
@@ -159,6 +185,13 @@ func (s *Service) addComponentText(ctx context.Context, block types.Block, block
 }
 
 func (s *Service) GetBasePage(ctx context.Context, doc documenttypes.File, index int) (textract.AnalyzeDocumentOutput, error) {
+	return s.GetPageWithFeatures(ctx, doc, index, []types.FeatureType{
+		types.FeatureTypeLayout,
+		types.FeatureTypeSignatures,
+	})
+}
+
+func (s *Service) GetPageWithFeatures(ctx context.Context, doc documenttypes.File, index int, features []types.FeatureType) (textract.AnalyzeDocumentOutput, error) {
 	pageBytes, err := doc.GetPageAsPNG(ctx, index)
 	if err != nil {
 		return textract.AnalyzeDocumentOutput{}, err
@@ -168,10 +201,7 @@ func (s *Service) GetBasePage(ctx context.Context, doc documenttypes.File, index
 		Document: &types.Document{
 			Bytes: pageBytes,
 		},
-		FeatureTypes: []types.FeatureType{
-			types.FeatureTypeLayout,
-			types.FeatureTypeSignatures,
-		},
+		FeatureTypes: features,
 	})
 	if err != nil {
 		return textract.AnalyzeDocumentOutput{}, err
@@ -179,3 +209,18 @@ func (s *Service) GetBasePage(ctx context.Context, doc documenttypes.File, index
 
 	return *res, err
 }
+
+func (s *Service) GetPageBlock(blocks []types.Block) (uuid.UUID, error) {
+	var pageBlock *types.Block
+	for _, block := range blocks {
+		if block.BlockType == types.BlockTypePage {
+			pageBlock = &block
+			break
+		}
+	}
+	if pageBlock == nil {
+		return uuid.Nil, errors.New("no page block found")
+	}
+
+	return uuid.Parse(*pageBlock.Id)
+}

internal/document/text/pageText_test.go

@@ -51,14 +51,14 @@ func TestGetBlockMap(t *testing.T) {
 	svc := Service{}
 	t.Run("0 blocks", func(t *testing.T) {
 		blocks := []types.Block{}
-		outBlocks := svc.getBlockMap(blocks)
+		outBlocks := svc.GetBlockMap(blocks)
 		assert.Equal(t, map[uuid.UUID]types.Block{}, outBlocks)
 	})
 	t.Run("nil id blocks", func(t *testing.T) {
 		blocks := []types.Block{
 			{},
 		}
-		outBlocks := svc.getBlockMap(blocks)
+		outBlocks := svc.GetBlockMap(blocks)
 		assert.Equal(t, map[uuid.UUID]types.Block{}, outBlocks)
 	})
 	t.Run("1 block", func(t *testing.T) {
@@ -69,7 +69,7 @@ func TestGetBlockMap(t *testing.T) {
 				Id: &blockIdStr,
 			},
 		}
-		outBlocks := svc.getBlockMap(blocks)
+		outBlocks := svc.GetBlockMap(blocks)
 		assert.Equal(t, map[uuid.UUID]types.Block{
 			blockId: {
 				Id: &blockIdStr,
@@ -83,7 +83,7 @@ func TestGetBlockMap(t *testing.T) {
 				Id: &blockIdStr,
 			},
 		}
-		outBlocks := svc.getBlockMap(blocks)
+		outBlocks := svc.GetBlockMap(blocks)
 		assert.Equal(t, map[uuid.UUID]types.Block{}, outBlocks)
 	})
 	t.Run("2 blocks", func(t *testing.T) {
@@ -99,7 +99,7 @@ func TestGetBlockMap(t *testing.T) {
 				Id: &blockTwoIdStr,
 			},
 		}
-		outBlocks := svc.getBlockMap(blocks)
+		outBlocks := svc.GetBlockMap(blocks)
 		assert.Equal(t, map[uuid.UUID]types.Block{
 			blockOneId: {
 				Id: &blockOneIdStr,

internal/document/text/tables.go

@@ -6,30 +6,33 @@ import (
 	"log/slog"
 	"strings"
 
-	documenttypes "queryorchestration/internal/document/types"
-
-	"github.com/aws/aws-sdk-go-v2/service/textract"
 	"github.com/aws/aws-sdk-go-v2/service/textract/types"
 	"github.com/google/uuid"
 )
 
 type cell struct {
-	value  string
-	row    int
-	column int
+	value      string
+	row        int
+	column     int
+	isSelected bool
 }
 
-func (s *Service) addTable(ctx context.Context, block types.Block, blockMap map[uuid.UUID]types.Block, builder *strings.Builder) error {
+func (s *Service) addTable(ctx context.Context, id uuid.UUID, blockMap map[uuid.UUID]types.Block, builder *strings.Builder) error {
 	params := tableParams{
 		maxRow:    0,
 		maxColumn: 0,
 		table:     []cell{},
 	}
 
-	directChildren := s.getDirectChildren(block, blockMap)
+	for _, uid := range s.GetDirectChildren(id, blockMap) {
+		s.addComponentTable(ctx, uid, blockMap, &params)
+	}
 
-	for _, uid := range directChildren {
-		s.addComponentTable(ctx, blockMap[uid], blockMap, &params)
+	invalidRows := make([]bool, params.maxRow+1)
+	for _, cell := range params.table {
+		if !cell.isSelected && cell.column == 0 {
+			invalidRows[cell.row] = true
+		}
 	}
 
 	tableArr := make([][]string, params.maxRow+1)
@@ -41,7 +44,14 @@ func (s *Service) addTable(ctx context.Context, block types.Block, blockMap map[
 		tableArr[cell.row][cell.column] = cell.value
 	}
 
-	jsonData, err := json.Marshal(tableArr)
+	finalTable := [][]string{}
+	for i, row := range tableArr {
+		if !invalidRows[i] {
+			finalTable = append(finalTable, row)
+		}
+	}
+
+	jsonData, err := json.Marshal(finalTable)
 	if err != nil {
 		return err
 	}
@@ -59,16 +69,19 @@ type tableParams struct {
 	table     []cell
 }
 
-func (s *Service) addComponentTable(ctx context.Context, block types.Block, blockMap map[uuid.UUID]types.Block, params *tableParams) {
+func (s *Service) addComponentTable(ctx context.Context, id uuid.UUID, blockMap map[uuid.UUID]types.Block, params *tableParams) {
+	block := blockMap[id]
+
 	switch block.BlockType {
 	case types.BlockTypeMergedCell:
 		content := ""
-		for _, uid := range s.getChildIds(block) {
+		addRow := true
+		for _, uid := range s.getRelationshipIds(block) {
 			if blockMap[uid].BlockType != types.BlockTypeCell {
 				slog.Warn("not a cell", "block_type", blockMap[uid].BlockType)
 				continue
 			}
-			content = s.getCellContent(blockMap[uid], blockMap)
+			content, addRow = s.getCellContent(uid, blockMap)
 			if content != "" {
 				break
 			}
@@ -83,16 +96,18 @@ func (s *Service) addComponentTable(ctx context.Context, block types.Block, bloc
 					row,
 					column,
 					content,
+					addRow,
 					params,
 				)
 			}
 		}
 	case types.BlockTypeCell:
-		content := s.getCellContent(block, blockMap)
+		content, addRow := s.getCellContent(id, blockMap)
 		s.addCellToTable(
 			int(*block.RowIndex)-1,
 			int(*block.ColumnIndex)-1,
 			content,
+			addRow,
 			params,
 		)
 	default:
@@ -100,21 +115,24 @@ func (s *Service) addComponentTable(ctx context.Context, block types.Block, bloc
 	}
 }
 
-func (s *Service) getCellContent(block types.Block, blockMap map[uuid.UUID]types.Block) string {
+func (s *Service) getCellContent(id uuid.UUID, blockMap map[uuid.UUID]types.Block) (string, bool) {
 	var str strings.Builder
-	for _, uid := range s.getChildIds(block) {
+	isSelected := true
+	for _, uid := range s.getRelationshipIds(blockMap[id]) {
 		if blockMap[uid].BlockType == types.BlockTypeWord {
 			if str.String() != "" {
 				str.WriteRune(' ')
 			}
 			str.WriteString(*blockMap[uid].Text)
+		} else if blockMap[uid].BlockType == types.BlockTypeSelectionElement && blockMap[uid].SelectionStatus == types.SelectionStatusNotSelected {
+			isSelected = false
 		}
 	}
 
-	return str.String()
+	return str.String(), isSelected
 }
 
-func (s *Service) addCellToTable(rowIndex int, columnIndex int, content string, params *tableParams) {
+func (s *Service) addCellToTable(rowIndex int, columnIndex int, content string, isSelected bool, params *tableParams) {
 	if params.maxRow < rowIndex {
 		params.maxRow = rowIndex
 	}
@@ -122,49 +140,9 @@ func (s *Service) addCellToTable(rowIndex int, columnIndex int, content string,
 		params.maxColumn = columnIndex
 	}
 	params.table = append(params.table, cell{
-		column: columnIndex,
-		row:    rowIndex,
-		value:  content,
+		column:     columnIndex,
+		row:        rowIndex,
+		value:      content,
+		isSelected: isSelected,
 	})
 }
-
-type specialisedFeatures struct {
-	BlockMap map[uuid.UUID]types.Block
-	Tables   []uuid.UUID
-}
-
-func (s *Service) getSpecialisedFeatures(ctx context.Context, document documenttypes.File, index int, features []types.FeatureType) (specialisedFeatures, error) {
-	pageBytes, err := document.GetPageAsPNG(ctx, index)
-	if err != nil {
-		return specialisedFeatures{}, err
-	}
-
-	res, err := s.cfg.GetTextractClient().AnalyzeDocument(ctx, &textract.AnalyzeDocumentInput{
-		Document: &types.Document{
-			Bytes: pageBytes,
-		},
-		FeatureTypes: features,
-	})
-	if err != nil {
-		return specialisedFeatures{}, err
-	}
-
-	tables := []uuid.UUID{}
-	for _, block := range res.Blocks {
-		if block.BlockType == types.BlockTypeTable {
-			id, err := uuid.Parse(*block.Id)
-			if err != nil {
-				slog.Warn("error parsing block id for tables", "error", err, "id", *block.Id)
-				continue
-			}
-
-			tables = append(tables, id)
-		}
-	}
-	blocks := s.getBlockMap(res.Blocks)
-
-	return specialisedFeatures{
-		BlockMap: blocks,
-		Tables:   tables,
-	}, nil
-}

internal/query/result/sync/trigger_test.go

@@ -5,7 +5,6 @@ import (
 	"fmt"
 	"testing"
 
-	"queryorchestration/internal/database/repository"
 	"queryorchestration/internal/serviceconfig"
 	queryc "queryorchestration/internal/serviceconfig/queue/query"
 	queuemock "queryorchestration/mocks/queue"
@@ -14,7 +13,6 @@ import (
 
 	"github.com/aws/aws-sdk-go-v2/service/sqs"
 	"github.com/google/uuid"
-	"github.com/pashagolub/pgxmock/v3"
 	"github.com/stretchr/testify/mock"
 )
 
@@ -26,11 +24,7 @@ type ResultSyncConfig struct {
 func TestTriggerSync(t *testing.T) {
 	ctx := context.Background()
 
-	pool, err := pgxmock.NewPool()
-	require.NoError(t, err)
 	cfg := &ResultSyncConfig{}
-	cfg.DBPool = pool
-	cfg.DBQueries = repository.New(pool)
 	mockSQS := queuemock.NewMockSQSClient(t)
 	cfg.QueueClient = mockSQS
 	cfg.QueryURL = "/i/am/here"
@@ -54,18 +48,14 @@ func TestTriggerSync(t *testing.T) {
 		).
 		Return(&sqs.SendMessageOutput{}, nil)
 
-	err = svc.TriggerSync(ctx, params)
+	err := svc.TriggerSync(ctx, params)
 	require.NoError(t, err)
 }
 
 func TestTriggerMultiSync(t *testing.T) {
 	ctx := context.Background()
 
-	pool, err := pgxmock.NewPool()
-	require.NoError(t, err)
 	cfg := &ResultSyncConfig{}
-	cfg.DBPool = pool
-	cfg.DBQueries = repository.New(pool)
 	mockSQS := queuemock.NewMockSQSClient(t)
 	cfg.QueueClient = mockSQS
 	cfg.QueryURL = "/i/am/here"
@@ -98,6 +88,6 @@ func TestTriggerMultiSync(t *testing.T) {
 		).
 		Return(&sqs.SendMessageOutput{}, nil)
 
-	err = svc.TriggerMultiSync(ctx, docId, []*uuid.UUID{&queryIDOne, &queryIDTwo})
+	err := svc.TriggerMultiSync(ctx, docId, []*uuid.UUID{&queryIDOne, &queryIDTwo})
 	require.NoError(t, err)
 }

internal/query/sync/service_test.go

@@ -1,29 +0,0 @@
-package querysync_test
-
-import (
-	"testing"
-
-	"queryorchestration/internal/database/repository"
-	querysync "queryorchestration/internal/query/sync"
-	"queryorchestration/internal/serviceconfig"
-	"queryorchestration/internal/serviceconfig/queue/query"
-
-	"github.com/pashagolub/pgxmock/v3"
-	"github.com/stretchr/testify/assert"
-	"github.com/stretchr/testify/require"
-)
-
-type QuerySyncConfig struct {
-	serviceconfig.BaseConfig
-	query.QueryConfig
-}
-
-func TestService(t *testing.T) {
-	pool, err := pgxmock.NewPool()
-	require.NoError(t, err)
-	cfg := &QuerySyncConfig{}
-	cfg.DBPool = pool
-	cfg.DBQueries = repository.New(pool)
-	svc := querysync.New(cfg, &querysync.Services{})
-	assert.NotNil(t, svc)
-}

internal/query/types/contextFull/creator_test.go

@@ -1,33 +1,21 @@
 package contextfull_test
 
 import (
-	"context"
 	"testing"
 
-	"queryorchestration/internal/database/repository"
 	resultprocessor "queryorchestration/internal/query/result/processor"
 	contextfull "queryorchestration/internal/query/types/contextFull"
-	"queryorchestration/internal/serviceconfig"
 
-	"github.com/pashagolub/pgxmock/v3"
 	"github.com/stretchr/testify/require"
 )
 
 func TestCreatorValidate(t *testing.T) {
-	ctx := context.Background()
-
-	pool, err := pgxmock.NewPool()
-	require.NoError(t, err)
-	cfg := &serviceconfig.BaseConfig{}
-	cfg.DBPool = pool
-	cfg.DBQueries = repository.New(pool)
-
 	svc := contextfull.NewCreator()
 
 	entity := &resultprocessor.Create{
 		Type: resultprocessor.TypeContextFull,
 	}
 
-	err = svc.Validate(ctx, entity)
+	err := svc.Validate(t.Context(), entity)
 	require.NoError(t, err)
 }

internal/query/types/contextFull/updator_test.go

@@ -1,29 +1,16 @@
 package contextfull_test
 
 import (
-	"context"
 	"testing"
 
-	"queryorchestration/internal/database/repository"
 	resultprocessor "queryorchestration/internal/query/result/processor"
 	contextfull "queryorchestration/internal/query/types/contextFull"
-	"queryorchestration/internal/serviceconfig"
 
 	"github.com/google/uuid"
 	"github.com/stretchr/testify/require"
-
-	"github.com/pashagolub/pgxmock/v3"
 )
 
 func TestUpdatorValidate(t *testing.T) {
-	ctx := context.Background()
-
-	pool, err := pgxmock.NewPool()
-	require.NoError(t, err)
-	cfg := &serviceconfig.BaseConfig{}
-	cfg.DBPool = pool
-	cfg.DBQueries = repository.New(pool)
-
 	svc := contextfull.NewUpdator()
 
 	current := &resultprocessor.Query{
@@ -36,6 +23,6 @@ func TestUpdatorValidate(t *testing.T) {
 		ID: current.ID,
 	}
 
-	err = svc.Validate(ctx, current, entity)
+	err := svc.Validate(t.Context(), current, entity)
 	require.NoError(t, err)
 }

internal/query/types/jsonExtractor/creator_test.go

@@ -1,28 +1,15 @@
 package jsonextractor_test
 
 import (
-	"context"
 	"testing"
 
-	"queryorchestration/internal/database/repository"
 	resultprocessor "queryorchestration/internal/query/result/processor"
 	jsonextractor "queryorchestration/internal/query/types/jsonExtractor"
-	"queryorchestration/internal/serviceconfig"
 
 	"github.com/stretchr/testify/require"
-
-	"github.com/pashagolub/pgxmock/v3"
 )
 
 func TestCreatorValidate(t *testing.T) {
-	ctx := context.Background()
-
-	pool, err := pgxmock.NewPool()
-	require.NoError(t, err)
-	cfg := &serviceconfig.BaseConfig{}
-	cfg.DBPool = pool
-	cfg.DBQueries = repository.New(pool)
-
 	svc := jsonextractor.NewCreator()
 
 	ccfg := "{}"
@@ -31,6 +18,6 @@ func TestCreatorValidate(t *testing.T) {
 		Config: &ccfg,
 	}
 
-	err = svc.Validate(ctx, entity)
+	err := svc.Validate(t.Context(), entity)
 	require.NoError(t, err)
 }

internal/query/types/jsonExtractor/process_test.go

@@ -168,11 +168,7 @@ func TestJSONProcessJSON(t *testing.T) {
 func TestJSONProcessResults(t *testing.T) {
 	ctx := context.Background()
 
-	pool, err := pgxmock.NewPool()
-	require.NoError(t, err)
 	cfg := &serviceconfig.BaseConfig{}
-	cfg.DBPool = pool
-	cfg.DBQueries = repository.New(pool)
 
 	extractor := jsonextractor.NewExtractor(cfg)
 

internal/query/types/jsonExtractor/updator_test.go

@@ -1,28 +1,16 @@
 package jsonextractor_test
 
 import (
-	"context"
 	"testing"
 
-	"queryorchestration/internal/database/repository"
 	resultprocessor "queryorchestration/internal/query/result/processor"
 	jsonextractor "queryorchestration/internal/query/types/jsonExtractor"
-	"queryorchestration/internal/serviceconfig"
 
 	"github.com/google/uuid"
-	"github.com/pashagolub/pgxmock/v3"
 	"github.com/stretchr/testify/require"
 )
 
 func TestUpdatorValidate(t *testing.T) {
-	ctx := context.Background()
-
-	pool, err := pgxmock.NewPool()
-	require.NoError(t, err)
-	cfg := &serviceconfig.BaseConfig{}
-	cfg.DBPool = pool
-	cfg.DBQueries = repository.New(pool)
-
 	svc := jsonextractor.NewUpdator()
 
 	current := &resultprocessor.Query{
@@ -35,6 +23,6 @@ func TestUpdatorValidate(t *testing.T) {
 		ID: current.ID,
 	}
 
-	err = svc.Validate(ctx, current, entity)
+	err := svc.Validate(t.Context(), current, entity)
 	require.NoError(t, err)
 }

internal/query/update/updateprivate_test.go

@@ -29,11 +29,7 @@ type QueryUpdateConfig struct {
 }
 
 func TestGetUpdator(t *testing.T) {
-	pool, err := pgxmock.NewPool()
-	require.NoError(t, err)
 	cfg := &QueryUpdateConfig{}
-	cfg.DBPool = pool
-	cfg.DBQueries = repository.New(pool)
 	svc := New(cfg, &Services{
 		Query: query.New(cfg),
 	})

internal/serviceconfig/database/transaction_test.go

@@ -2,40 +2,58 @@ package database_test
 
 import (
 	"context"
+	"errors"
 	"testing"
 
 	"queryorchestration/internal/database/repository"
 	"queryorchestration/internal/serviceconfig"
+	"queryorchestration/internal/test"
 
-	"github.com/pashagolub/pgxmock/v3"
+	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 )
 
 func TestExecuteTransaction(t *testing.T) {
-	ctx := context.Background()
+	t.Run("create client", func(t *testing.T) {
+		t.Parallel()
+		cfg := &serviceconfig.BaseConfig{}
+		net := test.GetNetwork(t)
+		test.CreateDB(t, cfg, net, &test.CreateDatabaseConfig{})
 
-	pool, err := pgxmock.NewPool()
-	require.NoError(t, err)
-	cfg := &serviceconfig.BaseConfig{}
-	cfg.DBPool = pool
-	cfg.DBQueries = repository.New(pool)
+		err := cfg.ExecuteDBTransaction(t.Context(), func(ctx context.Context, q *repository.Queries) error {
+			err := q.CreateClient(ctx, &repository.CreateClientParams{
+				Name:     "example_client",
+				Clientid: "ID",
+			})
+			require.NoError(t, err)
 
-	clientID := "ID"
-	clientName := "example_client"
-
-	pool.ExpectBegin()
-	pool.ExpectExec("name: CreateClient :exec").WithArgs(clientID, clientName).
-		WillReturnResult(pgxmock.NewResult("", 1))
-	pool.ExpectCommit()
-
-	err = cfg.ExecuteDBTransaction(ctx, func(ctx context.Context, q *repository.Queries) error {
-		err := q.CreateClient(ctx, &repository.CreateClientParams{
-			Name:     clientName,
-			Clientid: clientID,
+			return nil
 		})
 		require.NoError(t, err)
 
-		return nil
+		client, err := cfg.GetDBQueries().GetClient(t.Context(), "ID")
+		require.NoError(t, err)
+		assert.Equal(t, "example_client", client.Name)
+		assert.Equal(t, "ID", client.Clientid)
+	})
+	t.Run("error in transaction", func(t *testing.T) {
+		t.Parallel()
+		cfg := &serviceconfig.BaseConfig{}
+		net := test.GetNetwork(t)
+		test.CreateDB(t, cfg, net, &test.CreateDatabaseConfig{})
+
+		err := cfg.ExecuteDBTransaction(t.Context(), func(ctx context.Context, q *repository.Queries) error {
+			err := q.CreateClient(ctx, &repository.CreateClientParams{
+				Name:     "example_client",
+				Clientid: "ID",
+			})
+			require.NoError(t, err)
+
+			return errors.New("error")
+		})
+		require.Error(t, err)
+
+		_, err = cfg.GetDBQueries().GetClient(t.Context(), "ID")
+		require.Error(t, err)
 	})
-	require.NoError(t, err)
 }

internal/serviceconfig/threadpool/config_test.go

@@ -35,9 +35,11 @@ func TestStopAndWait(t *testing.T) {
 	assert.LessOrEqual(t, 10, cfg.MaxWorkers)
 	assert.NotNil(t, cfg.pool)
 
+	pool := cfg.GetThreadPool()
+
 	for i := range 2 {
 		index := i
-		cfg.pool.Submit(func() {
+		pool.Submit(func() {
 			sl.Info("new thread", "index", index)
 		})
 	}

internal/test/database.go

@@ -23,7 +23,7 @@ const (
 	dbPort  = 5432
 )
 
-func CreateDB(t testing.TB, cfg serviceconfig.ConfigProvider, network string, dcfg *CreateDatabaseConfig) testcontainers.Container {
+func CreateDB(t testing.TB, cfg serviceconfig.ConfigProvider, network string, dcfg *CreateDatabaseConfig) {
 	port, err := nat.NewPort("tcp", strconv.Itoa(dbPort))
 	require.NoError(t, err)
 
@@ -78,6 +78,4 @@ func CreateDB(t testing.TB, cfg serviceconfig.ConfigProvider, network string, dc
 		err = cfg.SetDBPool(t.Context())
 		require.NoError(t, err)
 	}
-
-	return container
 }