Update source acceptance test with auto

3 files changed, 40 insertions, 6 deletions

diff --git a/test/acceptance/helpers_test.go b/test/acceptance/helpers_test.go
index 3506e236..d228f787 100644
--- a/test/acceptance/helpers_test.go
+++ b/test/acceptance/helpers_test.go
@@ -223,7 +223,7 @@ func RunPagesProcessWithOutput(t *testing.T, pagesBinary string, listeners []Lis
 
 func RunPagesProcessWithStubGitLabServer(t *testing.T, wait bool, pagesBinary string, listeners []ListenSpec, promPort string, envs []string, extraArgs ...string) (teardown func()) {
 	var apiCalled bool
-	source := NewGitlabDomainsSourceStub(t, &apiCalled)
+	source := NewGitlabDomainsSourceStub(t, &apiCalled, 0)
 
 	gitLabAPISecretKey := CreateGitLabAPISecretKeyFixtureFile(t)
 	pagesArgs := append([]string{"-gitlab-server", source.URL, "-api-secret-key", gitLabAPISecretKey, "-domain-config-source", "gitlab"}, extraArgs...)
@@ -536,10 +536,16 @@ func waitForRoundtrips(t *testing.T, listeners []ListenSpec, timeout time.Durati
 	require.Equal(t, len(listeners), nListening, "all listeners must be accepting TCP connections")
 }
 
-func NewGitlabDomainsSourceStub(t *testing.T, apiCalled *bool) *httptest.Server {
+func NewGitlabDomainsSourceStub(t *testing.T, apiCalled *bool, readyCount int) *httptest.Server {
 	*apiCalled = false
+	currentStatusCount := 0
+
 	mux := http.NewServeMux()
 	mux.HandleFunc("/api/v4/internal/pages/status", func(w http.ResponseWriter, r *http.Request) {
+		if currentStatusCount < readyCount {
+			w.WriteHeader(http.StatusBadGateway)
+		}
+
 		w.WriteHeader(http.StatusNoContent)
 	})
 
diff --git a/test/acceptance/serving_test.go b/test/acceptance/serving_test.go
index 4ccdd8f4..da2843a0 100644
--- a/test/acceptance/serving_test.go
+++ b/test/acceptance/serving_test.go
@@ -379,6 +379,7 @@ func TestDomainsSource(t *testing.T) {
 		configSource string
 		domain       string
 		urlSuffix    string
+		readyCount   int
 	}
 	type want struct {
 		statusCode int
@@ -450,13 +451,40 @@ func TestDomainsSource(t *testing.T) {
 				apiCalled:  false,
 			},
 		},
-		// TODO: modify mock so we can test domain-config-source=auto when API/disk is not ready https://gitlab.com/gitlab-org/gitlab/-/issues/218358
+		{
+			name: "auto_source_gitlab_is_not_ready",
+			args: args{
+				configSource: "auto",
+				domain:       "test.domain.com",
+				urlSuffix:    "/",
+				readyCount:   100, // big number to ensure the API is in bad state for a while
+			},
+			want: want{
+				statusCode: http.StatusOK,
+				content:    "main-dir\n",
+				apiCalled:  false,
+			},
+		},
+		{
+			name: "auto_source_gitlab_is_ready",
+			args: args{
+				configSource: "auto",
+				domain:       "new-source-test.gitlab.io",
+				urlSuffix:    "/my/pages/project/",
+				readyCount:   0,
+			},
+			want: want{
+				statusCode: http.StatusOK,
+				content:    "New Pages GitLab Source TEST OK\n",
+				apiCalled:  true,
+			},
+		},
 	}
 
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
 			var apiCalled bool
-			source := NewGitlabDomainsSourceStub(t, &apiCalled)
+			source := NewGitlabDomainsSourceStub(t, &apiCalled, tt.args.readyCount)
 			defer source.Close()
 
 			gitLabAPISecretKey := CreateGitLabAPISecretKeyFixtureFile(t)
diff --git a/test/acceptance/zip_test.go b/test/acceptance/zip_test.go
index 3bfa21bf..5d3037c8 100644
--- a/test/acceptance/zip_test.go
+++ b/test/acceptance/zip_test.go
@@ -14,7 +14,7 @@ func TestZipServing(t *testing.T) {
 	skipUnlessEnabled(t)
 
 	var apiCalled bool
-	source := NewGitlabDomainsSourceStub(t, &apiCalled)
+	source := NewGitlabDomainsSourceStub(t, &apiCalled, 0)
 	defer source.Close()
 
 	gitLabAPISecretKey := CreateGitLabAPISecretKeyFixtureFile(t)
@@ -108,7 +108,7 @@ func TestZipServingConfigShortTimeout(t *testing.T) {
 	skipUnlessEnabled(t)
 
 	var apiCalled bool
-	source := NewGitlabDomainsSourceStub(t, &apiCalled)
+	source := NewGitlabDomainsSourceStub(t, &apiCalled, 0)
 	defer source.Close()
 
 	gitLabAPISecretKey := CreateGitLabAPISecretKeyFixtureFile(t)