Add tests for worktree cache update and remove unused initial_directory tracking

- Add test_traverse_updates_worktree_cache_on_checkout to verify cache updates when checking out in same worktree
- Improve test_traverse_cd_from_linked_to_main_worktree to better cover cache update scenarios
- Remove initial_directory tracking and os.chdir() back to initial directory in --return-to=here mode
  (as cwd changes don't propagate to calling shell anyway)

Note: Coverage shows 63->69 branch not taken, which is expected - we always find and delete
the old branch entry when updating the cache after checkout

Author: PawelLipski

git_machete/client/traverse.py

@@ -143,7 +143,6 @@ def traverse(
 
         # Store the main worktree path and initial directory for later restoration
         main_worktree_path = self._git.get_main_worktree_path()
-        initial_directory = os.getcwd()
         initial_branch = nearest_remaining_branch = self._git.get_current_branch()
 
         # Fetch worktrees once at the start to avoid repeated git worktree list calls
@@ -474,10 +473,9 @@ def traverse(
                 break
 
         if opt_return_to == TraverseReturnTo.HERE:
-            # Return to initial branch and initial directory
+            # Return to initial branch
+            # No point switching back to initial directory as cwd won't propagate back to the calling shell anyway
             self._switch_to_branch_worktree(initial_branch, main_worktree_path=main_worktree_path)
-            if os.getcwd() != initial_directory and os.path.exists(initial_directory):
-                os.chdir(initial_directory)
         elif opt_return_to == TraverseReturnTo.NEAREST_REMAINING:
             self._switch_to_branch_worktree(nearest_remaining_branch, main_worktree_path=main_worktree_path)
             # For NEAREST_REMAINING, we stay in the worktree where the branch is checked out

tests/test_traverse.py

@@ -2018,7 +2018,7 @@ def test_traverse_cd_from_linked_to_main_worktree(self) -> None:
         check_out("root")
         new_branch("branch-2")
         commit()
-        push()
+        # Don't push branch-2 so it will be pushed during traverse
 
         body = """
         root
@@ -2035,17 +2035,62 @@ def test_traverse_cd_from_linked_to_main_worktree(self) -> None:
         os.chdir(branch_1_worktree)
 
         # Now run traverse --start-from=first-root
-        # This should checkout root, which is NOT in any worktree
-        # We're currently in branch-1 linked worktree
-        # So it should cd to main worktree (lines 88-89)
-        output = launch_command("traverse", "-y", "--return-to=here", "--start-from=first-root")
+        # This will:
+        # 1. Checkout root (not in any worktree) in main worktree - triggers cache update (lines 63-66)
+        # 2. Checkout branch-1 (already in linked worktree where we are)
+        # 3. Checkout branch-2 (was in main worktree, now root is there) - triggers cache update again
+        # The cache updates should cover lines 63-66 where we delete the old branch entry
+        output = launch_command("traverse", "-y", "--start-from=first-root")
 
-        # Verify lines 88-89 were executed
+        # Verify lines 88-89 were executed (cd to main worktree for root)
         assert "Changing directory to main worktree at" in output, \
             f"Expected 'Changing directory to main worktree at' in output, but got:\n{output}"
 
-        # Verify we returned to branch-1 worktree (line 475)
-        current_dir = os.path.realpath(os.getcwd())
-        branch_1_worktree_real = os.path.realpath(branch_1_worktree)
-        assert current_dir == branch_1_worktree_real, \
-            f"Expected to be in {branch_1_worktree_real}, but in {current_dir}"
+        # Verify branch-2 was pushed (confirms traverse completed successfully)
+        assert "Pushing untracked branch branch-2" in output
+
+    def test_traverse_updates_worktree_cache_on_checkout(self) -> None:
+        """Test that the worktree cache is properly updated when checking out in the same worktree."""
+        if get_git_version() < (2, 5):
+            # git worktree command was introduced in git 2.5
+            return
+
+        (local_path, _) = create_repo_with_remote()
+        new_branch("root")
+        commit()
+        push()
+        new_branch("branch-1")
+        commit()
+        push()
+        new_branch("branch-2")
+        commit()
+        # Don't push branch-2 so it will be pushed during traverse
+
+        body: str = \
+            """
+            root
+              branch-1
+              branch-2
+            """
+        rewrite_branch_layout_file(body)
+
+        # Setup: Main worktree on root, create a linked worktree for branch-1
+        check_out("root")
+        add_worktree("branch-1")
+
+        # Run traverse from root in main worktree
+        # This will:
+        # 1. Initial cache: {root: main_worktree_path, branch-1: linked_worktree_path}
+        # 2. Visit root (already checked out in main worktree) - no operation
+        # 3. Visit branch-1 (in linked worktree) - cd to linked worktree, no checkout
+        # 4. Visit branch-2 (not checked out anywhere) - cd to main worktree, checkout branch-2
+        #    - When checking out branch-2, _update_worktrees_cache_after_checkout is called
+        #    - The cache has {root: main_worktree_path, branch-1: linked_worktree_path}
+        #    - current_worktree_path = main_worktree_path
+        #    - Loop finds root with main_worktree_path, deletes it (covers line 64-65)
+        #    - Adds branch-2: main_worktree_path
+        output = launch_command("traverse", "-y")
+
+        # Verify branch-2 was checked out and pushed
+        assert "Checking out branch-2" in output
+        assert "Pushing untracked branch branch-2" in output