fix: cookie test (#9608)

influxdata · Aug 17, 2021 · 02ccbec · 02ccbec
1 parent 41c384a
commit 02ccbec
Show file tree

Hide file tree

Showing 3 changed files with 23 additions and 19 deletions.
diff --git a/plugins/common/cookie/cookie.go b/plugins/common/cookie/cookie.go
@@ -9,6 +9,7 @@ import (
 	"strings"
 	"time"
 
+	clockutil "github.com/benbjohnson/clock"
 	"github.com/influxdata/telegraf"
 	"github.com/influxdata/telegraf/config"
 )
@@ -27,7 +28,7 @@ type CookieAuthConfig struct {
 	client *http.Client
 }
 
-func (c *CookieAuthConfig) Start(client *http.Client, log telegraf.Logger) (err error) {
+func (c *CookieAuthConfig) Start(client *http.Client, log telegraf.Logger, clock clockutil.Clock) (err error) {
 	c.client = client
 
 	if c.Method == "" {
@@ -45,7 +46,7 @@ func (c *CookieAuthConfig) Start(client *http.Client, log telegraf.Logger) (err
 
 	// continual auth renewal if set
 	if c.Renewal > 0 {
-		ticker := time.NewTicker(time.Duration(c.Renewal))
+		ticker := clock.Ticker(time.Duration(c.Renewal))
 		go func() {
 			for range ticker.C {
 				if err := c.auth(); err != nil && log != nil {

diff --git a/plugins/common/cookie/cookie_test.go b/plugins/common/cookie/cookie_test.go
@@ -9,6 +9,7 @@ import (
 	"testing"
 	"time"
 
+	clockutil "github.com/benbjohnson/clock"
 	"github.com/google/go-cmp/cmp"
 	"github.com/influxdata/telegraf/config"
 	"github.com/influxdata/telegraf/plugins/common/cookie"
@@ -121,20 +122,19 @@ func TestAuthConfig_Start(t *testing.T) {
 		fields  fields
 		args    args
 		wantErr error
-		assert  func(t *testing.T, c *cookie.CookieAuthConfig, srv fakeServer)
+		assert  func(t *testing.T, c *cookie.CookieAuthConfig, srv fakeServer, mock *clockutil.Mock)
 	}{
 		{
 			name: "zero renewal does not renew",
 			args: args{
 				renewal:  0,
 				endpoint: authEndpointNoCreds,
 			},
-			assert: func(t *testing.T, c *cookie.CookieAuthConfig, srv fakeServer) {
+			assert: func(t *testing.T, c *cookie.CookieAuthConfig, srv fakeServer, mock *clockutil.Mock) {
 				// should have Cookie Authed once
 				srv.checkAuthCount(t, 1)
 				srv.checkResp(t, http.StatusOK)
-				time.Sleep(renewalCheck)
-				// should have never Cookie Authed again
+				mock.Add(renewalCheck)
 				srv.checkAuthCount(t, 1)
 				srv.checkResp(t, http.StatusOK)
 			},
@@ -145,13 +145,13 @@ func TestAuthConfig_Start(t *testing.T) {
 				renewal:  renewal,
 				endpoint: authEndpointNoCreds,
 			},
-			assert: func(t *testing.T, c *cookie.CookieAuthConfig, srv fakeServer) {
+			assert: func(t *testing.T, c *cookie.CookieAuthConfig, srv fakeServer, mock *clockutil.Mock) {
 				// should have Cookie Authed once
 				srv.checkAuthCount(t, 1)
 				// default method set
 				require.Equal(t, http.MethodPost, c.Method)
 				srv.checkResp(t, http.StatusOK)
-				time.Sleep(renewalCheck)
+				mock.Add(renewalCheck)
 				// should have Cookie Authed at least twice more
 				srv.checkAuthCount(t, 3)
 				srv.checkResp(t, http.StatusOK)
@@ -168,11 +168,11 @@ func TestAuthConfig_Start(t *testing.T) {
 				renewal:  renewal,
 				endpoint: authEndpointWithBasicAuth,
 			},
-			assert: func(t *testing.T, c *cookie.CookieAuthConfig, srv fakeServer) {
+			assert: func(t *testing.T, c *cookie.CookieAuthConfig, srv fakeServer, mock *clockutil.Mock) {
 				// should have Cookie Authed once
 				srv.checkAuthCount(t, 1)
 				srv.checkResp(t, http.StatusOK)
-				time.Sleep(renewalCheck)
+				mock.Add(renewalCheck)
 				// should have Cookie Authed at least twice more
 				srv.checkAuthCount(t, 3)
 				srv.checkResp(t, http.StatusOK)
@@ -190,11 +190,11 @@ func TestAuthConfig_Start(t *testing.T) {
 				endpoint: authEndpointWithBasicAuth,
 			},
 			wantErr: fmt.Errorf("cookie auth renewal received status code: 401 (Unauthorized)"),
-			assert: func(t *testing.T, c *cookie.CookieAuthConfig, srv fakeServer) {
+			assert: func(t *testing.T, c *cookie.CookieAuthConfig, srv fakeServer, mock *clockutil.Mock) {
 				// should have never Cookie Authed
 				srv.checkAuthCount(t, 0)
 				srv.checkResp(t, http.StatusForbidden)
-				time.Sleep(renewalCheck)
+				mock.Add(renewalCheck)
 				// should have still never Cookie Authed
 				srv.checkAuthCount(t, 0)
 				srv.checkResp(t, http.StatusForbidden)
@@ -210,11 +210,11 @@ func TestAuthConfig_Start(t *testing.T) {
 				renewal:  renewal,
 				endpoint: authEndpointWithBody,
 			},
-			assert: func(t *testing.T, c *cookie.CookieAuthConfig, srv fakeServer) {
+			assert: func(t *testing.T, c *cookie.CookieAuthConfig, srv fakeServer, mock *clockutil.Mock) {
 				// should have Cookie Authed once
 				srv.checkAuthCount(t, 1)
 				srv.checkResp(t, http.StatusOK)
-				time.Sleep(renewalCheck)
+				mock.Add(renewalCheck)
 				// should have Cookie Authed at least twice more
 				srv.checkAuthCount(t, 3)
 				srv.checkResp(t, http.StatusOK)
@@ -231,11 +231,11 @@ func TestAuthConfig_Start(t *testing.T) {
 				endpoint: authEndpointWithBody,
 			},
 			wantErr: fmt.Errorf("cookie auth renewal received status code: 401 (Unauthorized)"),
-			assert: func(t *testing.T, c *cookie.CookieAuthConfig, srv fakeServer) {
+			assert: func(t *testing.T, c *cookie.CookieAuthConfig, srv fakeServer, mock *clockutil.Mock) {
 				// should have never Cookie Authed
 				srv.checkAuthCount(t, 0)
 				srv.checkResp(t, http.StatusForbidden)
-				time.Sleep(renewalCheck)
+				mock.Add(renewalCheck)
 				// should have still never Cookie Authed
 				srv.checkAuthCount(t, 0)
 				srv.checkResp(t, http.StatusForbidden)
@@ -255,15 +255,17 @@ func TestAuthConfig_Start(t *testing.T) {
 				Renewal:  config.Duration(tt.args.renewal),
 			}
 
-			if err := c.Start(srv.Client(), testutil.Logger{Name: "cookie_auth"}); tt.wantErr != nil {
+			mock := clockutil.NewMock()
+			if err := c.Start(srv.Client(), testutil.Logger{Name: "cookie_auth"}, mock); tt.wantErr != nil {
 				require.EqualError(t, err, tt.wantErr.Error())
 			} else {
 				require.NoError(t, err)
 			}
 
 			if tt.assert != nil {
-				tt.assert(t, c, srv)
+				tt.assert(t, c, srv, mock)
 			}
+			srv.Close()
 		})
 	}
 }
diff --git a/plugins/common/http/config.go b/plugins/common/http/config.go
@@ -5,6 +5,7 @@ import (
 	"net/http"
 	"time"
 
+	"github.com/benbjohnson/clock"
 	"github.com/influxdata/telegraf"
 	"github.com/influxdata/telegraf/config"
 	"github.com/influxdata/telegraf/plugins/common/cookie"
@@ -54,7 +55,7 @@ func (h *HTTPClientConfig) CreateClient(ctx context.Context, log telegraf.Logger
 	client = h.OAuth2Config.CreateOauth2Client(ctx, client)
 
 	if h.CookieAuthConfig.URL != "" {
-		if err := h.CookieAuthConfig.Start(client, log); err != nil {
+		if err := h.CookieAuthConfig.Start(client, log, clock.New()); err != nil {
 			return nil, err
 		}
 	}