From e672f4cccf25914e0430adefe5af7d331226ae0f Mon Sep 17 00:00:00 2001
From: Neel Dalsania <neel.dalsania@astronomer.io>
Date: Wed, 5 Jul 2023 21:23:20 +0530
Subject: [PATCH] Add timeout on github client (#1290)

* add timeout on github client

* implement linter recommendation
---
 cmd/root.go             |  6 ++++--
 version/version_test.go | 36 ++++++++++++++++++++++++++++++++++++
 2 files changed, 40 insertions(+), 2 deletions(-)
 create mode 100644 version/version_test.go

diff --git a/cmd/root.go b/cmd/root.go
index 11bc2bbc0..778487083 100644
--- a/cmd/root.go
+++ b/cmd/root.go
@@ -2,7 +2,9 @@ package cmd
 
 import (
 	"fmt"
+	"net/http"
 	"os"
+	"time"
 
 	"github.com/astronomer/astro-cli/cmd/registry"
 
@@ -70,8 +72,8 @@ Welcome to the Astro CLI, the modern command line interface for data orchestrati
 		PersistentPreRunE: func(cmd *cobra.Command, args []string) error {
 			// Check for latest version
 			if config.CFG.UpgradeMessage.GetBool() {
-				// create github client
-				githubClient := github.NewClient(nil)
+				// create github client with 3 second timeout, setting an aggressive timeout since its not mandatory to get a response in each command execution
+				githubClient := github.NewClient(&http.Client{Timeout: 3 * time.Second})
 				// compare current version to latest
 				err = version.CompareVersions(githubClient, "astronomer", "astro-cli")
 				if err != nil {
diff --git a/version/version_test.go b/version/version_test.go
new file mode 100644
index 000000000..caef40d0a
--- /dev/null
+++ b/version/version_test.go
@@ -0,0 +1,36 @@
+package version
+
+import (
+	"fmt"
+	"net/http"
+	"net/http/httptest"
+	"net/url"
+	"testing"
+	"time"
+
+	"github.com/google/go-github/v48/github"
+	"github.com/stretchr/testify/assert"
+)
+
+func TestGithubAPITimeout(t *testing.T) {
+	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		time.Sleep(5 * time.Second) // sleeping and doing nothing
+	}))
+	defer ts.Close()
+	githubURL, err := url.Parse(fmt.Sprintf("%s/", ts.URL))
+	assert.NoError(t, err)
+
+	githubClient := github.NewClient(&http.Client{Timeout: 1 * time.Second}) // client side timeout should be less than server side sleep defined above
+	githubClient.BaseURL = githubURL
+
+	start := time.Now()
+	release, err := getLatestRelease(githubClient, "test", "test")
+	elapsed := time.Since(start)
+	// assert time to get a response from the function is only slightly greater than client timeout
+	assert.GreaterOrEqual(t, elapsed, 1*time.Second)
+	assert.Less(t, elapsed, 2*time.Second)
+	// assert error returned is related to client timeout
+	assert.Nil(t, release)
+	assert.Error(t, err)
+	assert.Contains(t, err.Error(), "context deadline exceeded")
+}