From 164fd0ad3f26e25cc401957dcb2b76b76dc5e9f8 Mon Sep 17 00:00:00 2001
From: Jacob McCann <jacob.mccann2@target.com>
Date: Mon, 8 Jan 2018 12:47:08 -0600
Subject: [PATCH] Update keepalive config to use DurationFlag

---
 cmd/drone-agent/agent.go | 25 ++++---------------------
 cmd/drone-agent/main.go  |  4 ++--
 2 files changed, 6 insertions(+), 23 deletions(-)

diff --git a/cmd/drone-agent/agent.go b/cmd/drone-agent/agent.go
index fdb017aa5..c1f121912 100644
--- a/cmd/drone-agent/agent.go
+++ b/cmd/drone-agent/agent.go
@@ -69,26 +69,6 @@ func loop(c *cli.Context) error {
 
 	// grpc.Dial(target, ))
 
-	agentKeepalive := keepalive.ClientParameters{}
-	if c.String("keepalive-time") != "" {
-		d, err := time.ParseDuration(c.String("keepalive-time"))
-
-		if err != nil {
-			return err
-		}
-
-		agentKeepalive.Time = d
-	}
-	if c.String("keepalive-timeout") != "" {
-		d, err := time.ParseDuration(c.String("keepalive-timeout"))
-
-		if err != nil {
-			return err
-		}
-
-		agentKeepalive.Timeout = d
-	}
-
 	conn, err := grpc.Dial(
 		c.String("server"),
 		grpc.WithInsecure(),
@@ -96,7 +76,10 @@ func loop(c *cli.Context) error {
 			username: c.String("username"),
 			password: c.String("password"),
 		}),
-		grpc.WithKeepaliveParams(agentKeepalive),
+		grpc.WithKeepaliveParams(keepalive.ClientParameters{
+			Time: c.Duration("keepalive-time"),
+			Timeout: c.Duration("keepalive-timeout"),
+		}),
 	)
 
 	if err != nil {
diff --git a/cmd/drone-agent/main.go b/cmd/drone-agent/main.go
index 0ecd2919b..33a93d6a2 100644
--- a/cmd/drone-agent/main.go
+++ b/cmd/drone-agent/main.go
@@ -80,12 +80,12 @@ func main() {
 			Name:   "healthcheck",
 			Usage:  "enables the healthcheck endpoint",
 		},
-		cli.StringFlag{
+		cli.DurationFlag{
 			EnvVar: "DRONE_KEEPALIVE_TIME",
 			Name:   "keepalive-time",
 			Usage:  "after a duration of this time if the agent doesn't see any activity it pings the server to see if the transport is still alive",
 		},
-		cli.StringFlag{
+		cli.DurationFlag{
 			EnvVar: "DRONE_KEEPALIVE_TIMEOUT",
 			Name:   "keepalive-timeout",
 			Usage:  "after having pinged for keepalive check, the client waits for a duration of Timeout and if no activity is seen even after that the connection is closed.",