cu checkout: optional branch name flag (#126)

Signed-off-by: Andrea Luzzardi <[email protected]>

Author: aluzzardi

cmd/cu/checkout.go

@@ -22,7 +22,12 @@ var checkoutCmd = &cobra.Command{
 			return err
 		}
 
-		branch, err := repo.Checkout(ctx, envID)
+		branchName, err := app.Flags().GetString("branch")
+		if err != nil {
+			return err
+		}
+
+		branch, err := repo.Checkout(ctx, envID, branchName)
 		if err != nil {
 			return err
 		}
@@ -33,5 +38,6 @@ var checkoutCmd = &cobra.Command{
 }
 
 func init() {
+	checkoutCmd.Flags().StringP("branch", "b", "", "Local branch name to use")
 	rootCmd.AddCommand(checkoutCmd)
 }

environment/integration/repository_test.go

@@ -17,13 +17,13 @@ func TestRepositoryCreate(t *testing.T) {
 	WithRepository(t, "repository-create", SetupEmptyRepo, func(t *testing.T, repo *repository.Repository, user *UserActions) {
 		// Create an environment
 		env := user.CreateEnvironment("Test Create", "Testing repository create")
-		
+
 		// Verify environment was created properly
 		assert.NotNil(t, env)
 		assert.NotEmpty(t, env.ID)
 		assert.Equal(t, "Test Create", env.State.Title)
 		assert.NotEmpty(t, env.Worktree)
-		
+
 		// Verify worktree was created
 		_, err := os.Stat(env.Worktree)
 		assert.NoError(t, err)
@@ -35,17 +35,17 @@ func TestRepositoryGet(t *testing.T) {
 	t.Parallel()
 	WithRepository(t, "repository-get", SetupEmptyRepo, func(t *testing.T, repo *repository.Repository, user *UserActions) {
 		ctx := context.Background()
-		
+
 		// Create an environment
 		env := user.CreateEnvironment("Test Get", "Testing repository get")
-		
+
 		// Get the environment using repository directly
 		retrieved, err := repo.Get(ctx, user.dag, env.ID)
 		require.NoError(t, err)
 		assert.NotNil(t, retrieved)
 		assert.Equal(t, env.ID, retrieved.ID)
 		assert.Equal(t, env.State.Title, retrieved.State.Title)
-		
+
 		// Test getting non-existent environment
 		_, err = repo.Get(ctx, user.dag, "non-existent-env")
 		assert.Error(t, err)
@@ -57,16 +57,16 @@ func TestRepositoryList(t *testing.T) {
 	t.Parallel()
 	WithRepository(t, "repository-list", SetupEmptyRepo, func(t *testing.T, repo *repository.Repository, user *UserActions) {
 		ctx := context.Background()
-		
+
 		// Create two environments
 		env1 := user.CreateEnvironment("Environment 1", "First test environment")
 		env2 := user.CreateEnvironment("Environment 2", "Second test environment")
-		
+
 		// List should return at least 2
 		envs, err := repo.List(ctx)
 		require.NoError(t, err)
 		assert.GreaterOrEqual(t, len(envs), 2)
-		
+
 		// Verify the environments are in the list
 		var foundIDs []string
 		for _, e := range envs {
@@ -82,20 +82,20 @@ func TestRepositoryDelete(t *testing.T) {
 	t.Parallel()
 	WithRepository(t, "repository-delete", SetupEmptyRepo, func(t *testing.T, repo *repository.Repository, user *UserActions) {
 		ctx := context.Background()
-		
+
 		// Create an environment
 		env := user.CreateEnvironment("Test Delete", "Testing repository delete")
 		worktreePath := env.Worktree
 		envID := env.ID
-		
+
 		// Delete it
 		err := repo.Delete(ctx, envID)
 		require.NoError(t, err)
-		
+
 		// Verify it's gone
 		_, err = repo.Get(ctx, user.dag, envID)
 		assert.Error(t, err)
-		
+
 		// Verify worktree is deleted
 		_, err = os.Stat(worktreePath)
 		assert.True(t, os.IsNotExist(err))
@@ -107,22 +107,22 @@ func TestRepositoryCheckout(t *testing.T) {
 	t.Parallel()
 	WithRepository(t, "repository-checkout", SetupEmptyRepo, func(t *testing.T, repo *repository.Repository, user *UserActions) {
 		ctx := context.Background()
-		
+
 		// Create an environment and add content
 		env := user.CreateEnvironment("Test Checkout", "Testing repository checkout")
 		user.FileWrite(env.ID, "test.txt", "test content", "Add test file")
-		
+
 		// Checkout the environment branch in the source repo
-		branch, err := repo.Checkout(ctx, env.ID)
+		branch, err := repo.Checkout(ctx, env.ID, "")
 		require.NoError(t, err)
 		assert.NotEmpty(t, branch)
-		
+
 		// Verify we're on the correct branch
 		currentBranch, err := repository.RunGitCommand(ctx, repo.SourcePath(), "branch", "--show-current")
 		require.NoError(t, err)
 		// Branch name could be either env.ID or cu-env.ID depending on the logic
 		actualBranch := strings.TrimSpace(currentBranch)
-		assert.True(t, actualBranch == env.ID || actualBranch == "cu-"+env.ID, 
+		assert.True(t, actualBranch == env.ID || actualBranch == "cu-"+env.ID,
 			"Expected branch to be %s or cu-%s, got %s", env.ID, env.ID, actualBranch)
 	})
-}
+}

repository/repository.go

@@ -304,12 +304,14 @@ func (r *Repository) Delete(ctx context.Context, id string) error {
 
 // Checkout changes the user's current branch to that of the identified environment.
 // It attempts to get the most recent commit from the environment without discarding any user changes.
-func (r *Repository) Checkout(ctx context.Context, id string) (string, error) {
+func (r *Repository) Checkout(ctx context.Context, id, branch string) (string, error) {
 	if err := r.exists(ctx, id); err != nil {
 		return "", err
 	}
 
-	branch := "cu-" + id
+	if branch == "" {
+		branch = "cu-" + id
+	}
 
 	// set up remote tracking branch if it's not already there
 	_, err := RunGitCommand(ctx, r.userRepoPath, "show-ref", "--verify", "--quiet", fmt.Sprintf("refs/heads/%s", branch))
@@ -321,7 +323,7 @@ func (r *Repository) Checkout(ctx context.Context, id string) (string, error) {
 		}
 	}
 
-	_, err = RunGitCommand(ctx, r.userRepoPath, "checkout", id)
+	_, err = RunGitCommand(ctx, r.userRepoPath, "checkout", branch)
 	if err != nil {
 		return "", err
 	}