github.com/treeverse/lakefs@v1.24.1-0.20240520134607-95648127bfb0/pkg/kv/postgres/main_test.go (about)

     1  package postgres_test
     2  
     3  import (
     4  	"context"
     5  	"fmt"
     6  	"log"
     7  	"net/url"
     8  	"os"
     9  	"testing"
    10  
    11  	"github.com/jackc/pgx/v5/pgxpool"
    12  	"github.com/ory/dockertest/v3"
    13  	"github.com/treeverse/lakefs/pkg/testutil"
    14  )
    15  
    16  const (
    17  	dbContainerTimeoutSeconds = 10 * 60 // 10 min
    18  )
    19  
    20  var (
    21  	pool *dockertest.Pool
    22  )
    23  
    24  func runDBInstance(dockerPool *dockertest.Pool, dbName string) (string, func()) {
    25  	ctx := context.Background()
    26  	resource, err := dockerPool.Run("postgres", "11", []string{
    27  		"POSTGRES_USER=lakefs",
    28  		"POSTGRES_PASSWORD=lakefs",
    29  		"POSTGRES_DB=" + dbName,
    30  	})
    31  	if err != nil {
    32  		panic("Could not start postgresql: " + err.Error())
    33  	}
    34  
    35  	// set cleanup
    36  	closer := func() {
    37  		err := dockerPool.Purge(resource)
    38  		if err != nil {
    39  			panic("could not kill postgres container")
    40  		}
    41  	}
    42  
    43  	// expire, just to make sure
    44  	err = resource.Expire(dbContainerTimeoutSeconds)
    45  	if err != nil {
    46  		panic("could not expire postgres container")
    47  	}
    48  
    49  	// create connection
    50  	var pgPool *pgxpool.Pool
    51  	port := resource.GetPort("5432/tcp")
    52  	uri := fmt.Sprintf("postgres://lakefs:lakefs@localhost:%s/%s?sslmode=disable", port, url.PathEscape(dbName))
    53  	err = dockerPool.Retry(func() error {
    54  		var err error
    55  		pgPool, err = pgxpool.New(ctx, uri)
    56  		if err != nil {
    57  			return err
    58  		}
    59  		return testutil.PingPG(ctx, pgPool)
    60  	})
    61  	if err != nil {
    62  		panic("could not connect to postgres: " + err.Error())
    63  	}
    64  	pgPool.Close()
    65  
    66  	// return DB URI
    67  	return uri, closer
    68  }
    69  
    70  func TestMain(m *testing.M) {
    71  	var err error
    72  	pool, err = dockertest.NewPool("")
    73  	if err != nil {
    74  		log.Fatalf("Could not connect to Docker: %s", err)
    75  	}
    76  	code := m.Run()
    77  	os.Exit(code)
    78  }