Merged in feature/queuetest (pull request #14)

Add tests for queuing

* startingtest

* queuerunnertest

* addedtestfuncs

* testtests

* passintegration

* unittests

api/queue/queryRunner_test.go

@@ -0,0 +1,84 @@
+package controllers_test
+
+import (
+	"context"
+	"encoding/json"
+	controllers "queryorchestration/api/queue"
+	"queryorchestration/internal/database"
+	"queryorchestration/internal/database/repository"
+	"queryorchestration/internal/document"
+	"queryorchestration/internal/queue"
+	"queryorchestration/internal/test"
+	"testing"
+
+	"github.com/aws/aws-sdk-go-v2/service/sqs/types"
+	"github.com/go-playground/validator/v10"
+	"github.com/google/uuid"
+	"github.com/jackc/pgx/v5/pgtype"
+	"github.com/pashagolub/pgxmock/v3"
+	"github.com/stretchr/testify/assert"
+)
+
+func TestQueryRunner(t *testing.T) {
+	ctx := context.Background()
+
+	qCfg, cleanup := test.CreateQueue(t, ctx, &test.CreateQueueConfig{})
+	defer cleanup()
+
+	pool, err := pgxmock.NewPool()
+	if err != nil {
+		t.Fatalf("failed to open pgxmock database: %v", err)
+	}
+	queries := repository.New(pool)
+	db := &database.Connection{
+		Queries: queries,
+		Pool:    pool,
+	}
+
+	svc := document.New(db)
+
+	runner := controllers.NewQueryRunner(svc, validator.New())
+	assert.NotNil(t, runner)
+
+	doc := document.Document{
+		ID:           uuid.New(),
+		JobID:        uuid.New(),
+		Name:         "document_name",
+		CleanVersion: 1,
+		TextVersion:  1,
+	}
+	bodyBytes, err := json.Marshal(doc)
+	assert.Nil(t, err)
+	body := string(bodyBytes)
+	msg := &types.Message{
+		Body: &body,
+	}
+	collectorID := database.MustToDBUUID(uuid.New())
+	queryID := database.MustToDBUUID(uuid.New())
+	minCleanVersion := int32(1)
+	minTextVersion := int32(1)
+
+	cfg := &queue.Config{
+		URL:    qCfg.URL,
+		Client: qCfg.Client,
+	}
+
+	pool.ExpectQuery("name: GetCollectorFromJobID :one").WithArgs(database.MustToDBUUID(doc.JobID)).
+		WillReturnRows(
+			pgxmock.NewRows([]string{"id", "jobId", "minCleanVersion", "minTextVersion"}).
+				AddRow(collectorID, database.MustToDBUUID(doc.JobID), minCleanVersion, minTextVersion),
+		)
+	pool.ExpectQuery("name: ListResultsByDocumentID :many").WithArgs(database.MustToDBUUID(doc.ID), minCleanVersion, minTextVersion).
+		WillReturnRows(
+			pgxmock.NewRows([]string{"id", "queryId", "queryVersion"}).
+				AddRow(collectorID, queryID, int32(1)),
+		)
+	pool.ExpectQuery("name: GetCollectorQueries :many").WithArgs(collectorID).
+		WillReturnRows(
+			pgxmock.NewRows([]string{"collectorId", "queryId", "type", "queryVersion", "requiredIds"}).
+				AddRow(collectorID, queryID, repository.NullQuerytype{Querytype: repository.QuerytypeContextFull, Valid: true}, pgtype.Int4{Int32: int32(1), Valid: true}, []pgtype.UUID{}),
+		)
+
+	err = runner.Process(ctx, cfg, msg)
+	assert.Nil(t, err)
+}

cmd/queryRunner/main.go

@@ -6,6 +6,8 @@ import (
 	"queryorchestration/internal/document"
 	"queryorchestration/internal/queue"
 	"queryorchestration/internal/server"
+
+	_ "github.com/lib/pq"
 )
 
 func main() {

internal/api/listener_test.go

@@ -2,23 +2,23 @@ package api
 
 import (
 	"context"
-	"fmt"
 	"queryorchestration/internal/database"
 	"queryorchestration/internal/server"
+	"queryorchestration/internal/test"
 	"testing"
 
-	"github.com/docker/go-connections/nat"
-	"github.com/jackc/pgx/v5/pgxpool"
 	"github.com/stretchr/testify/assert"
-	"github.com/testcontainers/testcontainers-go"
-	"github.com/testcontainers/testcontainers-go/wait"
 	"google.golang.org/grpc"
 )
 
 func TestNew(t *testing.T) {
 	ctx := context.Background()
 
-	_, cleanup := createDB(t, ctx)
+	_, cleanup := test.CreateDB(t, ctx, &test.CreateDatabaseConfig{
+		Migrations: &database.MigrationConfig{
+			BasePath: "../..",
+		},
+	})
 	defer cleanup()
 
 	getControllers := func(c *server.Config) *grpc.Server {
@@ -41,66 +41,3 @@ func TestListen(t *testing.T) {
 
 	assert.Panics(t, func() { server.Listen() })
 }
-
-type db struct {
-	pool      *pgxpool.Pool
-	container *testcontainers.Container
-}
-
-func createDB(t *testing.T, ctx context.Context) (*db, func()) {
-	port, err := nat.NewPort("tcp", "5432")
-	if err != nil {
-		t.Fatalf("Failed to create port: %v", err)
-	}
-
-	name := "queryorchestration"
-	pass := "pass"
-	user := "postgres"
-
-	req := testcontainers.ContainerRequest{
-		Image: "postgres:17.2-alpine3.21",
-		Env: map[string]string{
-			"POSTGRES_DB":       name,
-			"POSTGRES_USER":     user,
-			"POSTGRES_PASSWORD": pass,
-		},
-		ExposedPorts: []string{port.Port()},
-		WaitingFor:   wait.ForListeningPort(port),
-	}
-
-	container, err := testcontainers.GenericContainer(ctx, testcontainers.GenericContainerRequest{
-		ContainerRequest: req,
-		Started:          true,
-	})
-	if err != nil {
-		t.Fatalf("Failed to start container: %v", err)
-	}
-
-	host, err := container.Host(ctx)
-	if err != nil {
-		t.Fatalf("Failed to extract host: %v", err)
-	}
-	mappedPort, err := container.MappedPort(ctx, port)
-	if err != nil {
-		t.Fatalf("Failed to extract port: %v", err)
-	}
-
-	t.Setenv("DB_USER", user)
-	t.Setenv("DB_PASS", pass)
-	t.Setenv("DB_HOST", host)
-	t.Setenv("DB_PORT", fmt.Sprint(mappedPort.Int()))
-	t.Setenv("DB_NAME", name)
-	t.Setenv("DB_NOSSL", "1")
-
-	pool := database.GetDBPool(ctx)
-
-	return &db{
-			pool:      pool,
-			container: &container,
-		}, func() {
-			err := container.Terminate(ctx)
-			if err != nil {
-				t.Error(err)
-			}
-		}
-}

internal/database/connection.go

@@ -16,7 +16,7 @@ type config struct {
 	host       string
 	port       string
 	name       string
-	disableSSL string
+	disableSSL bool
 }
 
 func mustGetConfig() *config {
@@ -34,7 +34,7 @@ func mustGetConfig() *config {
 		host:       host,
 		port:       port,
 		name:       name,
-		disableSSL: disableSSL,
+		disableSSL: disableSSL == "1",
 	}
 }
 
@@ -43,7 +43,7 @@ func mustGetConnectionString() string {
 	conf := mustGetConfig()
 	opts := ""
 
-	if conf.disableSSL == "1" {
+	if conf.disableSSL {
 		opts += "sslmode=disable"
 	}
 

internal/database/connection_test.go

@@ -2,23 +2,27 @@ package database_test
 
 import (
 	"context"
-	"fmt"
 	"queryorchestration/internal/database"
+	"queryorchestration/internal/test"
 	"testing"
 
-	"github.com/docker/go-connections/nat"
 	"github.com/stretchr/testify/assert"
-	"github.com/testcontainers/testcontainers-go"
-	"github.com/testcontainers/testcontainers-go/wait"
 )
 
 func TestDBConn(t *testing.T) {
 	ctx := context.Background()
-	_, cleanup := createDB(t, ctx)
+	_, cleanup := test.CreateDB(t, ctx, &test.CreateDatabaseConfig{
+		Migrations: &database.MigrationConfig{
+			BasePath: "../..",
+		},
+	})
 	defer cleanup()
 
 	conn := database.GetDBConn(ctx)
 	assert.NotNil(t, conn)
+
+	pool := database.GetDBPool(ctx)
+	assert.NotNil(t, pool)
 }
 
 func TestDBConnNoDB(t *testing.T) {
@@ -31,89 +35,3 @@ func TestDBConnNoDB(t *testing.T) {
 
 	assert.Panics(t, func() { database.GetDBConn(ctx) })
 }
-
-func TestDBPool(t *testing.T) {
-	ctx := context.Background()
-	_, cleanup := createDB(t, ctx)
-	defer cleanup()
-
-	pool := database.GetDBPool(ctx)
-	assert.NotNil(t, pool)
-}
-
-type dbConfig struct {
-	Name     string
-	Host     string
-	Port     int
-	User     string
-	Password string
-}
-
-type db struct {
-	config    *dbConfig
-	container *testcontainers.Container
-}
-
-func createDB(t *testing.T, ctx context.Context) (*db, func()) {
-	alias := "postgres"
-	port, err := nat.NewPort("tcp", "5432")
-	if err != nil {
-		t.Fatalf("Failed to create port: %v", err)
-	}
-
-	config := dbConfig{
-		Name:     "queryorchestration",
-		Password: "pass",
-		User:     "postgres",
-		Port:     port.Int(),
-		Host:     alias,
-	}
-
-	req := testcontainers.ContainerRequest{
-		Image: "postgres:17.2-alpine3.21",
-		Env: map[string]string{
-			"POSTGRES_DB":       config.Name,
-			"POSTGRES_USER":     config.User,
-			"POSTGRES_PASSWORD": config.Password,
-		},
-		ExposedPorts: []string{port.Port()},
-		WaitingFor:   wait.ForListeningPort(port),
-	}
-
-	container, err := testcontainers.GenericContainer(ctx, testcontainers.GenericContainerRequest{
-		ContainerRequest: req,
-		Started:          true,
-	})
-	if err != nil {
-		t.Fatalf("Failed to start container: %v", err)
-	}
-
-	mappedHost, err := container.Host(ctx)
-	if err != nil {
-		t.Fatalf("Failed to map host: %v", err)
-	}
-
-	mappedPort, err := container.MappedPort(ctx, port)
-	if err != nil {
-		t.Fatalf("Failed to map port: %v", err)
-	}
-
-	config.Host = mappedHost
-	config.Port = mappedPort.Int()
-
-	t.Setenv("DB_USER", config.User)
-	t.Setenv("DB_PASS", config.Password)
-	t.Setenv("DB_HOST", config.Host)
-	t.Setenv("DB_PORT", fmt.Sprint(config.Port))
-	t.Setenv("DB_NAME", config.Name)
-
-	return &db{
-			config:    &config,
-			container: &container,
-		}, func() {
-			err := container.Terminate(ctx)
-			if err != nil {
-				t.Error(err)
-			}
-		}
-}

internal/database/migrations.go

@@ -1,6 +1,7 @@
 package database
 
 import (
+	"context"
 	"database/sql"
 	"fmt"
 	"log"
@@ -10,35 +11,37 @@ import (
 	"github.com/golang-migrate/migrate/v4"
 	"github.com/golang-migrate/migrate/v4/database/postgres"
 	_ "github.com/golang-migrate/migrate/v4/source/file"
+	"github.com/jackc/pgx/v5"
 )
 
-func createDB() {
+func createDB(ctx context.Context) {
 	driver := "postgres"
 	opts := ""
 	conf := mustGetConfig()
-	if conf.disableSSL == "1" {
+	if conf.disableSSL {
 		opts += "sslmode=disable"
 	}
 	connStr := fmt.Sprintf("%s://%s:%s@%s:%s?%s", driver, conf.user, conf.password, conf.host, conf.port, opts)
 
-	db, err := sql.Open(driver, connStr)
+	conn, err := pgx.Connect(ctx, connStr)
 	if err != nil {
 		log.Panicf("Error opening database: %v", err)
 	}
-	defer db.Close()
 
-	result, err := db.Exec(fmt.Sprintf("SELECT 'CREATE DATABASE %s' WHERE NOT EXISTS (SELECT FROM pg_database WHERE datname = '%s')", conf.name, conf.name))
+	err = conn.Ping(ctx)
+	if err != nil {
+		log.Panicf("Error pinging database: %v", err)
+	}
+
+	rs, err := conn.Query(ctx, fmt.Sprintf("SELECT 'CREATE DATABASE %s' WHERE NOT EXISTS (SELECT FROM pg_database WHERE datname = '%s')", conf.name, conf.name))
 	if err != nil {
 		log.Panicf("Error creating database: %v", err)
 	}
 
-	rows, err := result.RowsAffected()
+	if rs.Next() {
-	if err != nil {
-		log.Panic(err)
-	} else if rows == 0 {
-		log.Printf("Database already exists: %s", conf.name)
-	} else {
 		log.Printf("Database created: %s", conf.name)
+	} else {
+		log.Printf("Database already exists: %s", conf.name)
 	}
 }
 
@@ -47,8 +50,8 @@ type MigrationConfig struct {
 	ConnectionString string
 }
 
-func RunMigrations(config *MigrationConfig) {
+func RunMigrations(ctx context.Context, config *MigrationConfig) {
-	createDB()
+	// createDB(ctx) - Add back only if necessary
 
 	connStr := config.ConnectionString
 	if connStr == "" {

internal/database/migrations_test.go

@@ -5,6 +5,7 @@ import (
 	"os"
 	"path"
 	"queryorchestration/internal/database"
+	"queryorchestration/internal/test"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
@@ -12,19 +13,24 @@ import (
 
 func TestRunMigrations(t *testing.T) {
 	ctx := context.Background()
-	_, cleanup := createDB(t, ctx)
+
+	_, cleanup := test.CreateDB(t, ctx, &test.CreateDatabaseConfig{
+		NoPool: true,
+	})
 	defer cleanup()
 
 	t.Setenv("DB_NOSSL", "1")
 
-	database.RunMigrations(&database.MigrationConfig{
+	database.RunMigrations(ctx, &database.MigrationConfig{
 		BasePath: path.Join(os.Getenv("PWD"), "../.."),
 	})
 }
 
 func TestRunMigrationsNoDB(t *testing.T) {
+	ctx := context.Background()
+
 	assert.Panics(t, func() {
-		database.RunMigrations(&database.MigrationConfig{
+		database.RunMigrations(ctx, &database.MigrationConfig{
 			BasePath: path.Join(os.Getenv("PWD"), "../.."),
 		})
 	})

internal/database/migrationsprivate_test.go

@@ -0,0 +1,13 @@
+package database
+
+import (
+	"context"
+	"testing"
+)
+
+func TestCreateDB(t *testing.T) {
+	t.Skip("Must test appropriately")
+	ctx := context.Background()
+
+	createDB(ctx)
+}

internal/database/repository/collector_test.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"queryorchestration/internal/database"
 	"queryorchestration/internal/database/repository"
+	"queryorchestration/internal/test"
 	"testing"
 
 	"github.com/google/uuid"
@@ -12,9 +13,13 @@ import (
 
 func TestCollector(t *testing.T) {
 	ctx := context.Background()
-	db, cleanup := createDB(t, ctx)
+	db, cleanup := test.CreateDB(t, ctx, &test.CreateDatabaseConfig{
+		Migrations: &database.MigrationConfig{
+			BasePath: "../../..",
+		}})
 	defer cleanup()
-	queries := repository.New(db.pool)
+
+	queries := repository.New(db.Pool)
 
 	collectorID := database.MustToDBUUID(uuid.New())
 

internal/database/repository/query_test.go

@@ -3,7 +3,9 @@ package repository_test
 import (
 	"context"
 	"log"
+	"queryorchestration/internal/database"
 	"queryorchestration/internal/database/repository"
+	"queryorchestration/internal/test"
 	"testing"
 
 	"github.com/jackc/pgx/v5/pgtype"
@@ -12,9 +14,14 @@ import (
 
 func TestQueries(t *testing.T) {
 	ctx := context.Background()
-	db, cleanup := createDB(t, ctx)
+
+	db, cleanup := test.CreateDB(t, ctx, &test.CreateDatabaseConfig{
+		Migrations: &database.MigrationConfig{
+			BasePath: "../../..",
+		}})
 	defer cleanup()
-	queries := repository.New(db.pool)
+
+	queries := repository.New(db.Pool)
 
 	contextQueryID, err := queries.CreateQuery(ctx, repository.Querytype(repository.QuerytypeContextFull))
 	assert.Nil(t, err)

internal/database/repository/result_test.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"queryorchestration/internal/database"
 	"queryorchestration/internal/database/repository"
+	"queryorchestration/internal/test"
 	"testing"
 
 	"github.com/google/uuid"
@@ -13,9 +14,14 @@ import (
 
 func TestResults(t *testing.T) {
 	ctx := context.Background()
-	db, cleanup := createDB(t, ctx)
+
+	db, cleanup := test.CreateDB(t, ctx, &test.CreateDatabaseConfig{
+		Migrations: &database.MigrationConfig{
+			BasePath: "../../..",
+		}})
 	defer cleanup()
-	queries := repository.New(db.pool)
+
+	queries := repository.New(db.Pool)
 
 	jsonQueryID, err := queries.CreateQuery(ctx, repository.Querytype(repository.QuerytypeJsonExtractor))
 	assert.Nil(t, err)

internal/queue/delete_test.go

@@ -3,6 +3,7 @@ package queue_test
 import (
 	"context"
 	"queryorchestration/internal/queue"
+	"queryorchestration/internal/test"
 	"testing"
 
 	"github.com/aws/aws-sdk-go-v2/service/sqs/types"
@@ -11,22 +12,19 @@ import (
 
 func TestDelete(t *testing.T) {
 	ctx := context.Background()
-	queueConfig, cleanup := createQueue(t, ctx)
+	queueConfig, cleanup := test.CreateQueue(t, ctx, &test.CreateQueueConfig{})
 	defer cleanup()
 
-	err := queue.Send(ctx, queueConfig.Config, "{}", map[string]types.MessageAttributeValue{})
+	cfg := &queue.Config{
+		URL:    queueConfig.URL,
+		Client: queueConfig.Client,
+	}
+
+	err := queue.Send(ctx, cfg, "{}", map[string]types.MessageAttributeValue{})
 	assert.Nil(t, err)
 
-	result, err := queue.Receive(ctx, queueConfig.Config, []string{})
+	message := test.AssertMessageWait(t, ctx, cfg, []string{})
-	assert.Nil(t, err)
 
-	assert.Len(t, result.Messages, 1)
+	err = queue.Delete(ctx, cfg, &message)
-	message := result.Messages[0]
-
-	err = queue.Delete(ctx, queueConfig.Config, &message)
 	assert.Nil(t, err)
-
-	result, err = queue.Receive(ctx, queueConfig.Config, []string{})
-	assert.Nil(t, err)
-	assert.Len(t, result.Messages, 0)
 }

internal/queue/listener_test.go

@@ -2,25 +2,20 @@ package queue_test
 
 import (
 	"context"
-	"fmt"
-	"queryorchestration/internal/database"
 	"queryorchestration/internal/queue"
 	"queryorchestration/internal/server"
+	"queryorchestration/internal/test"
 	"testing"
 	"time"
 
-	"github.com/docker/go-connections/nat"
-	"github.com/jackc/pgx/v5/pgxpool"
 	"github.com/stretchr/testify/assert"
-	"github.com/testcontainers/testcontainers-go"
-	"github.com/testcontainers/testcontainers-go/wait"
 )
 
 func TestNew(t *testing.T) {
 	ctx := context.Background()
-	_, cleanup := createQueue(t, ctx)
+	_, cleanup := test.CreateQueue(t, ctx, &test.CreateQueueConfig{})
 	defer cleanup()
-	_, cleanup = createDB(t, ctx)
+	_, cleanup = test.CreateDB(t, ctx, &test.CreateDatabaseConfig{})
 	defer cleanup()
 
 	ctx, cancel := context.WithTimeout(ctx, time.Second)
@@ -43,66 +38,3 @@ func TestListen(t *testing.T) {
 
 	assert.Panics(t, func() { queue.Listen(ctx) })
 }
-
-type db struct {
-	pool      *pgxpool.Pool
-	container *testcontainers.Container
-}
-
-func createDB(t *testing.T, ctx context.Context) (*db, func()) {
-	port, err := nat.NewPort("tcp", "5432")
-	if err != nil {
-		t.Fatalf("Failed to create port: %v", err)
-	}
-
-	name := "queryorchestration"
-	pass := "pass"
-	user := "postgres"
-
-	req := testcontainers.ContainerRequest{
-		Image: "postgres:17.2-alpine3.21",
-		Env: map[string]string{
-			"POSTGRES_DB":       name,
-			"POSTGRES_USER":     user,
-			"POSTGRES_PASSWORD": pass,
-		},
-		ExposedPorts: []string{port.Port()},
-		WaitingFor:   wait.ForListeningPort(port),
-	}
-
-	container, err := testcontainers.GenericContainer(ctx, testcontainers.GenericContainerRequest{
-		ContainerRequest: req,
-		Started:          true,
-	})
-	if err != nil {
-		t.Fatalf("Failed to start container: %v", err)
-	}
-
-	host, err := container.Host(ctx)
-	if err != nil {
-		t.Fatalf("Failed to extract host: %v", err)
-	}
-	mappedPort, err := container.MappedPort(ctx, port)
-	if err != nil {
-		t.Fatalf("Failed to extract port: %v", err)
-	}
-
-	t.Setenv("DB_USER", user)
-	t.Setenv("DB_PASS", pass)
-	t.Setenv("DB_HOST", host)
-	t.Setenv("DB_PORT", fmt.Sprint(mappedPort.Int()))
-	t.Setenv("DB_NAME", name)
-	t.Setenv("DB_NOSSL", "1")
-
-	pool := database.GetDBPool(ctx)
-
-	return &db{
-			pool:      pool,
-			container: &container,
-		}, func() {
-			err := container.Terminate(ctx)
-			if err != nil {
-				t.Error(err)
-			}
-		}
-}

internal/queue/poll_test.go

@@ -3,6 +3,7 @@ package queue_test
 import (
 	"context"
 	"queryorchestration/internal/queue"
+	"queryorchestration/internal/test"
 	"testing"
 	"time"
 
@@ -18,7 +19,7 @@ func (s MockController) Process(ctx context.Context, config *queue.Config, msg *
 
 func TestPollMessages(t *testing.T) {
 	ctx := context.Background()
-	queueConfig, cleanup := createQueue(t, ctx)
+	queueConfig, cleanup := test.CreateQueue(t, ctx, &test.CreateQueueConfig{})
 	defer cleanup()
 
 	controller := MockController{}
@@ -26,21 +27,29 @@ func TestPollMessages(t *testing.T) {
 	ctx, cancel := context.WithTimeout(ctx, time.Second)
 	defer cancel()
 
+	cfg := &queue.Config{
+		URL:    queueConfig.URL,
+		Client: queueConfig.Client,
+	}
 	queue.PollMessages(ctx, &queue.PollConfig{
-		Config:     queueConfig.Config,
+		Config:     cfg,
 		Controller: controller,
 	})
 }
 
 func TestPollMessage(t *testing.T) {
 	ctx := context.Background()
-	queueConfig, cleanup := createQueue(t, ctx)
+	queueConfig, cleanup := test.CreateQueue(t, ctx, &test.CreateQueueConfig{})
 	defer cleanup()
 
 	controller := MockController{}
 
+	cfg := &queue.Config{
+		URL:    queueConfig.URL,
+		Client: queueConfig.Client,
+	}
 	err := queue.PollMessage(ctx, &queue.PollConfig{
-		Config:     queueConfig.Config,
+		Config:     cfg,
 		Controller: controller,
 	})
 	assert.Nil(t, err)

internal/queue/receive_test.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"log"
 	"queryorchestration/internal/queue"
+	"queryorchestration/internal/test"
 	"testing"
 
 	"github.com/aws/aws-sdk-go-v2/service/sqs/types"
@@ -12,15 +13,19 @@ import (
 
 func TestReceive(t *testing.T) {
 	ctx := context.Background()
-	queueConfig, cleanup := createQueue(t, ctx)
+	queueConfig, cleanup := test.CreateQueue(t, ctx, &test.CreateQueueConfig{})
 	defer cleanup()
+	cfg := &queue.Config{
+		URL:    queueConfig.URL,
+		Client: queueConfig.Client,
+	}
 
 	attributes := map[string]types.MessageAttributeValue{}
 
-	err := queue.Send(ctx, queueConfig.Config, "example_body", attributes)
+	err := queue.Send(ctx, cfg, "example_body", attributes)
 	assert.Nil(t, err)
 
-	result, err := queue.Receive(ctx, queueConfig.Config, []string{})
+	result, err := queue.Receive(ctx, cfg, []string{})
 	assert.Nil(t, err)
 
 	assert.Len(t, result.Messages, 1)

internal/queue/send_test.go

@@ -3,23 +3,22 @@ package queue_test
 import (
 	"context"
 	"queryorchestration/internal/queue"
+	"queryorchestration/internal/test"
 	"testing"
 
-	"github.com/aws/aws-sdk-go-v2/aws"
 	"github.com/aws/aws-sdk-go-v2/service/sqs/types"
 	"github.com/stretchr/testify/assert"
 )
 
 func TestSend(t *testing.T) {
 	ctx := context.Background()
-	queueConfig, cleanup := createQueue(t, ctx)
+	queueConfig, cleanup := test.CreateQueue(t, ctx, &test.CreateQueueConfig{})
 	defer cleanup()
+	cfg := &queue.Config{
+		URL:    queueConfig.URL,
+		Client: queueConfig.Client,
+	}
 
-	err := queue.Send(ctx, queueConfig.Config, "{}", map[string]types.MessageAttributeValue{
+	err := queue.Send(ctx, cfg, "{}", map[string]types.MessageAttributeValue{})
-		"type": {
-			DataType:    aws.String("String"),
-			StringValue: aws.String("EXAMPLE_TYPE"),
-		},
-	})
 	assert.Nil(t, err)
 }

internal/server/server.go

@@ -2,6 +2,7 @@ package server
 
 import (
 	"context"
+	"log"
 	"queryorchestration/internal/database"
 	"queryorchestration/internal/database/repository"
 	"queryorchestration/internal/otel"
@@ -22,17 +23,23 @@ func New(ctx context.Context, cfg *NewConfig) *Config {
 	closeTracer := otel.New(ctx)
 	defer closeTracer()
 
-	database.RunMigrations(&database.MigrationConfig{
+	database.RunMigrations(ctx, &database.MigrationConfig{
 		BasePath: cfg.BasePath,
 	})
 
+	valid := validator.New()
+
 	dbPool := database.GetDBPool(ctx)
 	dbQueries := repository.New(dbPool)
 	db := &database.Connection{
 		Pool:    dbPool,
 		Queries: dbQueries,
 	}
-	valid := validator.New()
+
+	err := dbPool.Ping(ctx)
+	if err != nil {
+		log.Panic("Unable to ping database")
+	}
 
 	return &Config{
 		Database:  db,

internal/server/server_test.go

@@ -2,22 +2,19 @@ package server_test
 
 import (
 	"context"
-	"fmt"
-	"queryorchestration/internal/database"
 	"queryorchestration/internal/server"
+	"queryorchestration/internal/test"
 	"testing"
 
-	"github.com/docker/go-connections/nat"
-	"github.com/jackc/pgx/v5/pgxpool"
 	"github.com/stretchr/testify/assert"
-	"github.com/testcontainers/testcontainers-go"
-	"github.com/testcontainers/testcontainers-go/wait"
 )
 
 func TestNew(t *testing.T) {
 	ctx := context.Background()
 
-	_, cleanup := createDB(t, ctx)
+	_, cleanup := test.CreateDB(t, ctx, &test.CreateDatabaseConfig{
+		NoPool: true,
+	})
 	defer cleanup()
 
 	newCfg := &server.NewConfig{
@@ -27,66 +24,3 @@ func TestNew(t *testing.T) {
 	cfg := server.New(ctx, newCfg)
 	assert.NotNil(t, cfg)
 }
-
-type db struct {
-	pool      *pgxpool.Pool
-	container *testcontainers.Container
-}
-
-func createDB(t *testing.T, ctx context.Context) (*db, func()) {
-	port, err := nat.NewPort("tcp", "5432")
-	if err != nil {
-		t.Fatalf("Failed to create port: %v", err)
-	}
-
-	name := "queryorchestration"
-	pass := "pass"
-	user := "postgres"
-
-	req := testcontainers.ContainerRequest{
-		Image: "postgres:17.2-alpine3.21",
-		Env: map[string]string{
-			"POSTGRES_DB":       name,
-			"POSTGRES_USER":     user,
-			"POSTGRES_PASSWORD": pass,
-		},
-		ExposedPorts: []string{port.Port()},
-		WaitingFor:   wait.ForListeningPort(port),
-	}
-
-	container, err := testcontainers.GenericContainer(ctx, testcontainers.GenericContainerRequest{
-		ContainerRequest: req,
-		Started:          true,
-	})
-	if err != nil {
-		t.Fatalf("Failed to start container: %v", err)
-	}
-
-	host, err := container.Host(ctx)
-	if err != nil {
-		t.Fatalf("Failed to extract host: %v", err)
-	}
-	mappedPort, err := container.MappedPort(ctx, port)
-	if err != nil {
-		t.Fatalf("Failed to extract port: %v", err)
-	}
-
-	t.Setenv("DB_USER", user)
-	t.Setenv("DB_PASS", pass)
-	t.Setenv("DB_HOST", host)
-	t.Setenv("DB_PORT", fmt.Sprint(mappedPort.Int()))
-	t.Setenv("DB_NAME", name)
-	t.Setenv("DB_NOSSL", "1")
-
-	pool := database.GetDBPool(ctx)
-
-	return &db{
-			pool:      pool,
-			container: &container,
-		}, func() {
-			err := container.Terminate(ctx)
-			if err != nil {
-				t.Error(err)
-			}
-		}
-}

internal/test/api.go

@@ -0,0 +1,79 @@
+package test
+
+import (
+	"context"
+	"fmt"
+	"testing"
+
+	"github.com/docker/go-connections/nat"
+	"github.com/testcontainers/testcontainers-go"
+	"google.golang.org/grpc"
+	"google.golang.org/grpc/credentials/insecure"
+)
+
+type APIContainerConfig struct {
+	ServiceName string
+	DB          *ExternalDatabase
+	Network     *testcontainers.DockerNetwork
+}
+
+func CreateAPIContainer(t *testing.T, ctx context.Context, config *APIContainerConfig) (*grpc.ClientConn, func()) {
+	port, err := nat.NewPort("tcp", "8080")
+	if err != nil {
+		t.Fatalf("Failed to create port: %v", err)
+	}
+
+	container, cleanup := createContainer(t, ctx, &containerConfig{
+		DB:           config.DB,
+		ServiceName:  config.ServiceName,
+		Network:      config.Network,
+		ExposedPorts: []nat.Port{port},
+		WaitForMsg:   "Listening on port 8080",
+	})
+
+	host, err := container.Host(ctx)
+	if err != nil {
+		t.Fatalf("Failed to extract host: %v", err)
+	}
+	mappedPort, err := container.MappedPort(ctx, port)
+	if err != nil {
+		t.Fatalf("Failed to extract port: %v", err)
+	}
+
+	conn, err := grpc.NewClient(
+		fmt.Sprintf("%s:%s", host, mappedPort.Port()),
+		grpc.WithTransportCredentials(insecure.NewCredentials()),
+	)
+	if err != nil {
+		t.Fatalf("Failed to connect to gRPC server: %v", err)
+	}
+
+	return conn, func() {
+		err = conn.Close()
+		if err != nil {
+			t.Error(err)
+		}
+
+		cleanup()
+	}
+}
+
+func CreateAPIWithDependencies(t *testing.T, ctx context.Context, serviceName string) (*grpc.ClientConn, func()) {
+	network, ncleanup := CreateNetwork(t, ctx)
+
+	dbconfig, dbcleanup := CreateDB(t, ctx, &CreateDatabaseConfig{
+		Network: network,
+	})
+
+	conn, ccleanup := CreateAPIContainer(t, ctx, &APIContainerConfig{
+		ServiceName: serviceName,
+		DB:          dbconfig.External,
+		Network:     network,
+	})
+
+	return conn, func() {
+		ncleanup()
+		dbcleanup()
+		ccleanup()
+	}
+}

internal/test/api_test.go

@@ -0,0 +1,46 @@
+package test_test
+
+import (
+	"context"
+	"queryorchestration/internal/database"
+	"queryorchestration/internal/test"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestCreateAPIContainer(t *testing.T) {
+	ctx := context.Background()
+
+	ncfg, ncleanup := test.CreateNetwork(t, ctx)
+	defer ncleanup()
+	dbcfg, dbcleanup := test.CreateDB(t, ctx, &test.CreateDatabaseConfig{
+		Network: ncfg,
+		Migrations: &database.MigrationConfig{
+			BasePath: "../..",
+		},
+	})
+	defer dbcleanup()
+
+	cfg := &test.APIContainerConfig{
+		ServiceName: "queryService",
+		DB:          dbcfg.External,
+		Network:     ncfg,
+	}
+
+	conn, cleanup := test.CreateAPIContainer(t, ctx, cfg)
+	assert.NotNil(t, conn)
+	assert.NotNil(t, cleanup)
+
+	cleanup()
+}
+func TestCreateAPIWithDependencies(t *testing.T) {
+	ctx := context.Background()
+
+	conn, cleanup := test.CreateAPIWithDependencies(t, ctx, "queryService")
+
+	assert.NotNil(t, conn)
+	assert.NotNil(t, cleanup)
+
+	cleanup()
+}

internal/test/container.go

@@ -0,0 +1,68 @@
+package test
+
+import (
+	"context"
+	"fmt"
+	"strconv"
+	"testing"
+
+	"github.com/docker/go-connections/nat"
+	"github.com/testcontainers/testcontainers-go"
+	"github.com/testcontainers/testcontainers-go/wait"
+)
+
+type containerConfig struct {
+	ServiceName  string
+	DB           *ExternalDatabase
+	Network      *testcontainers.DockerNetwork
+	Env          map[string]string
+	WaitForMsg   string
+	ExposedPorts []nat.Port
+}
+
+func createContainer(t *testing.T, ctx context.Context, cfg *containerConfig) (testcontainers.Container, func()) {
+	env := map[string]string{
+		"DB_USER":  cfg.DB.User,
+		"DB_PASS":  cfg.DB.Password,
+		"DB_HOST":  cfg.DB.Host,
+		"DB_NAME":  cfg.DB.Name,
+		"DB_PORT":  strconv.Itoa(cfg.DB.Port),
+		"DB_NOSSL": "1",
+	}
+	if cfg.Env != nil {
+		for k, v := range cfg.Env {
+			env[k] = v
+		}
+	}
+
+	req := testcontainers.ContainerRequest{
+		Image:      "queryorchestration:latest",
+		Env:        env,
+		Networks:   []string{cfg.Network.Name},
+		WaitingFor: wait.ForLog(cfg.WaitForMsg),
+		Entrypoint: []string{fmt.Sprintf("./%s", cfg.ServiceName)},
+	}
+
+	if cfg.ExposedPorts != nil {
+		ports := make([]string, len(cfg.ExposedPorts))
+		for index, port := range cfg.ExposedPorts {
+			ports[index] = port.Port()
+		}
+		req.ExposedPorts = ports
+	}
+
+	container, err := testcontainers.GenericContainer(ctx, testcontainers.GenericContainerRequest{
+		ContainerRequest: req,
+		Started:          true,
+	})
+	if err != nil {
+		t.Fatalf("Failed to start container: %v", err)
+	}
+
+	return container, func() {
+		err := container.Terminate(ctx)
+		if err != nil {
+			t.Error(err)
+		}
+	}
+}

internal/test/container_test.go

@@ -0,0 +1,35 @@
+package test
+
+import (
+	"context"
+	"queryorchestration/internal/database"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestCreateContainer(t *testing.T) {
+	ctx := context.Background()
+
+	ncfg, ncleanup := CreateNetwork(t, ctx)
+	defer ncleanup()
+	dbcfg, dbcleanup := CreateDB(t, ctx, &CreateDatabaseConfig{
+		Network: ncfg,
+		Migrations: &database.MigrationConfig{
+			BasePath: "../..",
+		},
+	})
+	defer dbcleanup()
+
+	cfg := &containerConfig{
+		ServiceName: "queryService",
+		DB:          dbcfg.External,
+		Network:     ncfg,
+	}
+
+	container, cleanup := createContainer(t, ctx, cfg)
+	assert.NotNil(t, container)
+	assert.NotNil(t, cleanup)
+
+	cleanup()
+}

internal/test/database.go

@@ -1,10 +1,8 @@
-package repository_test
+package test
 
 import (
 	"context"
 	"fmt"
-	"os"
-	"path"
 	"queryorchestration/internal/database"
 	"testing"
 
@@ -14,12 +12,29 @@ import (
 	"github.com/testcontainers/testcontainers-go/wait"
 )
 
-type db struct {
+type ExternalDatabase struct {
-	pool      *pgxpool.Pool
+	Name     string
-	container *testcontainers.Container
+	Host     string
+	Port     int
+	User     string
+	Password string
 }
 
-func createDB(t *testing.T, ctx context.Context) (*db, func()) {
+type Database struct {
+	Pool      *pgxpool.Pool
+	Container testcontainers.Container
+	External  *ExternalDatabase
+}
+
+type CreateDatabaseConfig struct {
+	Network    *testcontainers.DockerNetwork
+	Migrations *database.MigrationConfig
+	NoPool     bool
+}
+
+func CreateDB(t *testing.T, ctx context.Context, cfg *CreateDatabaseConfig) (*Database, func()) {
+	alias := "postgres"
+
 	port, err := nat.NewPort("tcp", "5432")
 	if err != nil {
 		t.Fatalf("Failed to create port: %v", err)
@@ -40,6 +55,13 @@ func createDB(t *testing.T, ctx context.Context) (*db, func()) {
 		WaitingFor:   wait.ForListeningPort(port),
 	}
 
+	if cfg.Network != nil {
+		req.Networks = []string{cfg.Network.Name}
+		req.NetworkAliases = map[string][]string{
+			cfg.Network.Name: {alias},
+		}
+	}
+
 	container, err := testcontainers.GenericContainer(ctx, testcontainers.GenericContainerRequest{
 		ContainerRequest: req,
 		Started:          true,
@@ -64,15 +86,32 @@ func createDB(t *testing.T, ctx context.Context) (*db, func()) {
 	t.Setenv("DB_NAME", name)
 	t.Setenv("DB_NOSSL", "1")
 
-	database.RunMigrations(&database.MigrationConfig{
+	if cfg.Migrations != nil {
-		BasePath: path.Join(os.Getenv("PWD"), "../../.."),
+		database.RunMigrations(ctx, cfg.Migrations)
-	})
+	}
 
-	pool := database.GetDBPool(ctx)
+	var pool *pgxpool.Pool
+	if !cfg.NoPool {
+		pool = database.GetDBPool(ctx)
 
-	return &db{
+		err = pool.Ping(ctx)
-			pool:      pool,
+		if err != nil {
-			container: &container,
+			t.Fatal("Unable to ping database")
+		}
+	}
+
+	external := &ExternalDatabase{
+		Name:     name,
+		Password: pass,
+		User:     user,
+		Port:     port.Int(),
+		Host:     alias,
+	}
+
+	return &Database{
+			Pool:      pool,
+			Container: container,
+			External:  external,
 		}, func() {
 			err := container.Terminate(ctx)
 			if err != nil {

internal/test/database_test.go

@@ -0,0 +1,21 @@
+package test_test
+
+import (
+	"context"
+	"queryorchestration/internal/test"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestCreateDB(t *testing.T) {
+	ctx := context.Background()
+
+	dbcfg, cleanup := test.CreateDB(t, ctx, &test.CreateDatabaseConfig{
+		NoPool: true,
+	})
+	assert.NotNil(t, dbcfg)
+	assert.NotNil(t, cleanup)
+
+	cleanup()
+}

internal/test/network.go

@@ -0,0 +1,20 @@
+package test
+
+import (
+	"context"
+	"testing"
+
+	"github.com/testcontainers/testcontainers-go"
+	"github.com/testcontainers/testcontainers-go/network"
+)
+
+func CreateNetwork(t *testing.T, ctx context.Context) (*testcontainers.DockerNetwork, func()) {
+	network, err := network.New(ctx, network.WithDriver("bridge"))
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	return network, func() {
+		testcontainers.CleanupNetwork(t, network)
+	}
+}

internal/test/network_test.go

@@ -0,0 +1,19 @@
+package test_test
+
+import (
+	"context"
+	"queryorchestration/internal/test"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestCreateNetwork(t *testing.T) {
+	ctx := context.Background()
+
+	ncfg, cleanup := test.CreateNetwork(t, ctx)
+	assert.NotNil(t, ncfg)
+	assert.NotNil(t, cleanup)
+
+	cleanup()
+}

internal/test/queue.go

@@ -1,26 +1,42 @@
-package queue_test
+package test
 
 import (
 	"context"
 	"fmt"
 	"queryorchestration/internal/queue"
 	"testing"
+	"time"
 
 	"github.com/aws/aws-sdk-go-v2/aws"
 	"github.com/aws/aws-sdk-go-v2/config"
 	"github.com/aws/aws-sdk-go-v2/credentials"
 	"github.com/aws/aws-sdk-go-v2/service/sqs"
+	"github.com/aws/aws-sdk-go-v2/service/sqs/types"
 	"github.com/docker/go-connections/nat"
+	"github.com/stretchr/testify/assert"
 	"github.com/testcontainers/testcontainers-go"
 	"github.com/testcontainers/testcontainers-go/wait"
 )
 
-type queueConfig struct {
+type QueueConfig struct {
-	Container *testcontainers.Container
+	Container testcontainers.Container
-	Config    *queue.Config
+	Client    *sqs.Client
+	URL       string
+	External  *ExternalQueue
 }
 
-func createQueue(t *testing.T, ctx context.Context) (*queueConfig, func()) {
+type ExternalQueue struct {
+	Endpoint    string
+	Credentials aws.CredentialsProvider
+	Region      string
+	URL         string
+}
+
+type CreateQueueConfig struct {
+	Network *testcontainers.DockerNetwork
+}
+
+func CreateQueue(t *testing.T, ctx context.Context, cfg *CreateQueueConfig) (*QueueConfig, func()) {
 	queueName := "test-queue"
 	region := "us-east-1"
 	alias := "localstack"
@@ -48,6 +64,13 @@ func createQueue(t *testing.T, ctx context.Context) (*queueConfig, func()) {
 		WaitingFor:   wait.ForListeningPort(port),
 	}
 
+	if cfg.Network != nil {
+		req.Networks = []string{cfg.Network.Name}
+		req.NetworkAliases = map[string][]string{
+			cfg.Network.Name: {alias},
+		}
+	}
+
 	container, err := testcontainers.GenericContainer(ctx, testcontainers.GenericContainerRequest{
 		ContainerRequest: req,
 		Started:          true,
@@ -66,12 +89,12 @@ func createQueue(t *testing.T, ctx context.Context) (*queueConfig, func()) {
 	}
 
 	endpoint := fmt.Sprintf("http://%s:%s", host, mappedPort.Port())
-	cfg, err := config.LoadDefaultConfig(ctx, config.WithRegion(region), config.WithCredentialsProvider(provider), config.WithBaseEndpoint(endpoint))
+	sqsCfg, err := config.LoadDefaultConfig(ctx, config.WithRegion(region), config.WithCredentialsProvider(provider), config.WithBaseEndpoint(endpoint))
 	if err != nil {
 		t.Fatal(err)
 	}
 
-	client := sqs.NewFromConfig(cfg)
+	client := sqs.NewFromConfig(sqsCfg)
 
 	queueM, err := client.CreateQueue(ctx, &sqs.CreateQueueInput{
 		QueueName: aws.String(queueName),
@@ -80,11 +103,16 @@ func createQueue(t *testing.T, ctx context.Context) (*queueConfig, func()) {
 		t.Fatal(err)
 	}
 
-	return &queueConfig{
+	endpoint = fmt.Sprintf("http://%s:%s", alias, port.Port())
-			Container: &container,
+
-			Config: &queue.Config{
+	return &QueueConfig{
-				Client: client,
+			Container: container,
-				URL:    *queueM.QueueUrl,
+			Client:    client,
+			URL:       *queueM.QueueUrl,
+			External: &ExternalQueue{
+				Region:      region,
+				Credentials: provider,
+				Endpoint:    endpoint,
 			},
 		}, func() {
 			err := container.Terminate(ctx)
@@ -93,3 +121,27 @@ func createQueue(t *testing.T, ctx context.Context) (*queueConfig, func()) {
 			}
 		}
 }
+
+func AssertMessageWait(t *testing.T, ctx context.Context, cfg *queue.Config, attrs []string) types.Message {
+	time.Sleep(time.Second)
+
+	result, err := queue.Receive(ctx, cfg, attrs)
+	assert.Nil(t, err)
+	assert.Len(t, result.Messages, 1)
+	assert.NotNil(t, result.Messages[0])
+
+	return result.Messages[0]
+}
+
+func AssertMessageBodyWait(t *testing.T, ctx context.Context, cfg *queue.Config, body string) {
+	message := AssertMessageWait(t, ctx, cfg, []string{})
+
+	assert.Equal(t, body, *message.Body)
+}
+
+func AssertMessageAttrWait(t *testing.T, ctx context.Context, cfg *queue.Config, name string, value string) {
+	message := AssertMessageWait(t, ctx, cfg, []string{name})
+
+	assert.NotNil(t, message.MessageAttributes[name])
+	assert.Equal(t, value, *(message.MessageAttributes[name]).StringValue)
+}

internal/test/queue_test.go

@@ -0,0 +1,82 @@
+package test_test
+
+import (
+	"context"
+	"queryorchestration/internal/queue"
+	"queryorchestration/internal/test"
+	"testing"
+
+	"github.com/aws/aws-sdk-go-v2/aws"
+	"github.com/aws/aws-sdk-go-v2/service/sqs/types"
+	"github.com/stretchr/testify/assert"
+)
+
+func TestCreateQueue(t *testing.T) {
+	ctx := context.Background()
+
+	qcfg, cleanup := test.CreateQueue(t, ctx, &test.CreateQueueConfig{})
+	assert.NotNil(t, qcfg)
+	assert.NotNil(t, cleanup)
+
+	cleanup()
+}
+
+func TestAssertMessageWait(t *testing.T) {
+	ctx := context.Background()
+
+	qcfg, cleanup := test.CreateQueue(t, ctx, &test.CreateQueueConfig{})
+	defer cleanup()
+
+	cfg := &queue.Config{
+		URL:    qcfg.URL,
+		Client: qcfg.Client,
+	}
+
+	err := queue.Send(ctx, cfg, "body", map[string]types.MessageAttributeValue{})
+	assert.Nil(t, err)
+
+	msg := test.AssertMessageWait(t, ctx, cfg, []string{})
+	assert.NotNil(t, msg)
+}
+
+func TestAssertMessageBodyWait(t *testing.T) {
+	ctx := context.Background()
+
+	qcfg, cleanup := test.CreateQueue(t, ctx, &test.CreateQueueConfig{})
+	defer cleanup()
+
+	cfg := &queue.Config{
+		URL:    qcfg.URL,
+		Client: qcfg.Client,
+	}
+
+	err := queue.Send(ctx, cfg, "body", map[string]types.MessageAttributeValue{})
+	assert.Nil(t, err)
+
+	test.AssertMessageBodyWait(t, ctx, cfg, "\"body\"")
+}
+
+func TestAssertMessageAttrWait(t *testing.T) {
+	ctx := context.Background()
+
+	qcfg, cleanup := test.CreateQueue(t, ctx, &test.CreateQueueConfig{})
+	defer cleanup()
+
+	cfg := &queue.Config{
+		URL:    qcfg.URL,
+		Client: qcfg.Client,
+	}
+
+	name := "name"
+	value := "value"
+
+	err := queue.Send(ctx, cfg, "body", map[string]types.MessageAttributeValue{
+		name: {
+			DataType:    aws.String("String"),
+			StringValue: aws.String(value),
+		},
+	})
+	assert.Nil(t, err)
+
+	test.AssertMessageAttrWait(t, ctx, cfg, name, value)
+}

internal/test/queuecontainer.go

@@ -0,0 +1,81 @@
+package test
+
+import (
+	"context"
+	"testing"
+
+	"github.com/aws/aws-sdk-go-v2/aws"
+	"github.com/testcontainers/testcontainers-go"
+)
+
+type QueueContainerConfig struct {
+	ServiceName string
+	Queue       *queueConfig
+	DB          *ExternalDatabase
+	Network     *testcontainers.DockerNetwork
+}
+
+func CreateQueueContainer(t *testing.T, ctx context.Context, config *QueueContainerConfig) func() {
+	queueCredentials, err := config.Queue.Credentials.Retrieve(ctx)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	_, cleanup := createContainer(t, ctx, &containerConfig{
+		Network:     config.Network,
+		ServiceName: config.ServiceName,
+		DB:          config.DB,
+		WaitForMsg:  "Listening to queue",
+		Env: map[string]string{
+			"QUEUE_URL":             config.Queue.URL,
+			"AWS_DEFAULT_REGION":    config.Queue.Region,
+			"AWS_REGION":            config.Queue.Region,
+			"AWS_ACCESS_KEY_ID":     queueCredentials.AccessKeyID,
+			"AWS_SECRET_ACCESS_KEY": queueCredentials.SecretAccessKey,
+			"AWS_SESSION_TOKEN":     queueCredentials.SessionToken,
+			"AWS_ENDPOINT_URL_SQS":  config.Queue.Endpoint,
+		},
+	})
+
+	return cleanup
+}
+
+type queueConfig struct {
+	URL         string
+	Region      string
+	Endpoint    string
+	Credentials aws.CredentialsProvider
+}
+
+func CreateQueueWithDependencies(t *testing.T, ctx context.Context, serviceName string) (*QueueConfig, func()) {
+	network, ncleanup := CreateNetwork(t, ctx)
+
+	dbconfig, dbcleanup := CreateDB(t, ctx, &CreateDatabaseConfig{
+		Network: network,
+	})
+
+	qCfg, qcleanup := CreateQueue(t, ctx, &CreateQueueConfig{
+		Network: network,
+	})
+
+	cfg := &queueConfig{
+		URL:         qCfg.URL,
+		Region:      qCfg.External.Region,
+		Endpoint:    qCfg.External.Endpoint,
+		Credentials: qCfg.External.Credentials,
+	}
+
+	ccleanup := CreateQueueContainer(t, ctx, &QueueContainerConfig{
+		ServiceName: serviceName,
+		Queue:       cfg,
+		DB:          dbconfig.External,
+		Network:     network,
+	})
+
+	return qCfg, func() {
+		ncleanup()
+		dbcleanup()
+		qcleanup()
+		ccleanup()
+	}
+}

internal/test/queuecontainer_test.go

@@ -0,0 +1,52 @@
+package test
+
+import (
+	"context"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestCreateQueueContainer(t *testing.T) {
+	ctx := context.Background()
+
+	ncfg, ncleanup := CreateNetwork(t, ctx)
+	defer ncleanup()
+	dbcfg, dbcleanup := CreateDB(t, ctx, &CreateDatabaseConfig{
+		Network: ncfg,
+	})
+	defer dbcleanup()
+	qcfg, qcleanup := CreateQueue(t, ctx, &CreateQueueConfig{
+		Network: ncfg,
+	})
+	defer qcleanup()
+
+	extcfg := &queueConfig{
+		URL:         qcfg.URL,
+		Region:      qcfg.External.Region,
+		Endpoint:    qcfg.External.Endpoint,
+		Credentials: qcfg.External.Credentials,
+	}
+	cfg := &QueueContainerConfig{
+		ServiceName: "queryRunner",
+		DB:          dbcfg.External,
+		Network:     ncfg,
+		Queue:       extcfg,
+	}
+
+	cleanup := CreateQueueContainer(t, ctx, cfg)
+	assert.NotNil(t, cleanup)
+
+	cleanup()
+}
+
+func TestCreateQueueWithDependencies(t *testing.T) {
+	ctx := context.Background()
+
+	conn, cleanup := CreateQueueWithDependencies(t, ctx, "queryRunner")
+
+	assert.NotNil(t, conn)
+	assert.NotNil(t, cleanup)
+
+	cleanup()
+}

test/apiContainer_test.go

@@ -1,103 +0,0 @@
-package integration_test
-
-import (
-	"context"
-	"fmt"
-	"strconv"
-	"testing"
-
-	"github.com/docker/go-connections/nat"
-	"github.com/testcontainers/testcontainers-go"
-	"github.com/testcontainers/testcontainers-go/wait"
-	"google.golang.org/grpc"
-	"google.golang.org/grpc/credentials/insecure"
-)
-
-type apiContainerConfig struct {
-	ServiceName string
-	DB          dbConfig
-	Network     *testcontainers.DockerNetwork
-}
-
-func createAPIContainer(t *testing.T, ctx context.Context, config *apiContainerConfig) (*grpc.ClientConn, func()) {
-	port, err := nat.NewPort("tcp", "8080")
-	if err != nil {
-		t.Fatalf("Failed to create port: %v", err)
-	}
-
-	req := testcontainers.ContainerRequest{
-		Image: "queryorchestration:latest",
-		Env: map[string]string{
-			"DB_USER":  config.DB.User,
-			"DB_PASS":  config.DB.Password,
-			"DB_HOST":  config.DB.Host,
-			"DB_NAME":  config.DB.Name,
-			"DB_PORT":  strconv.Itoa(config.DB.Port),
-			"DB_NOSSL": "1",
-		},
-		ExposedPorts: []string{port.Port()},
-		WaitingFor:   wait.ForLog("Listening on port 8080"),
-		Networks:     []string{config.Network.Name},
-		Entrypoint:   []string{fmt.Sprintf("./%s", config.ServiceName)},
-	}
-
-	container, err := testcontainers.GenericContainer(ctx, testcontainers.GenericContainerRequest{
-		ContainerRequest: req,
-		Started:          true,
-	})
-	if err != nil {
-		t.Fatalf("Failed to start container: %v", err)
-	}
-
-	host, err := container.Host(ctx)
-	if err != nil {
-		t.Fatalf("Failed to extract host: %v", err)
-	}
-	mappedPort, err := container.MappedPort(ctx, port)
-	if err != nil {
-		t.Fatalf("Failed to extract port: %v", err)
-	}
-
-	conn, err := grpc.NewClient(
-		fmt.Sprintf("%s:%s", host, mappedPort.Port()),
-		grpc.WithTransportCredentials(insecure.NewCredentials()),
-	)
-	if err != nil {
-		t.Fatalf("Failed to connect to gRPC server: %v", err)
-	}
-
-	return conn, func() {
-		err := container.Terminate(ctx)
-		if err != nil {
-			t.Error(err)
-		}
-	}
-}
-
-func createAPIDependencies(t *testing.T, ctx context.Context, serviceName string) (*grpc.ClientConn, func()) {
-	network := createNetwork(t, ctx)
-
-	dbconfig, dbContainer := createDB(t, ctx, network)
-
-	conn, containerCleanup := createAPIContainer(t, ctx, &apiContainerConfig{
-		ServiceName: serviceName,
-		DB:          *dbconfig,
-		Network:     network,
-	})
-
-	return conn, func() {
-		testcontainers.CleanupNetwork(t, network)
-
-		err := dbContainer.Terminate(ctx)
-		if err != nil {
-			t.Error(err)
-		}
-
-		err = conn.Close()
-		if err != nil {
-			t.Error(err)
-		}
-
-		containerCleanup()
-	}
-}

test/container_test.go

@@ -1,69 +0,0 @@
-package integration_test
-
-import (
-	"context"
-	"testing"
-
-	"github.com/docker/go-connections/nat"
-	"github.com/testcontainers/testcontainers-go"
-	"github.com/testcontainers/testcontainers-go/network"
-	"github.com/testcontainers/testcontainers-go/wait"
-)
-
-func createNetwork(t *testing.T, ctx context.Context) *testcontainers.DockerNetwork {
-	network, err := network.New(ctx, network.WithDriver("bridge"))
-	if err != nil {
-		t.Fatal(err)
-	}
-
-	return network
-}
-
-type dbConfig struct {
-	Name     string
-	Host     string
-	Port     int
-	User     string
-	Password string
-}
-
-func createDB(t *testing.T, ctx context.Context, network *testcontainers.DockerNetwork) (*dbConfig, testcontainers.Container) {
-	alias := "postgres"
-	port, err := nat.NewPort("tcp", "5432")
-	if err != nil {
-		t.Fatalf("Failed to create port: %v", err)
-	}
-
-	config := dbConfig{
-		Name:     "queryorchestration",
-		Password: "pass",
-		User:     "postgres",
-		Port:     port.Int(),
-		Host:     alias,
-	}
-
-	req := testcontainers.ContainerRequest{
-		Image: "postgres:17.2-alpine3.21",
-		Env: map[string]string{
-			"POSTGRES_DB":       config.Name,
-			"POSTGRES_USER":     config.User,
-			"POSTGRES_PASSWORD": config.Password,
-		},
-		Networks: []string{network.Name},
-		NetworkAliases: map[string][]string{
-			network.Name: {alias},
-		},
-		ExposedPorts: []string{port.Port()},
-		WaitingFor:   wait.ForListeningPort(port),
-	}
-
-	container, err := testcontainers.GenericContainer(ctx, testcontainers.GenericContainerRequest{
-		ContainerRequest: req,
-		Started:          true,
-	})
-	if err != nil {
-		t.Fatalf("Failed to start container: %v", err)
-	}
-
-	return &config, container
-}

test/queryrunner_test.go

@@ -4,8 +4,11 @@ import (
 	"context"
 	"encoding/json"
 	"queryorchestration/internal/document"
+	"queryorchestration/internal/queue"
+	"queryorchestration/internal/test"
 	"testing"
 
+	"github.com/aws/aws-sdk-go-v2/service/sqs/types"
 	"github.com/google/uuid"
 	"github.com/stretchr/testify/assert"
 )
@@ -13,7 +16,7 @@ import (
 func TestQueryRunner(t *testing.T) {
 	ctx := context.Background()
 
-	queue, cleanup := createQueueDependencies(t, ctx, "queryRunner")
+	qCfg, cleanup := test.CreateQueueWithDependencies(t, ctx, "queryRunner")
 	defer cleanup()
 
 	document := document.Document{
@@ -26,11 +29,10 @@ func TestQueryRunner(t *testing.T) {
 	docJson, err := json.Marshal(document)
 	assert.Nil(t, err)
 
-	sendMessage(t, ctx, queue.Client, message{
+	cfg := &queue.Config{
-		URL:  queue.Config.URL,
+		URL:    qCfg.URL,
-		Type: "DOCTEXT",
+		Client: qCfg.Client,
-		Body: string(docJson),
+	}
-	})
 
-	assertMessageWait(t, ctx, queue, "DOCTEXT") // Fails so send it back to the queue
+	queue.Send(ctx, cfg, string(docJson), map[string]types.MessageAttributeValue{})
 }

test/queryservice_test.go

@@ -3,6 +3,7 @@ package integration_test
 import (
 	"context"
 	serviceinterfaces "queryorchestration/api/serviceInterfaces"
+	"queryorchestration/internal/test"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
@@ -11,7 +12,7 @@ import (
 func TestQueryService(t *testing.T) {
 	ctx := context.Background()
 
-	conn, cleanup := createAPIDependencies(t, ctx, "queryService")
+	conn, cleanup := test.CreateAPIWithDependencies(t, ctx, "queryService")
 	defer cleanup()
 
 	client := serviceinterfaces.NewQueryServiceClient(conn)

test/queueContainer_test.go

@@ -1,232 +0,0 @@
-package integration_test
-
-import (
-	"context"
-	"fmt"
-	"strconv"
-	"testing"
-	"time"
-
-	"github.com/aws/aws-sdk-go-v2/aws"
-	"github.com/aws/aws-sdk-go-v2/config"
-	"github.com/aws/aws-sdk-go-v2/credentials"
-	"github.com/aws/aws-sdk-go-v2/service/sqs"
-	"github.com/aws/aws-sdk-go-v2/service/sqs/types"
-	"github.com/docker/go-connections/nat"
-	"github.com/stretchr/testify/assert"
-	"github.com/testcontainers/testcontainers-go"
-	"github.com/testcontainers/testcontainers-go/wait"
-)
-
-type queueContainerConfig struct {
-	ServiceName string
-	Queue       queueConfig
-	DB          *dbConfig
-	Network     *testcontainers.DockerNetwork
-}
-
-func createQueueContainer(t *testing.T, ctx context.Context, config *queueContainerConfig) testcontainers.Container {
-	queueCredentials, err := config.Queue.Credentials.Retrieve(ctx)
-	if err != nil {
-		t.Fatal(err)
-	}
-
-	req := testcontainers.ContainerRequest{
-		Image: "queryorchestration:latest",
-		Env: map[string]string{
-			"QUEUE_URL":             config.Queue.URL,
-			"AWS_DEFAULT_REGION":    config.Queue.Region,
-			"AWS_REGION":            config.Queue.Region,
-			"AWS_ACCESS_KEY_ID":     queueCredentials.AccessKeyID,
-			"AWS_SECRET_ACCESS_KEY": queueCredentials.SecretAccessKey,
-			"AWS_SESSION_TOKEN":     queueCredentials.SessionToken,
-			"AWS_ENDPOINT_URL_SQS":  config.Queue.Endpoint,
-			"DB_USER":               config.DB.User,
-			"DB_PASS":               config.DB.Password,
-			"DB_HOST":               config.DB.Host,
-			"DB_NAME":               config.DB.Name,
-			"DB_PORT":               strconv.Itoa(config.DB.Port),
-			"DB_NOSSL":              "1",
-		},
-		Networks:   []string{config.Network.Name},
-		WaitingFor: wait.ForLog("Listening to queue"),
-		Entrypoint: []string{fmt.Sprintf("./%s", config.ServiceName)},
-	}
-
-	container, err := testcontainers.GenericContainer(ctx, testcontainers.GenericContainerRequest{
-		ContainerRequest: req,
-		Started:          true,
-	})
-	if err != nil {
-		t.Fatalf("Failed to start container: %v", err)
-	}
-
-	return container
-}
-
-type queueConfig struct {
-	URL         string
-	Region      string
-	Endpoint    string
-	Credentials aws.CredentialsProvider
-}
-
-type queue struct {
-	Container *testcontainers.Container
-	Client    *sqs.Client
-	Config    *queueConfig
-}
-
-func createQueue(t *testing.T, ctx context.Context, network *testcontainers.DockerNetwork) *queue {
-	queueName := "test-queue"
-	region := "us-east-1"
-	alias := "localstack"
-	provider := credentials.NewStaticCredentialsProvider("test", "test", "")
-
-	port, err := nat.NewPort("tcp", "4566")
-	if err != nil {
-		t.Fatalf("Failed to create port: %v", err)
-	}
-
-	req := testcontainers.ContainerRequest{
-		Image: "localstack/localstack:4.0.3",
-		Env: map[string]string{
-			"AWS_ACCESS_KEY_ID":      provider.Value.AccessKeyID,
-			"AWS_SECRET_ACCESS_KEY":  provider.Value.SecretAccessKey,
-			"AWS_SESSION_TOKEN":      provider.Value.SessionToken,
-			"AWS_DEFAULT_REGION":     region,
-			"AWS_REGION":             region,
-			"SERVICES":               "sqs",
-			"SKIP_SSL_CERT_DOWNLOAD": "1",
-			"LOCALSTACK_HOST":        alias,
-			"SQS_ENDPOINT_STRATEGY":  "path",
-		},
-		Networks: []string{network.Name},
-		NetworkAliases: map[string][]string{
-			network.Name: {alias},
-		},
-		ExposedPorts: []string{port.Port()},
-		WaitingFor:   wait.ForListeningPort(port),
-	}
-
-	container, err := testcontainers.GenericContainer(ctx, testcontainers.GenericContainerRequest{
-		ContainerRequest: req,
-		Started:          true,
-	})
-	if err != nil {
-		t.Fatalf("Failed to start container: %v", err)
-	}
-
-	host, err := container.Host(ctx)
-	if err != nil {
-		t.Fatalf("Failed to extract host: %v", err)
-	}
-	mappedPort, err := container.MappedPort(ctx, port)
-	if err != nil {
-		t.Fatalf("Failed to extract port: %v", err)
-	}
-
-	endpoint := fmt.Sprintf("http://%s:%s", host, mappedPort.Port())
-	cfg, err := config.LoadDefaultConfig(ctx, config.WithRegion(region), config.WithCredentialsProvider(provider), config.WithBaseEndpoint(endpoint))
-	if err != nil {
-		t.Fatal(err)
-	}
-
-	client := sqs.NewFromConfig(cfg)
-
-	queueM, err := client.CreateQueue(ctx, &sqs.CreateQueueInput{
-		QueueName: aws.String(queueName),
-	})
-	if err != nil {
-		t.Fatal(err)
-	}
-
-	endpoint = fmt.Sprintf("http://%s:%s", alias, port.Port())
-	queueC := &queueConfig{
-		URL:         *queueM.QueueUrl,
-		Region:      region,
-		Credentials: provider,
-		Endpoint:    endpoint,
-	}
-
-	return &queue{
-		Container: &container,
-		Client:    client,
-		Config:    queueC,
-	}
-}
-
-type message struct {
-	URL  string
-	Type string
-	Body string
-}
-
-func sendMessage(t *testing.T, ctx context.Context, queue *sqs.Client, message message) {
-	_, err := queue.SendMessage(ctx, &sqs.SendMessageInput{
-		MessageAttributes: map[string]types.MessageAttributeValue{
-			"type": {
-				DataType:    aws.String("String"),
-				StringValue: aws.String(message.Type),
-			},
-		},
-		QueueUrl:    aws.String(message.URL),
-		MessageBody: aws.String(message.Body),
-	})
-	if err != nil {
-		t.Fatalf("Failed to send message: %v", err)
-	}
-}
-
-func assertMessageWait(t *testing.T, ctx context.Context, queue *queue, msgType string) {
-	time.Sleep(time.Second)
-
-	result, err := queue.Client.ReceiveMessage(ctx, &sqs.ReceiveMessageInput{
-		QueueUrl:            &queue.Config.URL,
-		MaxNumberOfMessages: 1,
-		WaitTimeSeconds:     5,
-		VisibilityTimeout:   2,
-		MessageAttributeNames: []string{
-			"type",
-		},
-	})
-
-	assert.Nil(t, err)
-	for _, message := range result.Messages {
-		assert.Equal(t, msgType, *message.MessageAttributes["type"].StringValue)
-	}
-}
-
-func createQueueDependencies(t *testing.T, ctx context.Context, serviceName string) (*queue, func()) {
-	network := createNetwork(t, ctx)
-
-	dbconfig, dbContainer := createDB(t, ctx, network)
-
-	queue := createQueue(t, ctx, network)
-
-	container := createQueueContainer(t, ctx, &queueContainerConfig{
-		ServiceName: serviceName,
-		Queue:       *queue.Config,
-		DB:          dbconfig,
-		Network:     network,
-	})
-
-	return queue, func() {
-		testcontainers.CleanupNetwork(t, network)
-
-		err := dbContainer.Terminate(ctx)
-		if err != nil {
-			t.Error(err)
-		}
-
-		err = (*queue.Container).Terminate(ctx)
-		if err != nil {
-			t.Error(err)
-		}
-
-		err = container.Terminate(ctx)
-		if err != nil {
-			t.Error(err)
-		}
-	}
-}