From 8b8b3dbbdbcff00a96353466d89b80d05127a00e Mon Sep 17 00:00:00 2001
From: Sam Coe <samcoe@users.noreply.github.com>
Date: Mon, 9 May 2022 10:53:36 +0200
Subject: [PATCH] Add default host as known host when environment auth token
 exists

---
 internal/config/config.go      |  3 ++
 internal/config/config_test.go | 59 ++++++++++++++++++++++++++++++++++
 internal/set/string_set.go     |  1 +
 3 files changed, 63 insertions(+)

diff --git a/internal/config/config.go b/internal/config/config.go
index 42ac321..b54d4d1 100644
--- a/internal/config/config.go
+++ b/internal/config/config.go
@@ -72,6 +72,9 @@ func (c config) Hosts() []string {
 	if host := os.Getenv(ghHost); host != "" {
 		hosts.Add(host)
 	}
+	if token, _ := c.AuthToken(defaultHost); token != "" {
+		hosts.Add(defaultHost)
+	}
 	entries := c.hosts.keys()
 	hosts.AddValues(entries)
 	return hosts.ToSlice()
diff --git a/internal/config/config_test.go b/internal/config/config_test.go
index ab21677..c2477cc 100644
--- a/internal/config/config_test.go
+++ b/internal/config/config_test.go
@@ -398,6 +398,65 @@ func TestConfigHost(t *testing.T) {
 	}
 }
 
+func TestConfigHosts(t *testing.T) {
+	tests := []struct {
+		name      string
+		cfg       Config
+		ghHost    string
+		ghToken   string
+		wantHosts []string
+	}{
+		{
+			name:      "no known hosts",
+			cfg:       testLoadedNoHostConfig(),
+			wantHosts: []string{},
+		},
+		{
+			name:      "includes GH_HOST",
+			cfg:       testLoadedNoHostConfig(),
+			ghHost:    "test.com",
+			wantHosts: []string{"test.com"},
+		},
+		{
+			name:      "includes authenticated hosts",
+			cfg:       testLoadedConfig(),
+			wantHosts: []string{"github.com", "enterprise.com"},
+		},
+		{
+			name:      "includes default host if environment auth token",
+			cfg:       testLoadedNoHostConfig(),
+			ghToken:   "TOKEN",
+			wantHosts: []string{"github.com"},
+		},
+		{
+			name:      "deduplicates hosts",
+			cfg:       testLoadedConfig(),
+			ghHost:    "test.com",
+			ghToken:   "TOKEN",
+			wantHosts: []string{"test.com", "github.com", "enterprise.com"},
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			if tt.ghHost != "" {
+				k := "GH_HOST"
+				old := os.Getenv(k)
+				os.Setenv(k, tt.ghHost)
+				defer os.Setenv(k, old)
+			}
+			if tt.ghToken != "" {
+				k := "GH_TOKEN"
+				old := os.Getenv(k)
+				os.Setenv(k, tt.ghToken)
+				defer os.Setenv(k, old)
+			}
+			hosts := tt.cfg.Hosts()
+			assert.Equal(t, tt.wantHosts, hosts)
+		})
+	}
+}
+
 func TestConfigAuthToken(t *testing.T) {
 	orig_GITHUB_TOKEN := os.Getenv("GITHUB_TOKEN")
 	orig_GITHUB_ENTERPRISE_TOKEN := os.Getenv("GITHUB_ENTERPRISE_TOKEN")
diff --git a/internal/set/string_set.go b/internal/set/string_set.go
index 2fc2c59..8be4492 100644
--- a/internal/set/string_set.go
+++ b/internal/set/string_set.go
@@ -10,6 +10,7 @@ type stringSet struct {
 func NewStringSet() *stringSet {
 	s := &stringSet{}
 	s.m = make(map[string]struct{})
+	s.v = []string{}
 	return s
 }