From b786b2e1f578028ee36edce895edde264d5b41c0 Mon Sep 17 00:00:00 2001
From: Amir Zarrinkafsh <nightah@me.com>
Date: Sat, 28 Nov 2020 11:06:42 +1100
Subject: [PATCH] [MISC] Refactor webdriver port initialization (#1491)

This change aims to factorize code introduced in #1467 for webdriver port customisation within the suites.

Co-authored-by: James Elliott <james-d-elliott@users.noreply.github.com>
---
 internal/suites/suite_network_acl_test.go | 18 ++----------------
 internal/suites/utils.go                  | 17 +++++++++++++++++
 internal/suites/webdriver.go              | 10 +---------
 3 files changed, 20 insertions(+), 25 deletions(-)

diff --git a/internal/suites/suite_network_acl_test.go b/internal/suites/suite_network_acl_test.go
index 3222681c..1dc1dc13 100644
--- a/internal/suites/suite_network_acl_test.go
+++ b/internal/suites/suite_network_acl_test.go
@@ -3,8 +3,6 @@ package suites
 import (
 	"context"
 	"fmt"
-	"os"
-	"strconv"
 	"testing"
 	"time"
 
@@ -41,13 +39,7 @@ func (s *NetworkACLSuite) TestShouldAccessSecretUpon1FA() {
 	ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
 	defer cancel()
 
-	driverPort := os.Getenv("CHROMEDRIVER_PORT")
-	if driverPort == "" {
-		driverPort = defaultChromeDriverPort
-	}
-
-	p, _ := strconv.Atoi(driverPort)
-	wds, err := StartWebDriverWithProxy("http://proxy-client1.example.com:3128", p)
+	wds, err := StartWebDriverWithProxy("http://proxy-client1.example.com:3128", GetWebDriverPort())
 	s.Require().NoError(err)
 
 	defer wds.Stop() //nolint:errcheck // TODO: Legacy code, consider refactoring time permitting.
@@ -66,13 +58,7 @@ func (s *NetworkACLSuite) TestShouldAccessSecretUpon0FA() {
 	ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
 	defer cancel()
 
-	driverPort := os.Getenv("CHROMEDRIVER_PORT")
-	if driverPort == "" {
-		driverPort = defaultChromeDriverPort
-	}
-
-	p, _ := strconv.Atoi(driverPort)
-	wds, err := StartWebDriverWithProxy("http://proxy-client2.example.com:3128", p)
+	wds, err := StartWebDriverWithProxy("http://proxy-client2.example.com:3128", GetWebDriverPort())
 	s.Require().NoError(err)
 
 	defer wds.Stop() //nolint:errcheck // TODO: Legacy code, consider refactoring time permitting.
diff --git a/internal/suites/utils.go b/internal/suites/utils.go
index 951ab296..770030c7 100644
--- a/internal/suites/utils.go
+++ b/internal/suites/utils.go
@@ -1,5 +1,10 @@
 package suites
 
+import (
+	"os"
+	"strconv"
+)
+
 // GetLoginBaseURL returns the URL of the login portal and the path prefix if specified.
 func GetLoginBaseURL() string {
 	if PathPrefix != "" {
@@ -8,3 +13,15 @@ func GetLoginBaseURL() string {
 
 	return LoginBaseURL
 }
+
+// GetWebDriverPort returns the port to initialize the webdriver with.
+func GetWebDriverPort() int {
+	driverPort := os.Getenv("CHROMEDRIVER_PORT")
+	if driverPort == "" {
+		driverPort = defaultChromeDriverPort
+	}
+
+	p, _ := strconv.Atoi(driverPort)
+
+	return p
+}
diff --git a/internal/suites/webdriver.go b/internal/suites/webdriver.go
index 25fc8b02..01dda26f 100644
--- a/internal/suites/webdriver.go
+++ b/internal/suites/webdriver.go
@@ -7,7 +7,6 @@ import (
 	"fmt"
 	"io/ioutil"
 	"os"
-	"strconv"
 	"strings"
 	"testing"
 	"time"
@@ -73,14 +72,7 @@ func StartWebDriverWithProxy(proxy string, port int) (*WebDriverSession, error)
 
 // StartWebDriver create a selenium session.
 func StartWebDriver() (*WebDriverSession, error) {
-	driverPort := os.Getenv("CHROMEDRIVER_PORT")
-	if driverPort == "" {
-		driverPort = defaultChromeDriverPort
-	}
-
-	p, _ := strconv.Atoi(driverPort)
-
-	return StartWebDriverWithProxy("", p)
+	return StartWebDriverWithProxy("", GetWebDriverPort())
 }
 
 // Stop stop the selenium session.