Merge pull request #34 from hsinhoyeh/issue/33

allow mysql to use customized database

README.md

@@ -37,6 +37,7 @@ Dockertest ships with support for these backends:
   - [My build is broken!](#my-build-is-broken)
   - [Out of disk space](#out-of-disk-space)
   - [Removing old containers](#removing-old-containers)
+  - [Customized database] (#Customized-database)
 
 <!-- END doctoc generated TOC please keep comment here to allow auto update -->
 
@@ -225,4 +226,31 @@ Try cleaning up the images with [docker-cleanup-volumes](https://github.com/chad
 Sometimes container clean up fails. Check out
 [this stackoverflow question](http://stackoverflow.com/questions/21398087/how-to-delete-dockers-images) on how to fix this.
 
+### Customized database
+
+I am using postgres (or mysql) driver, how do I use customized database instead of default one?
+You can alleviate this helper function to do that, see testcase or example below:
+
+```go
+
+func TestMain(m *testing.M) {
+	if c, err := dockertest.ConnectToPostgreSQL(15, time.Second, func(url string) bool {
+        customizedDB := "cherry" // here I am connecting cherry database
+        newURL, err := SetUpPostgreDatabase(customizedDB, url)
+
+        // or use SetUpMysqlDatabase for mysql driver
+
+        if err != nil {
+                log.Fatal(err)
+        }
+        db, err := sql.Open("postgres", newURL)
+        if err != nil {
+            return false
+        }
+        return db.Ping() == nil
+    }); err != nil {
+        log.Fatal(err)
+    }
+```
+
 *Thanks to our sponsors: Ory GmbH & Imarum GmbH*

docker_integration_test.go

@@ -48,6 +48,26 @@ func TestConnectToPostgreSQL(t *testing.T) {
 	defer c.KillRemove()
 }
 
+func TestConnectToPostgreSQLWithCustomizedDB(t *testing.T) {
+	c, err := ConnectToPostgreSQL(15, time.Millisecond*500, func(url string) bool {
+		customizedDB := "db0001"
+		gotURL, err := SetUpPostgreDatabase(customizedDB, url)
+		if err != nil {
+			return false
+		}
+		assert.True(t, strings.Contains(gotURL, customizedDB),
+			fmt.Sprintf("url(%s) should contains tag(%s)", gotURL, customizedDB))
+		db, err := sql.Open("postgres", gotURL)
+		if err != nil {
+			return false
+		}
+		defer db.Close()
+		return true
+	})
+	assert.Nil(t, err)
+	defer c.KillRemove()
+}
+
 func TestConnectToRabbitMQ(t *testing.T) {
 	c, err := ConnectToRabbitMQ(15, time.Millisecond*500, func(url string) bool {
 		amqp, err := amqp.Dial(fmt.Sprintf("amqp://%v", url))
@@ -74,6 +94,26 @@ func TestConnectToMySQL(t *testing.T) {
 	defer c.KillRemove()
 }
 
+func TestConnectToMySQLWithCustomizedDB(t *testing.T) {
+	customizedDB := "db0001"
+	c, err := ConnectToMySQL(20, time.Second, func(url string) bool {
+		gotURL, err := SetUpMySQLDatabase(customizedDB, url)
+		if err != nil {
+			return false
+		}
+		assert.True(t, strings.Contains(gotURL, customizedDB),
+			fmt.Sprintf("url(%s) should contains tag(%s)", gotURL, customizedDB))
+		db, err := sql.Open("mysql", gotURL)
+		if err != nil {
+			return false
+		}
+		defer db.Close()
+		return true
+	})
+	assert.Nil(t, err)
+	defer c.KillRemove()
+}
+
 func TestConnectToMongoDB(t *testing.T) {
 	c, err := ConnectToMongoDB(15, time.Millisecond*500, func(url string) bool {
 		db, err := mgo.Dial(url)

mysql.go

@@ -1,10 +1,17 @@
 package dockertest
 
 import (
+	"database/sql"
 	"errors"
 	"fmt"
 	"log"
 	"time"
+
+	mysql "github.com/go-sql-driver/mysql"
+)
+
+const (
+	defaultMySQLDBName = "mysql"
 )
 
 // SetupMySQLContainer sets up a real MySQL instance for testing purposes,
@@ -40,3 +47,29 @@ func ConnectToMySQL(tries int, delay time.Duration, connector func(url string) b
 	}
 	return c, errors.New("Could not set up MySQL container.")
 }
+
+// SetUpMySQLDatabase connects mysql container with given $connectURL and also creates a new database named $databaseName
+// A modified url used to connect the created database will be returned
+func SetUpMySQLDatabase(databaseName, connectURL string) (url string, err error) {
+	if databaseName == defaultMySQLDBName {
+		return connectURL, nil
+	}
+
+	db, err := sql.Open("mysql", connectURL)
+	if err != nil {
+		return "", err
+	}
+	defer db.Close()
+	_, err = db.Exec(fmt.Sprintf("CREATE DATABASE IF NOT EXISTS %s", databaseName))
+	if err != nil {
+		return "", err
+	}
+
+	// parse dsn
+	config, err := mysql.ParseDSN(connectURL)
+	if err != nil {
+		return "", err
+	}
+	config.DBName = databaseName // overwrite database name
+	return config.FormatDSN(), nil
+}

postgres.go

@@ -1,10 +1,14 @@
 package dockertest
 
 import (
+	"database/sql"
 	"errors"
 	"fmt"
 	"log"
+	"net/url"
 	"time"
+
+	_ "github.com/lib/pq"
 )
 
 // SetupPostgreSQLContainer sets up a real PostgreSQL instance for testing purposes,
@@ -39,3 +43,38 @@ func ConnectToPostgreSQL(tries int, delay time.Duration, connector func(url stri
 	}
 	return c, errors.New("Could not set up PostgreSQL container.")
 }
+
+// SetUpPostgreDatabase connects postgre container with given $connectURL and also creates a new database named $databaseName
+// A modified url used to connect the created database will be returned
+func SetUpPostgreDatabase(databaseName, connectURL string) (modifiedURL string, err error) {
+	db, err := sql.Open("postgres", connectURL)
+	if err != nil {
+		return "", err
+	}
+	defer db.Close()
+
+	count := 0
+	err = db.QueryRow(
+		fmt.Sprintf("SELECT COUNT(*) FROM pg_catalog.pg_database WHERE datname = '%s' ;", databaseName)).
+		Scan(&count)
+	if err != nil {
+		return "", err
+	}
+	if count == 0 {
+		// not found for $databaseName, create it
+		_, err = db.Exec(fmt.Sprintf("CREATE DATABASE %s", databaseName))
+		if err != nil {
+			return "", err
+		}
+	}
+
+	// replace dbname in url
+	// from: postgres://postgres:[email protected]:9071/postgres?sslmode=disable
+	// to: postgres://postgres:[email protected]:9071/$databaseName?sslmode=disable
+	u, err := url.Parse(connectURL)
+	if err != nil {
+		return "", err
+	}
+	u.Path = fmt.Sprintf("/%s", databaseName)
+	return u.String(), nil
+}