path: root/linaropy/git/worktree.py

import unittest
import logging
import os

import uuid

from ..proj import Proj

from ..cd import cd
from gitrepo import GitRepo

from clone import Clone

from sh import git, ErrorReturnCode, rm

import shutil


class Worktree(GitRepo):

    @classmethod
    def create(cls, proj, repo, path, localBranch, trackedBranch=None):
        """ Factory function for creating worktrees in new directories.

        Parameters
        ----------
        proj : Proj
            Temporary project directory
        repo : Clone
            Repo that we're creating a worktree for.
        path
            Path to the proposed worktree (must not exist).
        localBranch
            The name of the local branch that the created worktree should be on.
        trackedBranch
            Branch to base the local branch on (only valid when the local branch
            is a new branch).
        """
        if not isinstance(proj, Proj):
            raise TypeError('Unsupported input proj type')

        # Technically, we could support Worktree here as well, but that's
        # unlikely to be very useful in practice so there's no use complicating
        # our testing and our interfaces with this (repo needs to have a
        # clonedir method)
        if not isinstance(repo, Clone):
            raise TypeError('Unsupported input repo type')

        path = str(path)

        # An empty path would result in an empty variable expansion
        # and a malformed git worktree add expression.
        if path == "" or path == "None":
            raise TypeError(
                'You must specify a worktree directory when creating a Worktree')

        if not os.path.isabs(path):
            raise TypeError('Worktree path must be absolute')

        if os.path.exists(path):
            raise EnvironmentError('Worktree path %s already exists' % path)

        if localBranch is None:
            raise TypeError(
                'You must specify a local branch when creating a Worktree')

        localBranch = str(localBranch)

        if not GitRepo.is_valid_branch_name(localBranch):
            raise TypeError('Invalid local branch %s' % localBranch)

        if trackedBranch is None:
            trackedBranch = "master"
        else:
            if repo.branch_exists(localBranch):
                raise TypeError(
                    'Local branch %s already exists; it is invalid to provide a tracked branch' %
                    localBranch)

            trackedBranch = str(trackedBranch)

            if not GitRepo.is_valid_branch_name(trackedBranch):
                raise TypeError(
                    'Invalid tracked branch %s' % trackedBranch)

        if not repo.branch_exists(trackedBranch):
            raise EnvironmentError(
                'Tracked branch %s does not exist in repo %s' %
                (trackedBranch, repo.clonedir()))

        try:
            # worktree add needs to be called inside the repo directory
            with cd(repo.clonedir()):
                logging.info(
                    "Worktree(): calling git worktree add in %s with branch %s tracking %s " %
                    (path, localBranch, trackedBranch))
                if repo.branch_exists(localBranch):
                    git("worktree", "add", path, localBranch)
                else:
                    git("worktree", "add", "-b", localBranch, path,
                        trackedBranch)
        except ErrorReturnCode as exc:
            raise EnvironmentError("Unable to create a git worktree")

        return cls(proj, path)

    def __init__(self, proj, path):
        """
        Create a worktree object representing the worktree at the given path.
        The worktree must already exist.

        Parameters
        ----------
        proj : Proj
            Temporary project directory.
        path
            Path to the worktree.
        """
        super(Worktree, self).__init__(proj)

        if path is None:
            raise TypeError(
                'Must specify worktree path when creating a worktree')

        if not os.path.isdir(path):
            raise EnvironmentError('%s does not name a directory' % path)

        if not os.path.isfile(os.path.join(path, ".git")):
            # A worktree always contains a .git file in its root (as opposed to
            # a .git directory, as contained by the main working tree obtained
            # with git init or git clone).
            # Note that we are VERY specific about where we want that .git file
            # to live - using git rev-parse could mislead us into thinking that
            # any subdirectory within a worktree is itself a worktree, which
            # makes it impossible to detect complex hierarchies of worktrees.
            raise EnvironmentError('%s is not a worktree' % path)
        self.repodir = path

    def get_original_repo(self):
        # The git common dir should point to the .git directory of the repo that
        # the worktree was created from. We can then strip the .git to get the
        # path to the original repo.
        with cd(self.repodir):
            commonDir = str(git("rev-parse", "--git-common-dir"))[:-1]
            repo, _ = os.path.split(commonDir)
            return repo

    def clean(self, deleteBranch, forceBranchDelete=False):
        """ Clean up the current worktree.

        Delete its directory and run prune on the original repo. If deleteBranch
        is true, also delete the branch that is checked out in the worktree, but
        only if it is merged. To delete an unmerged branch, set forceBranchDelete
        to true.

        It is an error to invoke this method on a worktree whose directory has
        already been deleted (either by another call to clean or through any
        other means).
        """
        if forceBranchDelete and not deleteBranch:
            raise TypeError(
                "Can't force branch deletion if deleteBranch is false")

        if not os.path.isdir(self.repodir):
            raise EnvironmentError('Worktree directory not found: %s' %
                                   self.repodir)

        try:
            branch = self.getbranch()
        except Exception as exc:
            raise EnvironmentError('Failed to get current branch for %s' %
                                   self.repodir)

        originalRepo = self.get_original_repo()
        try:
            shutil.rmtree(self.repodir)
            logging.info("Worktree clean: removed worktree directory %s" %
                         self.repodir)
        except Exception as exc:
            raise EnvironmentError('Failed to remove worktree directory: %s' %
                                   str(exc))

        with cd(originalRepo):
            try:
                git("worktree", "prune")
                logging.info("Worktree clean: pruned repo %s" % originalRepo)
            except ErrorReturnCode:
                raise EnvironmentError(
                    'Worktree directory was removed, but git prune failed')

            if deleteBranch:
                try:
                    self.delete_branch(branch, forceBranchDelete)
                    logging.info("Worktree clean: deleted branch %s" % branch)
                except EnvironmentError as exc:
                    raise EnvironmentError(
                        'Worktree directory %s was removed, but branch deletion failed: %s' %
                        (self.repodir, str(exc)))


