diff --git a/src/codegen/git/repo_operator/remote_repo_operator.py b/src/codegen/git/repo_operator/remote_repo_operator.py index fe7579e35..7f17f66b1 100644 --- a/src/codegen/git/repo_operator/remote_repo_operator.py +++ b/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""" diff --git a/src/codegen/git/repo_operator/repo_operator.py b/src/codegen/git/repo_operator/repo_operator.py index 0d064ec3e..8353403a6 100644 --- a/src/codegen/git/repo_operator/repo_operator.py +++ b/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 diff --git a/tests/integration/codegen/git/repo_operator/test_remote_repo_operator.py b/tests/integration/codegen/git/repo_operator/test_remote_repo_operator.py index 99604e0d4..6d4363111 100644 --- a/tests/integration/codegen/git/repo_operator/test_remote_repo_operator.py +++ b/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"