test_helper.go

Documentation: github.com/ory/x/sqlcon/dockertest

     1  package dockertest
     2  
     3  import (
     4  	"context"
     5  	"fmt"
     6  	"io/ioutil"
     7  	"log"
     8  	"os"
     9  	"regexp"
    10  	"strings"
    11  	"sync"
    12  	"testing"
    13  	"time"
    14  
    15  	"github.com/ory/x/stringsx"
    16  
    17  	"github.com/gobuffalo/pop/v5"
    18  
    19  	"github.com/jmoiron/sqlx"
    20  	"github.com/pkg/errors"
    21  	"github.com/stretchr/testify/require"
    22  
    23  	"github.com/docker/docker/api/types"
    24  	"github.com/docker/docker/api/types/filters"
    25  	"github.com/docker/docker/client"
    26  
    27  	"github.com/ory/dockertest/v3"
    28  
    29  	"github.com/ory/x/logrusx"
    30  	"github.com/ory/x/resilience"
    31  )
    32  
    33  // atexit := atexit.NewOnExit()
    34  // atexit.Add(func() {
    35  //	dockertest.KillAll()
    36  // })
    37  // atexit.Exit(testMain(m))
    38  
    39  // func WrapCleanup
    40  
    41  var resources = []*dockertest.Resource{}
    42  var pool *dockertest.Pool
    43  
    44  // KillAllTestDatabases deletes all test databases.
    45  func KillAllTestDatabases() {
    46  	pool, err := dockertest.NewPool("")
    47  	if err != nil {
    48  		panic(err)
    49  	}
    50  
    51  	for _, r := range resources {
    52  		if err := pool.Purge(r); err != nil {
    53  			panic(err)
    54  		}
    55  	}
    56  }
    57  
    58  // Register sets up OnExit.
    59  func Register() *OnExit {
    60  	onexit := NewOnExit()
    61  	onexit.Add(func() {
    62  		KillAllTestDatabases()
    63  	})
    64  	return onexit
    65  }
    66  
    67  // Parallel runs tasks in parallel.
    68  func Parallel(fs []func()) {
    69  	wg := sync.WaitGroup{}
    70  
    71  	wg.Add(len(fs))
    72  	for _, f := range fs {
    73  		go func(ff func()) {
    74  			defer wg.Done()
    75  			ff()
    76  		}(f)
    77  	}
    78  
    79  	wg.Wait()
    80  }
    81  
    82  func connect(dialect, driver, dsn string) (db *sqlx.DB, err error) {
    83  	if scheme := strings.Split(dsn, "://")[0]; scheme == "mysql" {
    84  		dsn = strings.Replace(dsn, "mysql://", "", -1)
    85  	} else if scheme == "cockroach" {
    86  		dsn = strings.Replace(dsn, "cockroach://", "postgres://", 1)
    87  	}
    88  	err = resilience.Retry(
    89  		logrusx.New("", ""),
    90  		time.Second*5,
    91  		time.Minute*5,
    92  		func() (err error) {
    93  			db, err = sqlx.Open(dialect, dsn)
    94  			if err != nil {
    95  				log.Printf("Connecting to database %s failed: %s", driver, err)
    96  				return err
    97  			}
    98  
    99  			if err := db.Ping(); err != nil {
   100  				log.Printf("Pinging database %s failed: %s", driver, err)
   101  				return err
   102  			}
   103  
   104  			return nil
   105  		},
   106  	)
   107  	if err != nil {
   108  		return nil, errors.Errorf("Unable to connect to %s (%s): %s", driver, dsn, err)
   109  	}
   110  	log.Printf("Connected to database %s", driver)
   111  	return db, nil
   112  }
   113  
   114  func connectPop(t require.TestingT, url string) (c *pop.Connection) {
   115  	require.NoError(t, resilience.Retry(logrusx.New("", ""), time.Second*5, time.Minute*5, func() error {
   116  		var err error
   117  		c, err = pop.NewConnection(&pop.ConnectionDetails{
   118  			URL: url,
   119  		})
   120  		if err != nil {
   121  			log.Printf("could not create pop connection")
   122  			return err
   123  		}
   124  		if err := c.Open(); err != nil {
   125  			// an Open error probably means we have a problem with the connections config
   126  			log.Printf("could not open pop connection: %+v", err)
   127  			return err
   128  		}
   129  		return c.RawQuery("select version()").Exec()
   130  	}))
   131  	return
   132  }
   133  
   134  // ## PostgreSQL ##
   135  
   136  func startPostgreSQL() (*dockertest.Resource, error) {
   137  	pool, err := dockertest.NewPool("")
   138  	if err != nil {
   139  		return nil, errors.Wrap(err, "Could not connect to docker")
   140  	}
   141  
   142  	resource, err := pool.Run("postgres", "11.8", []string{"POSTGRES_PASSWORD=secret", "POSTGRES_DB=postgres"})
   143  	if err == nil {
   144  		resources = append(resources, resource)
   145  	}
   146  	return resource, err
   147  }
   148  
   149  // RunTestPostgreSQL runs a PostgreSQL database and returns the URL to it.
   150  func RunTestPostgreSQL(t testing.TB) string {
   151  	if dsn := os.Getenv("TEST_DATABASE_POSTGRESQL"); dsn != "" {
   152  		t.Logf("Skipping Docker setup because environment variable TEST_DATABASE_POSTGRESQL is set to: %s", dsn)
   153  		return dsn
   154  	}
   155  
   156  	u, err := RunPostgreSQL()
   157  	require.NoError(t, err)
   158  
   159  	return u
   160  }
   161  
   162  // RunPostgreSQL runs a PostgreSQL database and returns the URL to it.
   163  func RunPostgreSQL() (string, error) {
   164  	resource, err := startPostgreSQL()
   165  	if err != nil {
   166  		return "", err
   167  	}
   168  
   169  	return fmt.Sprintf("postgres://postgres:secret@127.0.0.1:%s/postgres?sslmode=disable", resource.GetPort("5432/tcp")), nil
   170  }
   171  
   172  // ConnectToTestPostgreSQL connects to a PostgreSQL database.
   173  func ConnectToTestPostgreSQL() (*sqlx.DB, error) {
   174  	if dsn := os.Getenv("TEST_DATABASE_POSTGRESQL"); dsn != "" {
   175  		return connect("pgx", "postgres", dsn)
   176  	}
   177  
   178  	resource, err := startPostgreSQL()
   179  	if err != nil {
   180  		return nil, errors.Wrap(err, "Could not start resource")
   181  	}
   182  
   183  	db := bootstrap("postgres://postgres:secret@localhost:%s/postgres?sslmode=disable", "5432/tcp", "pgx", pool, resource)
   184  	return db, nil
   185  }
   186  
   187  func ConnectToTestPostgreSQLPop(t testing.TB) *pop.Connection {
   188  	url := RunTestPostgreSQL(t)
   189  	return connectPop(t, url)
   190  }
   191  
   192  // ## MySQL ##
   193  
   194  func startMySQL() (*dockertest.Resource, error) {
   195  	pool, err := dockertest.NewPool("")
   196  	if err != nil {
   197  		return nil, errors.Wrap(err, "Could not connect to docker")
   198  	}
   199  
   200  	resource, err := pool.Run("mysql", "8.0", []string{"MYSQL_ROOT_PASSWORD=secret"})
   201  	if err == nil {
   202  		resources = append(resources, resource)
   203  	}
   204  	return resource, err
   205  }
   206  
   207  // RunMySQL runs a RunMySQL database and returns the URL to it.
   208  func RunMySQL() (string, error) {
   209  	resource, err := startMySQL()
   210  	if err != nil {
   211  		return "", err
   212  	}
   213  
   214  	return fmt.Sprintf("mysql://root:secret@(localhost:%s)/mysql?parseTime=true&multiStatements=true", resource.GetPort("3306/tcp")), nil
   215  }
   216  
   217  // RunTestMySQL runs a MySQL database and returns the URL to it.
   218  func RunTestMySQL(t testing.TB) string {
   219  	if dsn := os.Getenv("TEST_DATABASE_MYSQL"); dsn != "" {
   220  		t.Logf("Skipping Docker setup because environment variable TEST_DATABASE_MYSQL is set to: %s", dsn)
   221  		return dsn
   222  	}
   223  
   224  	u, err := RunMySQL()
   225  	require.NoError(t, err)
   226  
   227  	return u
   228  }
   229  
   230  // ConnectToTestMySQL connects to a MySQL database.
   231  func ConnectToTestMySQL() (*sqlx.DB, error) {
   232  	if dsn := os.Getenv("TEST_DATABASE_MYSQL"); dsn != "" {
   233  		log.Println("Found mysql test database config, skipping dockertest...")
   234  		return connect("mysql", "mysql", dsn)
   235  	}
   236  
   237  	resource, err := startMySQL()
   238  	if err != nil {
   239  		return nil, errors.Wrap(err, "Could not start resource")
   240  	}
   241  
   242  	db := bootstrap("root:secret@(localhost:%s)/mysql?parseTime=true", "3306/tcp", "mysql", pool, resource)
   243  	return db, nil
   244  }
   245  
   246  func ConnectToTestMySQLPop(t testing.TB) *pop.Connection {
   247  	url := RunTestMySQL(t)
   248  	return connectPop(t, url)
   249  }
   250  
   251  // ## CockroachDB
   252  
   253  func startCockroachDB(version string) (*dockertest.Resource, error) {
   254  	pool, err := dockertest.NewPool("")
   255  	if err != nil {
   256  		return nil, errors.Wrap(err, "Could not connect to docker")
   257  	}
   258  
   259  	resource, err := pool.RunWithOptions(&dockertest.RunOptions{
   260  		Repository: "cockroachdb/cockroach",
   261  		Tag:        stringsx.Coalesce(version, "v20.2.5"),
   262  		Cmd:        []string{"start-single-node", "--insecure"},
   263  	})
   264  	if err == nil {
   265  		resources = append(resources, resource)
   266  	}
   267  	return resource, err
   268  }
   269  
   270  // RunCockroachDB runs a CockroachDB database and returns the URL to it.
   271  func RunCockroachDB() (string, error) {
   272  	return RunCockroachDBWithVersion("")
   273  }
   274  
   275  // RunCockroachDB runs a CockroachDB database and returns the URL to it.
   276  func RunCockroachDBWithVersion(version string) (string, error) {
   277  	resource, err := startCockroachDB(version)
   278  	if err != nil {
   279  		return "", err
   280  	}
   281  
   282  	return fmt.Sprintf("cockroach://root@localhost:%s/defaultdb?sslmode=disable", resource.GetPort("26257/tcp")), nil
   283  }
   284  
   285  // RunTestCockroachDB runs a CockroachDB database and returns the URL to it.
   286  func RunTestCockroachDB(t testing.TB) string {
   287  	return RunTestCockroachDBWithVersion(t, "")
   288  }
   289  
   290  // RunTestCockroachDB runs a CockroachDB database and returns the URL to it.
   291  func RunTestCockroachDBWithVersion(t testing.TB, version string) string {
   292  	if dsn := os.Getenv("TEST_DATABASE_COCKROACHDB"); dsn != "" {
   293  		t.Logf("Skipping Docker setup because environment variable TEST_DATABASE_COCKROACHDB is set to: %s", dsn)
   294  		return dsn
   295  	}
   296  
   297  	u, err := RunCockroachDBWithVersion(version)
   298  	require.NoError(t, err)
   299  
   300  	return u
   301  }
   302  
   303  // ConnectToTestCockroachDB connects to a CockroachDB database.
   304  func ConnectToTestCockroachDB() (*sqlx.DB, error) {
   305  	if dsn := os.Getenv("TEST_DATABASE_COCKROACHDB"); dsn != "" {
   306  		log.Println("Found cockroachdb test database config, skipping dockertest...")
   307  		return connect("pgx", "cockroach", dsn)
   308  	}
   309  
   310  	resource, err := startCockroachDB("")
   311  	if err != nil {
   312  		return nil, errors.Wrap(err, "Could not start resource")
   313  	}
   314  
   315  	db := bootstrap("postgres://root@localhost:%s/defaultdb?sslmode=disable", "26257/tcp", "pgx", pool, resource)
   316  	return db, nil
   317  }
   318  
   319  func ConnectToTestCockroachDBPop(t testing.TB) *pop.Connection {
   320  	url := RunTestCockroachDB(t)
   321  	return connectPop(t, url)
   322  }
   323  
   324  func bootstrap(u, port, d string, pool *dockertest.Pool, resource *dockertest.Resource) (db *sqlx.DB) {
   325  	if err := resilience.Retry(logrusx.New("", ""), time.Second*5, time.Minute*5, func() error {
   326  		var err error
   327  		db, err = sqlx.Open(d, fmt.Sprintf(u, resource.GetPort(port)))
   328  		if err != nil {
   329  			return err
   330  		}
   331  
   332  		return db.Ping()
   333  	}); err != nil {
   334  		if pErr := pool.Purge(resource); pErr != nil {
   335  			log.Fatalf("Could not connect to docker and unable to remove image: %s - %s", err, pErr)
   336  		}
   337  		log.Fatalf("Could not connect to docker: %s", err)
   338  	}
   339  	return
   340  }
   341  
   342  var comments = regexp.MustCompile("(--[^\n]*\n)|(?s:/\\*.+\\*/)")
   343  
   344  func StripDump(d string) string {
   345  	d = comments.ReplaceAllLiteralString(d, "")
   346  	d = strings.TrimPrefix(d, "Command \"dump\" is deprecated, cockroach dump will be removed in a subsequent release.\r\nFor details, see: https://github.com/cockroachdb/cockroach/issues/54040\r\n")
   347  	d = strings.ReplaceAll(d, "\r\n", "")
   348  	d = strings.ReplaceAll(d, "\t", " ")
   349  	d = strings.ReplaceAll(d, "\n", " ")
   350  	return d
   351  }
   352  
   353  func DumpSchema(ctx context.Context, t *testing.T, db string) string {
   354  	var containerPort string
   355  	var cmd []string
   356  	cases := stringsx.RegisteredCases{}
   357  	switch db {
   358  	case cases.AddCase("postgres"):
   359  		containerPort = "5432"
   360  		cmd = []string{"pg_dump", "-U", "postgres", "-s", "-T", "hydra_*_migration", "-T", "schema_migration"}
   361  	case cases.AddCase("mysql"):
   362  		containerPort = "3306"
   363  		cmd = []string{"/usr/bin/mysqldump", "-u", "root", "--password=secret", "mysql"}
   364  	case cases.AddCase("cockroach"):
   365  		containerPort = "26257"
   366  		cmd = []string{"./cockroach", "dump", "defaultdb", "--insecure", "--dump-mode=schema"}
   367  	default:
   368  		t.Log(cases.ToUnknownCaseErr(db))
   369  		t.FailNow()
   370  		return ""
   371  	}
   372  
   373  	cli, err := client.NewClientWithOpts(client.FromEnv)
   374  	require.NoError(t, err)
   375  	containers, err := cli.ContainerList(ctx, types.ContainerListOptions{
   376  		Quiet:   true,
   377  		Filters: filters.NewArgs(filters.Arg("expose", containerPort)),
   378  	})
   379  	require.NoError(t, err)
   380  
   381  	if len(containers) != 1 {
   382  		t.Logf("Ambiguous amount of %s containers: %d", db, len(containers))
   383  		t.FailNow()
   384  	}
   385  
   386  	process, err := cli.ContainerExecCreate(ctx, containers[0].ID, types.ExecConfig{
   387  		Tty:          true,
   388  		AttachStdout: true,
   389  		Cmd:          cmd,
   390  	})
   391  	require.NoError(t, err)
   392  
   393  	resp, err := cli.ContainerExecAttach(ctx, process.ID, types.ExecStartCheck{
   394  		Tty: true,
   395  	})
   396  	require.NoError(t, err)
   397  	dump, err := ioutil.ReadAll(resp.Reader)
   398  	require.NoError(t, err, "%s", dump)
   399  
   400  	return StripDump(string(dump))
   401  }
   402
View as plain text