fix CG-9440 clean repo - clears from the default branch (#398)

Author: christinewangcw

src/codegen/git/repo_operator/remote_repo_operator.py

@@ -79,6 +79,22 @@ def codeowners_parser(self) -> CodeOwnersParser | None:
     # SET UP
     ####################################################################################################################
 
+    @override
+    def clean_repo(self) -> None:
+        """Cleans the repo by:
+        1. Discards any changes (tracked/untracked)
+        2. Checks out the default branch (+ makes sure it's up to date with the remote)
+        3. Deletes all branches except the default branch
+        4. Deletes all remotes except origin
+
+        Used in SetupOption.PULL_OR_CLONE to allow people to re-use existing repos and start from a clean state.
+        """
+        logger.info(f"Cleaning repo at {self.repo_path} ...")
+        self.discard_changes()
+        self.checkout_branch(self.default_branch, remote=True)
+        self.clean_branches()
+        self.clean_remotes()
+
     @override
     def pull_repo(self) -> None:
         """Pull the latest commit down to an existing local repo"""

src/codegen/git/repo_operator/repo_operator.py

@@ -159,15 +159,11 @@ def repo_exists(self) -> bool:
     def clean_repo(self) -> None:
         """Cleans the repo by:
         1. Discards any changes (tracked/untracked)
-        2. Checks out the default branch (+ makes sure it's up to date with the remote)
-        3. Deletes all branches except the default branch
-        4. Deletes all remotes except origin
-
-        Used in SetupOption.PULL_OR_CLONE to allow people to re-use existing repos and start from a clean state.
+        2. Deletes all branches except the checked out branch
+        3. Deletes all remotes except origin
         """
         logger.info(f"Cleaning repo at {self.repo_path} ...")
         self.discard_changes()
-        self.checkout_branch(self.default_branch)  # TODO(CG-9440): add back remote=True
         self.clean_branches()
         self.clean_remotes()
 
@@ -277,20 +273,6 @@ def is_branch_checked_out(self, branch_name: str) -> bool:
             return False
         return self.git_cli.active_branch.name == branch_name
 
-    def delete_local_branch(self, branch_name: str) -> None:
-        if branch_name not in self.git_cli.branches:
-            logger.info(f"Branch {branch_name} does not exist locally. Skipping delete_local_branch.")
-            return
-        if branch_name is self.default_branch:
-            msg = "Deleting the default branch is not implemented yet."
-            raise NotImplementedError(msg)
-
-        if self.is_branch_checked_out(branch_name):
-            self.checkout_branch(self.default_branch)
-
-        logger.info(f"Deleting local branch: {branch_name} ...")
-        self.git_cli.delete_head(branch_name, force=True)  # force deletes even if the branch has unmerged changes
-
     def checkout_branch(self, branch_name: str | None, *, remote: bool = False, remote_name: str = "origin", create_if_missing: bool = True) -> CheckoutResult:
         """Attempts to check out the branch in the following order:
         - Check out the local branch by name

tests/integration/codegen/git/repo_operator/test_remote_repo_operator.py

@@ -12,7 +12,12 @@
 
 @pytest.fixture
 def op(repo_config, request, tmpdir):
-    op = RemoteRepoOperator(repo_config, shallow=request.param, base_dir=tmpdir, bot_commit=False)
+    op = RemoteRepoOperator(
+        repo_config,
+        shallow=request.param if hasattr(request, "param") else True,
+        base_dir=tmpdir,
+        bot_commit=False,
+    )
     yield op
 
 
@@ -76,3 +81,24 @@ def test_checkout_branch_remote_already_checked_out_resets_branch(mock_git_clien
     assert res == CheckoutResult.SUCCESS
     assert len(op.git_cli.heads) == 1
     assert op.head_commit.hexsha == original_commit_head.hexsha
+
+
+def test_clean_repo(op: RemoteRepoOperator):
+    num_branches = len(op.git_cli.branches)
+    op.checkout_branch(branch_name="test_branch", create_if_missing=True)
+    with open(f"{op.repo_path}/test.txt", "w") as f:
+        f.write("test")
+    op.git_cli.git.add(A=True)
+    op.git_cli.create_remote(name="other-remote", url=op.clone_url)
+
+    assert op.git_cli.active_branch.name == "test_branch"
+    assert len(op.git_cli.branches) == num_branches + 1
+    assert len(op.git_cli.remotes) == 2
+    assert op.git_cli.is_dirty()
+
+    op.clean_repo()
+    assert not op.git_cli.is_dirty()  # discards changes
+    assert len(op.git_cli.branches) == 1  # deletes only the checked out branch
+    assert op.git_cli.active_branch.name == op.default_branch
+    assert len(op.git_cli.remotes) == 1  # deletes all remotes except origin
+    assert op.git_cli.remotes[0].name == "origin"