From 13f1b6d1fff291070c2aac725b644e3adcd9c651 Mon Sep 17 00:00:00 2001
From: Maciej Iwanowski <maciej.iwanowski@critical.today>
Date: Sat, 21 Jan 2017 21:53:59 +0100
Subject: [PATCH] Adding timeout to Snap REST client

---
 mgmt/rest/client/client.go           |  7 +++++++
 mgmt/rest/client/client_func_test.go | 17 +++++++++++++++++
 2 files changed, 24 insertions(+)

diff --git a/mgmt/rest/client/client.go b/mgmt/rest/client/client.go
index efc66aecf..50001a288 100644
--- a/mgmt/rest/client/client.go
+++ b/mgmt/rest/client/client.go
@@ -102,6 +102,13 @@ func Username(u string) metaOp {
 	}
 }
 
+//Timeout is an option that can be provided to the func client.New in order to set HTTP connection timeout.
+func Timeout(t time.Duration) metaOp {
+	return func(c *Client) {
+		c.http.Timeout = t
+	}
+}
+
 var (
 	secureTransport = &http.Transport{
 		TLSClientConfig: &tls.Config{InsecureSkipVerify: false},
diff --git a/mgmt/rest/client/client_func_test.go b/mgmt/rest/client/client_func_test.go
index 35b48377b..2c861e6d0 100644
--- a/mgmt/rest/client/client_func_test.go
+++ b/mgmt/rest/client/client_func_test.go
@@ -26,6 +26,7 @@ package client
 import (
 	"fmt"
 	"io/ioutil"
+	"net/http"
 	"sync"
 	"testing"
 	"time"
@@ -627,4 +628,20 @@ func TestSnapClient(t *testing.T) {
 		So(err, ShouldNotBeNil)
 		So(c, ShouldBeNil)
 	})
+
+	go http.ListenAndServe("127.0.0.1:65000", timeoutHandler{})
+	c, err = New("http://127.0.0.1:65000", "", true, Timeout(time.Second))
+	Convey("Client should timeout", t, func() {
+		So(err, ShouldBeNil)
+		r := c.GetTasks()
+		So(r.Err, ShouldNotBeNil)
+	})
+}
+
+type timeoutHandler struct{}
+
+//ServeHTTP implements http.Handler interface
+func (th timeoutHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
+	time.Sleep(3 * time.Second)
+	w.Write([]byte("Hello!"))
 }