From b40bfc56ce95bfc5ab529a4703f4e90c67c62992 Mon Sep 17 00:00:00 2001
From: Josh Komoroske <josh@stackrox.com>
Date: Mon, 5 Mar 2018 21:46:02 -0800
Subject: [PATCH] Connection timeout tests

---
 plugin_test.go | 52 ++++++++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 52 insertions(+)

diff --git a/plugin_test.go b/plugin_test.go
index 563c06b..26160bb 100644
--- a/plugin_test.go
+++ b/plugin_test.go
@@ -8,6 +8,8 @@ import (
 	"github.com/appleboy/easyssh-proxy"
 	"github.com/stretchr/testify/assert"
 	"strings"
+
+	"time"
 )
 
 func TestMissingHostOrUser(t *testing.T) {
@@ -457,6 +459,56 @@ func TestEnvOutput(t *testing.T) {
 	assert.Equal(t, unindent(expected), unindent(buffer.String()))
 }
 
+// TestConnectTimeoutRetry tests that when a network error occurs, the connect
+// is retried until it either succeeds or the configured timeout is hit.
+func TestConnectTimeoutRetry(t *testing.T) {
+	start := time.Now()
+
+	plugin := Plugin{
+		Config: Config{
+			Host:         []string{"localhost"},
+			UserName:     "drone-scp",
+			Port:         2200,
+			KeyPath:      "./tests/.ssh/id_rsa",
+			Script:       []string{"exit"},
+			RetryTimeout: 15 * time.Second,
+			Sync:         true,
+		},
+	}
+
+	err := plugin.Exec()
+	assert.NotNil(t, err)
+
+	end := time.Now()
+	assert.WithinDuration(t, start.Local().Add(15*time.Second), end, 2*time.Second)
+
+}
+
+// TestConnectTimeoutRetry tests that when a non-network error occurs, the connect
+// is not retried and instead returns the error immediately.
+func TestConnectTimeoutImmediate(t *testing.T) {
+	start := time.Now()
+
+	plugin := Plugin{
+		Config: Config{
+			Host:         []string{"localhost"},
+			UserName:     "drone-scp",
+			Port:         22,
+			Key:          "",
+			Script:       []string{"exit"},
+			RetryTimeout: 60 * time.Second,
+			Sync:         true,
+		},
+	}
+
+	err := plugin.Exec()
+	assert.NotNil(t, err)
+
+	end := time.Now()
+	assert.WithinDuration(t, start.Local().Add(time.Second), end, 2*time.Second)
+
+}
+
 func unindent(text string) string {
 	return strings.TrimSpace(strings.Replace(text, "\t", "", -1))
 }