feat: add Restore function (#2028)

Updates #1682

Author: JoeWang1127

internal/gitrepo/gitrepo.go

@@ -20,6 +20,7 @@ import (
 	"fmt"
 	"log/slog"
 	"os"
+	"os/exec"
 	"strings"
 	"time"
 
@@ -44,6 +45,7 @@ type Repository interface {
 	GetCommitsForPathsSinceCommit(paths []string, sinceCommit string) ([]*Commit, error)
 	CreateBranchAndCheckout(name string) error
 	Push(branchName string) error
+	Restore(paths []string) error
 }
 
 // LocalRepository represents a git repository.
@@ -449,3 +451,19 @@ func (r *LocalRepository) Push(branchName string) error {
 	slog.Info("Successfully pushed branch to remote 'origin", "branch", branchName)
 	return nil
 }
+
+// Restore restores changes in the working tree, leaving staged area untouched.
+// Note that untracked files, if any, are not touched.
+//
+// Wrap git operations in exec, because [git.Worktree.Restore] does not support
+// this operation.
+func (r *LocalRepository) Restore(paths []string) error {
+	args := []string{"restore"}
+	args = append(args, paths...)
+	slog.Info("Restoring uncommitted changes", "paths", strings.Join(paths, ","))
+	cmd := exec.Command("git", args...)
+	cmd.Stderr = os.Stderr
+	cmd.Stdout = os.Stdout
+	cmd.Dir = r.Dir
+	return cmd.Run()
+}

internal/gitrepo/gitrepo_test.go

@@ -15,6 +15,7 @@
 package gitrepo
 
 import (
+	"fmt"
 	"os"
 	"path/filepath"
 	"strings"
@@ -1117,6 +1118,80 @@ func TestCreateBranchAndCheckout(t *testing.T) {
 	}
 }
 
+func TestRestore(t *testing.T) {
+	for _, test := range []struct {
+		name          string
+		paths         []string
+		wantErr       bool
+		wantErrPhrase string
+	}{
+		{
+			name: "restore files in paths",
+			paths: []string{
+				"first/path",
+				"second/path",
+			},
+		},
+	} {
+		t.Run(test.name, func(t *testing.T) {
+			repo, dir := initTestRepo(t)
+			localRepo := &LocalRepository{
+				Dir:  dir,
+				repo: repo,
+			}
+			// Create files in test.paths and commit the change.
+			for _, path := range test.paths {
+				file := filepath.Join(path, "example.txt")
+				createAndCommit(t, repo, file, []byte("old content"), fmt.Sprintf("commit path, %s", path))
+			}
+
+			// Change file contents.
+			for _, path := range test.paths {
+				file := filepath.Join(dir, path, "example.txt")
+				if err := os.WriteFile(file, []byte("new content"), 0755); err != nil {
+					t.Fatal(err)
+				}
+				// Create untracked files.
+				untrackedFile := filepath.Join(dir, path, "untracked.txt")
+				if err := os.WriteFile(untrackedFile, []byte("new content"), 0755); err != nil {
+					t.Fatal(err)
+				}
+			}
+
+			err := localRepo.Restore(test.paths)
+			if test.wantErr {
+				if err == nil {
+					t.Fatalf("%s should return error", test.name)
+				}
+				if !strings.Contains(err.Error(), test.wantErrPhrase) {
+					t.Errorf("Restore() returned error %q, want to contain %q", err.Error(), test.wantErrPhrase)
+				}
+				return
+			}
+
+			if err != nil {
+				t.Fatal(err)
+			}
+
+			for _, path := range test.paths {
+				got, err := os.ReadFile(filepath.Join(dir, path, "example.txt"))
+				if err != nil {
+					t.Fatal(err)
+				}
+				// Verify file contents are restored.
+				if diff := cmp.Diff("old content", string(got)); diff != "" {
+					t.Errorf("Restore() mismatch (-want +got):\n%s", diff)
+				}
+				// Verify the untracked files are untouched.
+				untrackedFile := filepath.Join(dir, path, "untracked.txt")
+				if _, err := os.Stat(untrackedFile); err != nil {
+					t.Errorf("untracked file, %s should not be removed", untrackedFile)
+				}
+			}
+		})
+	}
+}
+
 // initTestRepo creates a new git repository in a temporary directory.
 func initTestRepo(t *testing.T) (*git.Repository, string) {
 	t.Helper()

internal/librarian/command.go

@@ -555,7 +555,7 @@ func findSubDirRelPaths(dir, subDir string) ([]string, error) {
 		return nil, fmt.Errorf("subDir is not nested within the dir: %s, %s", subDir, dir)
 	}
 
-	paths := []string{}
+	var paths []string
 	err = filepath.WalkDir(subDir, func(path string, d fs.DirEntry, err error) error {
 		if err != nil {
 			return err

internal/librarian/librarian_test.go

@@ -45,11 +45,11 @@ func TestRun(t *testing.T) {
 func TestGenerate_DefaultBehavior(t *testing.T) {
 	ctx := context.Background()
 
-	// 1. Setup a mock repository with a state file
+	// 1. Set up a mock repository with a state file
 	repo := newTestGitRepo(t)
 	repoDir := repo.GetDir()
 
-	// Setup a dummy API Source repo to prevent cloning googleapis/googleapis
+	// Set up a dummy API Source repo to prevent cloning googleapis/googleapis
 	apiSourceDir := t.TempDir()
 	runGit(t, apiSourceDir, "init")
 	runGit(t, apiSourceDir, "config", "user.email", "[email protected]")

internal/librarian/mocks_test.go

@@ -317,6 +317,7 @@ type MockRepository struct {
 	ChangedFilesInCommitError              error
 	CreateBranchAndCheckoutError           error
 	PushError                              error
+	RestoreError                           error
 }
 
 func (m *MockRepository) IsClean() (bool, error) {
@@ -424,3 +425,7 @@ func (m *MockRepository) Push(name string) error {
 	}
 	return nil
 }
+
+func (m *MockRepository) Restore(paths []string) error {
+	return m.RestoreError
+}