class TestWorktree(unittest.TestCase):
    testdirprefix = "WorktreeUT"

    # TODO: these are duplicated in the GitRepo tests - reuse them
    def __create_dummy_commit(self):
        filename = "file" + str(uuid.uuid4())
        open(filename, "a").close()
        git("add", filename)
        git("commit", "-m", "Branches without commits confuse git")

    def __get_current_branch(self):
        branch = str(git("rev-parse", "--abbrev-ref", "HEAD"))
        # git rev-parse returns a trailing newline that we must get rid of
        return branch[:-1]

    def setUp(self):
        self.proj = Proj(prefix=TestWorktree.testdirprefix)

        repoPath = str(os.path.join(self.proj.projdir, "repo"))
        os.makedirs(repoPath)

        with cd(repoPath):
            git("init")
            self.__create_dummy_commit()

        self.repo = Clone(self.proj, repoPath)

    def tearDown(self):
        # We clean up the entire proj directory between tests in order to ensure
        # that no state survives between tests. The proj directory contains not
        # only the clone that we use for each test (where we don't want leftover
        # branches) but also various worktree directories (which may have
        # different names between tests, e.g. worktreedir, worktree1,
        # unimportantworktreedir etc). We could rely on each test to clean up
        # after itself, but it's safer to just nuke everything out of
        # existence.
        self.proj.cleanup()

    def test_worktree(self):
        worktreePath = str(os.path.join(self.proj.projdir, "worktreedir"))
        worktreeBranch = "worktreebranch"

        self.worktree = Worktree.create(self.proj, self.repo, worktreePath,
                                        worktreeBranch, "master")

        self.assertTrue(os.path.isdir(worktreePath),
                        "Failed to create worktree directory")

        with cd(worktreePath):
            self.assertEqual(self.__get_current_branch(),
                             worktreeBranch,
                             "Worktree is on the wrong branch")

    def test_worktree_track_existing(self):
        worktreePath = str(os.path.join(self.proj.projdir, "worktreedir"))
        worktreeBranch = "worktreebranch"
        parentBranch = "parentbranch"

        with cd(self.repo.clonedir()):
            git("checkout", "-b", parentBranch)
            self.__create_dummy_commit()

        self.worktree = Worktree.create(self.proj, self.repo, worktreePath,
                                        worktreeBranch, parentBranch)

        with cd(worktreePath):
            self.assertEqual(git("rev-parse", worktreeBranch),
                             git("rev-parse", parentBranch),
                             "Worktree branch is not based on the parent branch")

    def test_worktree_track_new(self):
        worktreePath = str(os.path.join(self.proj.projdir, "worktreedir"))
        worktreeBranch = "worktreebranch"
        parentBranch = "parentbranch"

        with self.assertRaises(EnvironmentError) as context:
            self.worktree = Worktree.create(self.proj, self.repo, worktreePath,
                                            worktreeBranch, parentBranch)

        self.assertEqual(str(context.exception),
                         "Tracked branch %s does not exist in repo %s" %
                         (parentBranch, self.repo.clonedir()))

    def test_worktree_track_default(self):
        worktreePath = str(os.path.join(self.proj.projdir, "worktreedir"))
        worktreeBranch = "worktreebranch"

        self.worktree = Worktree.create(self.proj, self.repo, worktreePath,
                                        worktreeBranch)

        self.assertTrue(os.path.isdir(worktreePath),
                        "Failed to create worktree directory")

        with cd(worktreePath):
            self.assertEqual(self.__get_current_branch(),
                             worktreeBranch,
                             "Worktree is on the wrong branch")
            self.assertEqual(git("rev-parse", worktreeBranch),
                             git("rev-parse", "master"),
                             "Worktree branch is not based on master")

    def test_worktree_track_invalid(self):
        worktreePath = os.path.join(
            self.proj.projdir, "unimportantworktreedir")
        with self.assertRaises(TypeError) as context:
            self.worktree = Worktree.create(self.proj, self.repo, worktreePath,
                                            "worktreebranch",
                                            "invalid branch name")

        self.assertEqual(str(context.exception),
                         "Invalid tracked branch invalid branch name")

    def test_worktree_local_invalid(self):
        worktreePath = os.path.join(
            self.proj.projdir, "unimportantworktreedir")
        with self.assertRaises(TypeError) as context:
            self.worktree = Worktree.create(self.proj, self.repo, worktreePath,
                                            "invalid branch name")

        self.assertEqual(str(context.exception),
                         "Invalid local branch invalid branch name")

        with self.assertRaises(TypeError) as context:
            self.worktree = Worktree.create(self.proj, self.repo, worktreePath,
                                            None)

        self.assertEqual(str(context.exception),
                         "You must specify a local branch when creating a Worktree")

    def test_worktree_local_existing(self):
        worktreePath = str(os.path.join(self.proj.projdir, "worktreedir"))
        worktreeBranch = "worktreebranch"

        with cd(self.repo.clonedir()):
            git("checkout", "-b", worktreeBranch)
            self.__create_dummy_commit()
            git("checkout", "master")

        self.worktree = Worktree.create(self.proj, self.repo, worktreePath,
                                        worktreeBranch)

        with cd(worktreePath):
            self.assertEqual(self.__get_current_branch(),
                             worktreeBranch,
                             "Worktree is on the wrong branch")

    def test_worktree_local_checked_out(self):
        worktreePath = str(os.path.join(self.proj.projdir, "worktreedir"))
        worktreeBranch = "worktreebranch"

        with cd(self.repo.clonedir()):
            git("checkout", "-b", worktreeBranch)
            self.__create_dummy_commit()

        with self.assertRaises(EnvironmentError) as context:
            self.worktree = Worktree.create(self.proj, self.repo, worktreePath,
                                            worktreeBranch)

        self.assertEqual(str(context.exception),
                         "Unable to create a git worktree")

    def test_worktree_track_with_existing_local(self):
        worktreePath = str(os.path.join(self.proj.projdir, "worktreedir"))
        worktreeBranch = "worktreebranch"

        with cd(self.repo.clonedir()):
            git("checkout", "-b", worktreeBranch)
            self.__create_dummy_commit()
            git("checkout", "master")

        with self.assertRaises(TypeError) as context:
            self.worktree = Worktree.create(self.proj, self.repo, worktreePath,
                                            worktreeBranch, "master")

        self.assertEqual(str(context.exception),
                         "Local branch worktreebranch already exists; it is invalid to provide a tracked branch")

    def test_worktree_dir_existing(self):
        worktreePath = str(os.path.join(self.proj.projdir, "exists"))

        os.makedirs(worktreePath)

        with self.assertRaises(EnvironmentError) as context:
            self.worktree = Worktree.create(self.proj, self.repo, worktreePath,
                                            "existingdir")

        self.assertEqual(str(context.exception),
                         "Worktree path %s already exists" %
                         worktreePath)

    def test_worktree_dir_invalid(self):
        with self.assertRaises(TypeError) as context:
            self.worktree = Worktree.create(self.proj, self.repo, None,
                                            "invalidpath")

        self.assertEqual(str(context.exception),
                         "You must specify a worktree directory when creating a Worktree")

        with self.assertRaises(TypeError) as context:
            self.worktree = Worktree.create(self.proj, self.repo, "",
                                            "invalidpath")

        self.assertEqual(str(context.exception),
                         "You must specify a worktree directory when creating a Worktree")

    def test_worktree_dir_absolute(self):
        worktreePath = os.path.join(self.proj.projdir, "worktreedir")
        worktreePath = str(os.path.abspath(worktreePath))

        self.worktree = Worktree.create(self.proj, self.repo, worktreePath,
                                        "worktreebranch")

        self.assertTrue(os.path.isdir(worktreePath),
                        "Failed to create worktree directory")

    def test_worktree_dir_relative(self):
        startPath = os.path.join(self.proj.projdir, "start", "here")
        worktreePath = os.path.join(self.proj.projdir, "worktreedir")

        os.makedirs(startPath)
        with cd(startPath):
            relativePath = os.path.relpath(worktreePath, startPath)

            with self.assertRaises(TypeError) as context:
                self.worktree = Worktree.create(self.proj, self.repo,
                                                relativePath, "worktreebranch")

            self.assertTrue(str(context.exception),
                            "Worktree path must be absolute")

    def test_worktree_dir_missing_hops(self):
        worktreePath = str(os.path.join(self.proj.projdir, "none", "of",
                                        "these", "exist", "yet"))

        self.worktree = Worktree.create(self.proj, self.repo, worktreePath,
                                        "worktreebranch")

        self.assertTrue(os.path.isdir(worktreePath),
                        "Failed to create worktree directory")

    def test_worktree_clone_invalid(self):
        worktreePath = str(os.path.join(self.proj.projdir, "worktreedir"))

        with self.assertRaises(TypeError) as context:
            self.worktree = Worktree.create(self.proj, "not a clone",
                                            worktreePath, "worktreebranch")

        self.assertEqual(str(context.exception), "Unsupported input repo type")

        with self.assertRaises(TypeError) as context:
            self.worktree = Worktree.create(self.proj, None, worktreePath,
                                            "worktreebranch")

        self.assertEqual(str(context.exception), "Unsupported input repo type")

    def test_worktree_proj_invalid(self):
        worktreePath = str(os.path.join(self.proj.projdir, "worktreedir"))

        with self.assertRaises(TypeError) as context:
            self.worktree = Worktree.create("not a proj", self.repo,
                                            worktreePath, "worktreebranch")

        self.assertEqual(str(context.exception), "Unsupported input proj type")

        with self.assertRaises(TypeError) as context:
            self.worktree = Worktree.create(None, self.repo, worktreePath,
                                            "worktreebranch")

        self.assertEqual(str(context.exception), "Unsupported input proj type")

    def test_worktree_multiple_calls(self):
        worktreePath1 = os.path.join(self.proj.projdir, "worktree1")
        branch1 = "branch1"
        track1 = "track1"

        worktreePath2 = os.path.join(self.proj.projdir, "worktree2")
        branch2 = "branch2"
        track2 = "track2"

        with cd(self.repo.clonedir()):
            git("checkout", "-b", track1)
            self.__create_dummy_commit()

            git("checkout", "-b", track2)
            self.__create_dummy_commit()

        self.worktree1 = Worktree.create(self.proj, self.repo, worktreePath1,
                                         branch1, track1)

        self.worktree2 = Worktree.create(self.proj, self.repo, worktreePath2,
                                         branch2, track2)

        self.assertTrue(os.path.isdir(worktreePath1),
                        "Failed to create worktree directory")
        self.assertTrue(os.path.isdir(worktreePath2),
                        "Failed to create worktree directory")

        with cd(worktreePath1):
            self.assertEqual(self.__get_current_branch(),
                             branch1,
                             "Worktree is on the wrong branch")

            self.assertEqual(git("rev-parse", branch1),
                             git("rev-parse", track1),
                             "Worktree branch is not based on the correct branch")

        with cd(worktreePath2):
            self.assertEqual(self.__get_current_branch(),
                             branch2,
                             "Worktree is on the wrong branch")

            self.assertEqual(git("rev-parse", branch2),
                             git("rev-parse", track2),
                             "Worktree branch is not based on the correct branch")

    def test_worktree_no_path(self):
        with self.assertRaises(TypeError) as context:
            self.worktree = Worktree(self.proj, None)

        self.assertEqual(str(context.exception),
                         "Must specify worktree path when creating a worktree")

    def test_worktree_invalid_path(self):
        worktreePath = os.path.join(self.proj.projdir, "does", "not", "exist")

        with self.assertRaises(EnvironmentError) as context:
            self.worktree = Worktree(self.proj, worktreePath)

        self.assertEqual(str(context.exception),
                         "%s does not name a directory" %
                         worktreePath)

        worktreePath = os.path.join(self.proj.projdir, "file")

        open(worktreePath, "a").close()

        with self.assertRaises(EnvironmentError) as context:
            self.worktree = Worktree(self.proj, worktreePath)

        self.assertEqual(str(context.exception),
                         "%s does not name a directory" % worktreePath)

    def test_not_a_worktree(self):
        worktreePath = os.path.join(self.proj.projdir, "worktreedir")

        # Test that an empty directory is not mistaken for a worktree.
        os.makedirs(worktreePath)

        with self.assertRaises(EnvironmentError) as context:
            self.worktree = Worktree(self.proj, worktreePath)

        self.assertEqual(str(context.exception),
                         "%s is not a worktree" % worktreePath)

        # Test that a non-empty directory that doesn't contain anything
        # git-related is not mistaken for a worktree.
        os.makedirs(os.path.join(worktreePath, "notempty"))

        with self.assertRaises(EnvironmentError) as context:
            self.worktree = Worktree(self.proj, worktreePath)

        self.assertEqual(str(context.exception),
                         "%s is not a worktree" % worktreePath)

        # Test that a main git workting tree (obtained with git init or git
        # clone) isn't mistaken for a worktree.
        with self.assertRaises(EnvironmentError) as context:
            self.worktree = Worktree(self.proj, self.repo.clonedir())

        self.assertEqual(str(context.exception),
                         "%s is not a worktree" % self.repo.clonedir())

        # Test that a directory that contains a worktree as a subdirectory is
        # not mistaken for a worktree.
        childWorktreePath = os.path.join(worktreePath, "aworktree")
        with cd(self.repo.clonedir()):
            git("worktree", "add", "-b", "wokrtreebranch1", childWorktreePath)

        with self.assertRaises(EnvironmentError) as context:
            self.worktree = Worktree(self.proj, worktreePath)

        self.assertEqual(str(context.exception),
                         "%s is not a worktree" % worktreePath)

        # Test that a subdirectory of a worktree is not mistaken for a
        # worktree.
        worktreeSubdirPath = os.path.join(childWorktreePath, "notaworktree")
        os.makedirs(worktreeSubdirPath)

        with self.assertRaises(EnvironmentError) as context:
            self.worktree = Worktree(self.proj, worktreeSubdirPath)

        self.assertEqual(str(context.exception),
                         "%s is not a worktree" % worktreeSubdirPath)

    def test_get_original_repo(self):
        worktreePath = os.path.join(self.proj.projdir, "worktreedir")
        worktreeBranch = "worktreebranch"

        with cd(self.repo.clonedir()):
            git("worktree", "add", "-b", worktreeBranch, worktreePath)

        worktree = Worktree(self.proj, worktreePath)
        self.assertEqual(worktree.get_original_repo(), self.repo.clonedir())

    def test_cleanup(self):
        worktreePath = os.path.join(self.proj.projdir, "worktreedir")
        worktreeBranch = "worktreebranch"

        with cd(self.repo.clonedir()):
            git("worktree", "add", "-b", worktreeBranch, worktreePath)
            self.assertTrue(os.path.isdir(worktreePath))

        self.worktree = Worktree(self.proj, worktreePath)
        self.worktree.clean(False)

        self.assertFalse(os.path.isdir(worktreePath))

        with cd(self.repo.clonedir()):
            git("rev-parse", worktreeBranch)

    def test_cleanup_delete_merged_branch(self):
        worktreePath = os.path.join(self.proj.projdir, "worktreedir")
        worktreeBranch = "worktreebranch"

        with cd(self.repo.clonedir()):
            git("worktree", "add", "-b", worktreeBranch, worktreePath)
            self.assertTrue(os.path.isdir(worktreePath))

            with cd(worktreePath):
                self.__create_dummy_commit()

            git("merge", worktreeBranch)

        self.worktree = Worktree(self.proj, worktreePath)
        self.worktree.clean(True)

        self.assertFalse(os.path.isdir(worktreePath))

        with cd(self.repo.clonedir()):
            with self.assertRaises(ErrorReturnCode) as context:
                git("rev-parse", worktreeBranch)

    def test_cleanup_delete_unmerged_branch(self):
        worktreePath = os.path.join(self.proj.projdir, "worktreedir")
        worktreeBranch = "worktreebranch"

        with cd(self.repo.clonedir()):
            git("worktree", "add", "-b", worktreeBranch, worktreePath)
            self.assertTrue(os.path.isdir(worktreePath))

            with cd(worktreePath):
                self.__create_dummy_commit()

        self.worktree = Worktree(self.proj, worktreePath)

        with self.assertRaises(EnvironmentError) as context:
            self.worktree.clean(True)

        self.assertRegexpMatches(str(context.exception),
                                 "branch deletion failed:")
        self.assertRegexpMatches(str(context.exception),
                                 "not fully merged")

        self.assertFalse(os.path.isdir(worktreePath))

        with cd(self.repo.clonedir()):
            git("rev-parse", worktreeBranch)

    def test_cleanup_force_delete_branch(self):
        worktreePath = os.path.join(self.proj.projdir, "worktreedir")
        worktreeBranch = "worktreebranch"

        with cd(self.repo.clonedir()):
            git("worktree", "add", "-b", worktreeBranch, worktreePath)
            self.assertTrue(os.path.isdir(worktreePath))

            with cd(worktreePath):
                self.__create_dummy_commit()

        self.worktree = Worktree(self.proj, worktreePath)
        self.worktree.clean(True, True)

        self.assertFalse(os.path.isdir(worktreePath))

        with cd(self.repo.clonedir()):
            with self.assertRaises(ErrorReturnCode) as context:
                git("rev-parse", worktreeBranch)

    def test_cleanup_invalid_force(self):
        worktreePath = os.path.join(self.proj.projdir, "worktreedir")
        worktreeBranch = "worktreebranch"

        with cd(self.repo.clonedir()):
            git("worktree", "add", "-b", worktreeBranch, worktreePath)
            self.assertTrue(os.path.isdir(worktreePath))

            with cd(worktreePath):
                self.__create_dummy_commit()

        self.worktree = Worktree(self.proj, worktreePath)

        with self.assertRaises(TypeError) as context:
            self.worktree.clean(False, True)

        self.assertEqual(str(context.exception),
                         "Can't force branch deletion if deleteBranch is false")

        self.assertTrue(os.path.isdir(worktreePath))

        with cd(self.repo.clonedir()):
            git("rev-parse", worktreeBranch)

    def test_cleanup_already_clean(self):
        worktreePath = os.path.join(self.proj.projdir, "worktreedir")
        worktreeBranch = "worktreebranch"

        with cd(self.repo.clonedir()):
            git("worktree", "add", "-b", worktreeBranch, worktreePath)
            self.assertTrue(os.path.isdir(worktreePath))

        self.worktree = Worktree(self.proj, worktreePath)

        rm("-rf", worktreePath)

        with self.assertRaises(EnvironmentError) as context:
            self.worktree.clean(False)

        self.assertRegexpMatches(str(context.exception),
                                 'Worktree directory not found: %s' %
                                 worktreePath)

if __name__ == '__main__':
    # logging.basicConfig(level="INFO")
    unittest.main()