Merge pull request #8 from opsmatic/customclient

Added HTTP client override support

example_test.go

@@ -6,8 +6,11 @@ package goes_test
 
 import (
 	"fmt"
-	"github.com/belogik/goes"
+	"net/http"
 	"net/url"
+	"time"
+
+	"github.com/belogik/goes"
 )
 
 func ExampleConnection_CreateIndex() {
@@ -86,7 +89,7 @@ func ExampleConnection_Search() {
 			},
 		},
 	}
-	
+
 	extraArgs := make(url.Values, 1)
 
 	searchResults, err := conn.Search(query, []string{"someindex"}, []string{""}, extraArgs)
@@ -143,3 +146,15 @@ func ExampleConnection_Delete() {
 
 	fmt.Printf("%s", response)
 }
+
+func ExampleConnectionOverrideHttpClient() {
+	tr := &http.Transport{
+		ResponseHeaderTimeout: 1 * time.Second,
+	}
+	cl := &http.Client{
+		Transport: tr,
+	}
+	conn := goes.NewConnection("localhost", "9200").WithClient(cl)
+
+	fmt.Printf("%v\n", conn.Client)
+}

goes.go

@@ -31,7 +31,12 @@ func (err *SearchError) Error() string {
 // This function is pretty useless for now but might be useful in a near future
 // if wee need more features like connection pooling or load balancing.
 func NewConnection(host string, port string) *Connection {
-	return &Connection{host, port}
+	return &Connection{host, port, http.DefaultClient}
+}
+
+func (c *Connection) WithClient(cl *http.Client) *Connection {
+	c.Client = cl
+	return c
 }
 
 // CreateIndex creates a new index represented by a name and a mapping
@@ -285,8 +290,6 @@ func (req *Request) Run() (Response, error) {
 
 	reader := bytes.NewReader(postData)
 
-	client := http.DefaultClient
-
 	newReq, err := http.NewRequest(req.method, req.Url(), reader)
 	if err != nil {
 		return Response{}, err
@@ -296,7 +299,7 @@ func (req *Request) Run() (Response, error) {
 		newReq.Header.Set("Content-Type", "application/x-www-form-urlencoded")
 	}
 
-	resp, err := client.Do(newReq)
+	resp, err := req.Conn.Client.Do(newReq)
 	if err != nil {
 		return Response{}, err
 	}

goes_test.go

@@ -6,6 +6,7 @@ package goes
 
 import (
 	. "launchpad.net/gocheck"
+	"net/http"
 	"net/url"
 	"os"
 	"testing"
@@ -38,7 +39,22 @@ func (s *GoesTestSuite) SetUpTest(c *C) {
 
 func (s *GoesTestSuite) TestNewConnection(c *C) {
 	conn := NewConnection(ES_HOST, ES_PORT)
-	c.Assert(conn, DeepEquals, &Connection{ES_HOST, ES_PORT})
+	c.Assert(conn, DeepEquals, &Connection{ES_HOST, ES_PORT, http.DefaultClient})
+}
+
+func (s *GoesTestSuite) TestWithClient(c *C) {
+	tr := &http.Transport{
+		DisableCompression:    true,
+		ResponseHeaderTimeout: 1 * time.Second,
+	}
+	cl := &http.Client{
+		Transport: tr,
+	}
+	conn := NewConnection(ES_HOST, ES_PORT).WithClient(cl)
+
+	c.Assert(conn, DeepEquals, &Connection{ES_HOST, ES_PORT, cl})
+	c.Assert(conn.Client.Transport.(*http.Transport).DisableCompression, Equals, true)
+	c.Assert(conn.Client.Transport.(*http.Transport).ResponseHeaderTimeout, Equals, 1*time.Second)
 }
 
 func (s *GoesTestSuite) TestUrl(c *C) {

structs.go

@@ -5,6 +5,7 @@
 package goes
 
 import (
+	"net/http"
 	"net/url"
 )
 
@@ -15,6 +16,10 @@ type Connection struct {
 
 	// The port to use
 	Port string
+
+	// Client is the http client used to make requests, allowing settings things
+	// such as timeouts etc
+	Client *http.Client
 }
 
 // Represents a Request to elasticsearch