From 9f887f8b9b06f66579e2989ee01c5b3d68e8a5b9 Mon Sep 17 00:00:00 2001
From: Luke Massa <lukefrederickmassa@gmail.com>
Date: Mon, 12 Aug 2024 22:20:25 -0400
Subject: [PATCH] chore: Make e2e tests more vcs agnostic (#4836)

---
 e2e/e2e.go    | 20 +++++---------------
 e2e/github.go | 16 +++++++++++++++-
 e2e/vcs.go    |  2 ++
 3 files changed, 22 insertions(+), 16 deletions(-)

diff --git a/e2e/e2e.go b/e2e/e2e.go
index 0c2247018..079c329df 100644
--- a/e2e/e2e.go
+++ b/e2e/e2e.go
@@ -128,7 +128,7 @@ func (t *E2ETester) Start(ctx context.Context) (*E2EResult, error) {
 	// waiting for atlantis run and finish
 	maxLoops := 20
 	i := 0
-	for ; i < maxLoops && checkStatus(state); i++ {
+	for ; i < maxLoops && t.vcsClient.IsAtlantisInProgress(state); i++ {
 		time.Sleep(2 * time.Second)
 		state, _ = t.vcsClient.GetAtlantisStatus(ctx, branchName)
 		if state == "" {
@@ -144,22 +144,13 @@ func (t *E2ETester) Start(ctx context.Context) (*E2EResult, error) {
 	log.Printf("atlantis run finished with status %q", state)
 	e2eResult.testResult = state
 	// check if atlantis run was a success
-	if state != "success" {
+	if !t.vcsClient.DidAtlantisSucceed(state) {
 		return e2eResult, fmt.Errorf("atlantis run project type %q failed with %q status", t.projectType.Name, state)
 	}
 
 	return e2eResult, nil
 }
 
-func checkStatus(state string) bool {
-	for _, s := range []string{"success", "error", "failure"} {
-		if state == s {
-			return false
-		}
-	}
-	return true
-}
-
 func cleanUp(ctx context.Context, t *E2ETester, pullRequestNumber int, branchName string) error {
 	// clean up
 	err := t.vcsClient.ClosePullRequest(ctx, pullRequestNumber)
@@ -168,12 +159,11 @@ func cleanUp(ctx context.Context, t *E2ETester, pullRequestNumber int, branchNam
 	}
 	log.Printf("closed pull request %d", pullRequestNumber)
 
-	deleteBranchName := fmt.Sprintf("%s/%s", "heads", branchName)
-	err = t.vcsClient.DeleteBranch(ctx, deleteBranchName)
+	err = t.vcsClient.DeleteBranch(ctx, branchName)
 	if err != nil {
-		return fmt.Errorf("error while deleting branch %s: %v", deleteBranchName, err)
+		return fmt.Errorf("error while deleting branch %s: %v", branchName, err)
 	}
-	log.Printf("deleted branch %s", deleteBranchName)
+	log.Printf("deleted branch %s", branchName)
 
 	return nil
 }
diff --git a/e2e/github.go b/e2e/github.go
index 40b64f7ef..06cc2ebbc 100644
--- a/e2e/github.go
+++ b/e2e/github.go
@@ -153,9 +153,23 @@ func (g GithubClient) ClosePullRequest(ctx context.Context, pullRequestNumber in
 }
 func (g GithubClient) DeleteBranch(ctx context.Context, branchName string) error {
 
-	_, err := g.client.Git.DeleteRef(ctx, g.ownerName, g.repoName, branchName)
+	deleteBranchName := fmt.Sprintf("%s/%s", "heads", branchName)
+	_, err := g.client.Git.DeleteRef(ctx, g.ownerName, g.repoName, deleteBranchName)
 	if err != nil {
 		return fmt.Errorf("error while deleting branch %s: %v", branchName, err)
 	}
 	return nil
 }
+
+func (g GithubClient) IsAtlantisInProgress(state string) bool {
+	for _, s := range []string{"success", "error", "failure"} {
+		if state == s {
+			return false
+		}
+	}
+	return true
+}
+
+func (g GithubClient) DidAtlantisSucceed(state string) bool {
+	return state == "success"
+}
diff --git a/e2e/vcs.go b/e2e/vcs.go
index 79f3deb41..2648a913f 100644
--- a/e2e/vcs.go
+++ b/e2e/vcs.go
@@ -23,4 +23,6 @@ type VCSClient interface {
 	GetAtlantisStatus(ctx context.Context, branchName string) (string, error)
 	ClosePullRequest(ctx context.Context, pullRequestNumber int) error
 	DeleteBranch(ctx context.Context, branchName string) error
+	IsAtlantisInProgress(state string) bool
+	DidAtlantisSucceed(state string) bool
 }