added elasticsearch support

docker.go

@@ -23,13 +23,15 @@ import (
 	"fmt"
 	"log"
 	"os/exec"
+	"strconv"
 	"strings"
 	"time"
 
 	"database/sql"
 	"math/rand"
 	"regexp"
 
+	"github.com/mattbaird/elastigo/lib"
 	"github.com/pborman/uuid"
 	"gopkg.in/mgo.v2"
 
@@ -260,6 +262,21 @@ func SetupPostgreSQLContainer() (c ContainerID, ip string, port int, err error)
 	return
 }
 
+// SetupElasticSearchContainer sets up a real ElasticSearch instance for testing purposes
+// using a Docker container. It returns the container ID and its IP address,
+// or makes the test fail on error.
+func SetupElasticSearchContainer() (c ContainerID, ip string, port int, err error) {
+	port = randInt(1024, 49150)
+	forward := fmt.Sprintf("%d:%d", port, 9200)
+	if BindDockerToLocalhost != "" {
+		forward = "127.0.0.1:" + forward
+	}
+	c, ip, err = setupContainer(elasticsearchImage, port, 15*time.Second, func() (string, error) {
+		return run("--name", uuid.New(), "-d", "-P", "-p", forward, elasticsearchImage)
+	})
+	return
+}
+
 // OpenPostgreSQLContainerConnection is supported for legacy reasons. Don't use it.
 func OpenPostgreSQLContainerConnection(tries int, delay time.Duration) (c ContainerID, db *sql.DB, err error) {
 	c, ip, port, err := SetupPostgreSQLContainer()
@@ -329,6 +346,33 @@ func OpenMySQLContainerConnection(tries int, delay time.Duration) (c ContainerID
 	return c, nil, errors.New("Could not set up MySQL container.")
 }
 
+// OpenElasticSearchContainerConnection is supported for legacy reasons. Don't use it.
+func OpenElasticSearchContainerConnection(tries int, delay time.Duration) (c ContainerID, con *elastigo.Conn, err error) {
+	c, ip, port, err := SetupElasticSearchContainer()
+	if err != nil {
+		return c, nil, fmt.Errorf("Could not set up ElasticSearch container: %v", err)
+	}
+
+	for try := 0; try <= tries; try++ {
+		time.Sleep(delay)
+		url := fmt.Sprintf("%s:%d", ip, port)
+		log.Printf("Try %d: Connecting %s", try, url)
+
+		conn := elastigo.NewConn()
+		conn.Domain = ip
+		conn.Port = strconv.Itoa(port)
+
+		resp, err := conn.Health("")
+		if err == nil && resp.Status != "" {
+			log.Printf("Try %d: Successfully connected to %v", try, conn.Domain)
+			return c, conn, nil
+		}
+
+		log.Printf("Try %d: Could not set up ElasticSearch container: %v", try, err)
+	}
+	return c, nil, errors.New("Could not set up ElasticSearch container.")
+}
+
 // ConnectToPostgreSQL starts a PostgreSQL image and passes the database url to the connector callback.
 func ConnectToPostgreSQL(tries int, delay time.Duration, connector func(url string) bool) (c ContainerID, err error) {
 	c, ip, port, err := SetupPostgreSQLContainer()
@@ -384,3 +428,22 @@ func ConnectToMySQL(tries int, delay time.Duration, connector func(url string) b
 	}
 	return c, errors.New("Could not set up MySQL container.")
 }
+
+// ConnectToElasticSearch starts an ElasticSearch image and passes the database url to the connector callback function.
+// The url will match the ip:port pattern (e.g. 123.123.123.123:4241)
+func ConnectToElasticSearch(tries int, delay time.Duration, connector func(url string) bool) (c ContainerID, err error) {
+	c, ip, port, err := SetupElasticSearchContainer()
+	if err != nil {
+		return c, fmt.Errorf("Could not set up ElasticSearch container: %v", err)
+	}
+
+	for try := 0; try <= tries; try++ {
+		time.Sleep(delay)
+		url := fmt.Sprintf("%s:%d", ip, port)
+		if connector(url) {
+			return c, nil
+		}
+		log.Printf("Try %d failed. Retrying.", try)
+	}
+	return c, errors.New("Could not set up ElasticSearch container.")
+}

docker_test.go

@@ -2,11 +2,15 @@ package dockertest_test
 
 import (
 	"database/sql"
-	. "github.com/ory-am/dockertest"
-	"github.com/stretchr/testify/require"
-	"gopkg.in/mgo.v2"
+	"strings"
 	"testing"
 	"time"
+
+	"gopkg.in/mgo.v2"
+
+	"github.com/mattbaird/elastigo/lib"
+	. "github.com/ory-am/dockertest"
+	"github.com/stretchr/testify/require"
 )
 
 func TestOpenPostgreSQLContainerConnection(t *testing.T) {
@@ -37,6 +41,16 @@ func TestOpenMongoDBContainerConnection(t *testing.T) {
 	defer db.Close()
 }
 
+func TestOpenElasticSearchContainerConnection(t *testing.T) {
+	c, conn, err := OpenElasticSearchContainerConnection(15, time.Millisecond*500)
+	require.Nil(t, err)
+	defer c.KillRemove()
+	require.NotNil(t, conn)
+	_, err = conn.Health("")
+	require.Nil(t, err)
+	defer conn.Close()
+}
+
 func TestConnectToPostgreSQL(t *testing.T) {
 	c, err := ConnectToPostgreSQL(15, time.Millisecond*500, func(url string) bool {
 		db, err := sql.Open("postgres", url)
@@ -75,3 +89,27 @@ func TestConnectToMongoDB(t *testing.T) {
 	require.Nil(t, err)
 	defer c.KillRemove()
 }
+
+func TestConnectToElasticSearch(t *testing.T) {
+	c, err := ConnectToElasticSearch(15, time.Millisecond*500, func(url string) bool {
+		segs := strings.Split(url, ":")
+		if len(segs) != 2 {
+			return false
+		}
+
+		conn := elastigo.NewConn()
+		conn.Domain = segs[0]
+		conn.Port = segs[1]
+		resp, err := conn.Health()
+		if err != nil {
+			return false
+		}
+		if resp.Status != "green" {
+			return false
+		}
+		defer conn.Close()
+		return true
+	})
+	require.Nil(t, err)
+	defer c.KillRemove()
+}

vars.go

@@ -20,9 +20,10 @@ var (
 )
 
 const (
-	mongoImage    = "mongo"
-	mysqlImage    = "mysql"
-	postgresImage = "postgres"
+	mongoImage         = "mongo"
+	mysqlImage         = "mysql"
+	postgresImage      = "postgres"
+	elasticsearchImage = "elasticsearch"
 
 	// MySQLUsername must be passed as username when connecting to mysql
 	MySQLUsername = "root"