Cleanup branch

Author: kalverra

tools/flakeguard/cmd/make_pr.go

@@ -10,7 +10,6 @@ import (
 	"time"
 
 	"github.com/go-git/go-git/v5"
-	"github.com/go-git/go-git/v5/config"
 	"github.com/go-git/go-git/v5/plumbing"
 	"github.com/google/go-github/v72/github"
 	"github.com/spf13/cobra"
@@ -88,28 +87,18 @@ func makePR(cmd *cobra.Command, args []string) error {
 	}
 	fmt.Println(" ✅")
 
-	// Create and checkout new branch
-	branchName := fmt.Sprintf("flakeguard-skip-%s", time.Now().Format("20060102150405"))
-	err = targetRepoWorktree.Checkout(&git.CheckoutOptions{
-		Branch: plumbing.NewBranchReferenceName(branchName),
-		Create: true,
-	})
-	if err != nil {
-		return fmt.Errorf("failed to checkout new branch: %w", err)
-	}
-
-	// Push the new branch to GitHub before making any commits
-	fmt.Printf("Pushing new branch '%s' to GitHub, tap your yubikey if it's blinking...", branchName)
-	err = repo.Push(&git.PushOptions{
-		RemoteName: "origin",
-		RefSpecs: []config.RefSpec{
-			config.RefSpec(fmt.Sprintf("refs/heads/%s:refs/heads/%s", branchName, branchName)),
-		},
-	})
-	if err != nil {
-		return fmt.Errorf("failed to push new branch to GitHub: %w", err)
-	}
-	fmt.Println(" ✅")
+	// Cleanup default branch when we're done
+	defer func() {
+		fmt.Printf("Cleaning up branch %s...", defaultBranch)
+		err = targetRepoWorktree.Checkout(&git.CheckoutOptions{
+			Branch: plumbing.NewBranchReferenceName(defaultBranch),
+			Force:  true,
+		})
+		if err != nil {
+			fmt.Printf("Failed to checkout default branch: %v\n", err)
+		}
+		fmt.Println(" ✅")
+	}()
 
 	if len(currentlyFlakyEntries) == 0 {
 		fmt.Println("No flaky tests found!")
@@ -132,6 +121,7 @@ func makePR(cmd *cobra.Command, args []string) error {
 		return fmt.Errorf("failed to modify code to skip tests: %w", err)
 	}
 
+	branchName := fmt.Sprintf("flakeguard-skip-%s", time.Now().Format("20060102150405"))
 	fmt.Print("Committing changes, tap your yubikey if it's blinking...")
 	sha, err := flake_git.MakeSignedCommit(repoPath, fmt.Sprintf("Skips flaky %d tests: %s", len(testsToSkip), strings.Join(jiraTickets, ", ")), branchName, githubToken)
 	if err != nil {

tools/flakeguard/git/git.go

@@ -279,6 +279,37 @@ func MakeSignedCommit(repoPath, commitMessage, branch, githubToken string) (stri
 		return "", err
 	}
 
+	// Create the branch on GitHub using REST API
+	ctx := context.Background()
+	ts := oauth2.StaticTokenSource(&oauth2.Token{AccessToken: githubToken})
+	tc := oauth2.NewClient(ctx, ts)
+	restClient := github.NewClient(tc)
+
+	owner, repoName, _, err := GetOwnerRepoDefaultBranchFromLocalRepo(repoPath)
+	if err != nil {
+		return "", fmt.Errorf("failed to get repo info: %w", err)
+	}
+
+	// Get the reference of the current HEAD (which should be the base branch)
+	headRef, err := repo.Head()
+	if err != nil {
+		return "", fmt.Errorf("failed to get HEAD reference: %w", err)
+	}
+	baseSHA := headRef.Hash().String()
+
+	// Create new branch reference on GitHub
+	newRef := &github.Reference{
+		Ref: github.Ptr("refs/heads/" + branch),
+		Object: &github.GitObject{
+			SHA: github.Ptr(baseSHA),
+		},
+	}
+
+	_, _, err = restClient.Git.CreateRef(ctx, owner, repoName, newRef)
+	if err != nil {
+		return "", fmt.Errorf("failed to create branch on GitHub: %w", err)
+	}
+
 	// Inspired by https://github.com/planetscale/ghcommit/tree/main
 
 	// process added / modified files:
@@ -358,17 +389,8 @@ func MakeSignedCommit(repoPath, commitMessage, branch, githubToken string) (stri
 		body = splitMsg[1]
 	}
 
-	owner, repoName, _, err := GetOwnerRepoDefaultBranchFromLocalRepo(repoPath)
-	if err != nil {
-		return "", err
-	}
-
-	// Get HEAD reference to get the current commit hash
-	headRef, err := repo.Head()
-	if err != nil {
-		return "", fmt.Errorf("failed to get HEAD reference: %w", err)
-	}
-	expectedHeadOid := headRef.Hash().String()
+	// Use the baseSHA from the branch creation as the expected head OID
+	expectedHeadOid := baseSHA
 	// create the $input struct for the graphQL createCommitOnBranch mutation request:
 	input := githubv4.CreateCommitOnBranchInput{
 		Branch: githubv4.CommittableBranch{
@@ -387,7 +409,7 @@ func MakeSignedCommit(repoPath, commitMessage, branch, githubToken string) (stri
 	}
 
 	if err := graphqlClient.Mutate(context.Background(), &m, input, nil); err != nil {
-		return "", err
+		return "", fmt.Errorf("failed to create commit: %w", err)
 	}
 
 	return m.CreateCommitOnBranch.Commit.OID, nil