allow max connections to be set on test db container (#35)

* allow max connections to be set on test db container Signed-off-by: Sarah Funkhouser <[email protected]> * no more magic numbers Signed-off-by: Sarah Funkhouser <[email protected]> --------- Signed-off-by: Sarah Funkhouser <[email protected]>
theopenlane · Nov 19, 2024 · 311368d · 311368d
1 parent 86273ac
commit 311368d
Showing 1 changed file with 66 additions and 16 deletions.
diff --git a/testutils/db.go b/testutils/db.go
@@ -16,21 +16,64 @@ import (
 	_ "modernc.org/sqlite"
 )
 
+const (
+	defaultMaxConn       = 100
+	defaultExpiryMinutes = 5
+)
+
+// Option is a functional option for configuring the DBConfig
+type Option func(*DBConfig)
+
+// TestFixture is a struct that holds the pool, resource and URI for the database
 type TestFixture struct {
 	Pool     *dockertest.Pool
 	resource *dockertest.Resource
 	URI      string
 	Dialect  string
 }
 
+// DBConfig is a struct that holds the configuration for the database
+type DBConfig struct {
+	// image is the docker image to use
+	image string
+
+	// expiry is the duration after which the container will be killed
+	expiry time.Duration
+
+	// maxConn is the maximum number of connections to the database
+	// defaults to 100
+	maxConn int
+}
+
+// WithImage sets the image used for docker tests
+func WithImage(image string) Option {
+	return func(c *DBConfig) {
+		c.image = image
+	}
+}
+
+// WithExpiryMinutes sets the expiration of the container, defaults to 10 minutes
+func WithExpiryMinutes(expiry int) Option {
+	return func(c *DBConfig) {
+		c.expiry = time.Duration(expiry) * time.Minute
+	}
+}
+
+// WithMaxConn sets the max connections to the db container, defaults to 100
+func WithMaxConn(maxConn int) Option {
+	return func(c *DBConfig) {
+		c.maxConn = maxConn
+	}
+}
+
 func TeardownFixture(tf *TestFixture) {
 	if tf.Pool != nil {
 		if err := tf.Pool.Purge(tf.resource); err != nil {
 			log.Fatalf("Could not purge resource: %s", err)
 		}
 	}
 }
-func GetPostgresDockerTest(image string, expiry time.Duration) (*TestFixture, error) {
+func GetPostgresDockerTest(image string, expiry time.Duration, maxConn int) (*TestFixture, error) {
 	pool, err := dockertest.NewPool("")
 	if err != nil {
 		return nil, err
@@ -53,6 +96,7 @@ func GetPostgresDockerTest(image string, expiry time.Duration) (*TestFixture, er
 		&dockertest.RunOptions{
 			Repository: defaultImg,
 			Tag:        imgTag,
+			Cmd:        []string{"postgres", "-c", fmt.Sprintf("max_connections=%d", maxConn)},
 			Env: []string{
 				fmt.Sprintf("POSTGRES_PASSWORD=%s", password),
 				fmt.Sprintf("POSTGRES_USER=%s", user),
@@ -106,32 +150,38 @@ func GetPostgresDockerTest(image string, expiry time.Duration) (*TestFixture, er
 	}, nil
 }
 
-func getTestDB(u string, expiry time.Duration) (*TestFixture, error) {
+func getTestDB(u string, expiry time.Duration, maxConn int) (*TestFixture, error) {
 	switch {
 	case strings.HasPrefix(u, "postgres"):
-		return GetPostgresDockerTest(u, expiry)
+		return GetPostgresDockerTest(u, expiry, maxConn)
 	default:
 		return nil, newURIError(u)
 	}
 }
 
 // GetTestURI returns the dialect, connection string and if used a testcontainer for database connectivity in tests
-func GetTestURI(u string, expiryMinutes int) *TestFixture {
+func GetTestURI(opts ...Option) *TestFixture {
+	config := &DBConfig{
+		maxConn: defaultMaxConn,                                    // default to 100
+		expiry:  time.Duration(defaultExpiryMinutes) * time.Minute, // default to 5 minutes
+	}
+
+	for _, opt := range opts {
+		opt(config)
+	}
+
 	switch {
-	case strings.HasPrefix(u, "sqlite://"):
+	case strings.HasPrefix(config.image, "sqlite://"):
 		// return dialect.SQLite, strings.TrimPrefix(u, "sqlite://")
-		return &TestFixture{Dialect: dialect.SQLite, URI: strings.TrimPrefix(u, "sqlite://")}
-	case strings.HasPrefix(u, "libsql://"):
+		return &TestFixture{Dialect: dialect.SQLite, URI: strings.TrimPrefix(config.image, "sqlite://")}
+	case strings.HasPrefix(config.image, "libsql://"):
 		// return dialect.SQLite, strings.TrimPrefix(u, "libsql://")
-		return &TestFixture{Dialect: "libsql", URI: strings.TrimPrefix(u, "libsql://")}
-	case strings.HasPrefix(u, "postgres://"), strings.HasPrefix(u, "postgresql://"):
+		return &TestFixture{Dialect: "libsql", URI: strings.TrimPrefix(config.image, "libsql://")}
+	case strings.HasPrefix(config.image, "postgres://"), strings.HasPrefix(config.image, "postgresql://"):
 		// return dialect.Postgres, u
-		return &TestFixture{Dialect: dialect.Postgres, URI: u}
-	case strings.HasPrefix(u, "docker://"):
-		// set reasonable expiry for docker test containers
-		expiry := time.Duration(expiryMinutes) * time.Minute
-
-		tf, err := getTestDB(strings.TrimPrefix(u, "docker://"), expiry)
+		return &TestFixture{Dialect: dialect.Postgres, URI: config.image}
+	case strings.HasPrefix(config.image, "docker://"):
+		tf, err := getTestDB(strings.TrimPrefix(config.image, "docker://"), config.expiry, config.maxConn)
 		if err != nil {
 			panic(err)
 		}
@@ -140,6 +190,6 @@ func GetTestURI(u string, expiryMinutes int) *TestFixture {
 
 		return tf
 	default:
-		panic("invalid DB URI, uri: " + u)
+		panic("invalid DB URI, uri: " + config.image)
 	}
 }