GitRepository: Implement set_upstream_branch and get_upstream_branch methods

set_upstream_branch sets upstream branch for the local branch using
git branch --set-upstream
get_upstream_branch returns info about upstream branches

Signed-off-by: Ed Bartosh <eduard.bartosh@intel.com>

2 files changed, 75 insertions, 0 deletions

diff --git a/gbp/git/repository.py b/gbp/git/repository.py
index f3cde12d..3f6f5cb7 100644
--- a/gbp/git/repository.py
+++ b/gbp/git/repository.py
@@ -451,6 +451,44 @@ class GitRepository(object):
                 return True
         return False
 
+    def set_upstream_branch(self, local_branch, upstream):
+        """
+        Set upstream branches for local branch
+
+        @param local_branch: name of the local branch
+        @type local_branch: C{str}
+        @param upstream: remote/branch, for example origin/master
+        @type upstream: C{str}
+        """
+
+        # check if both branches exist
+        for branch, remote in [(local_branch, False), (upstream, True)]:
+            if not self.has_branch(branch, remote=remote):
+                raise GitRepositoryError("Branch %s doesn't exist!" % branch)
+
+            self._git_getoutput('branch', ["--set-upstream", local_branch, upstream])
+
+
+    def get_upstream_branch(self, local_branch):
+        """
+        Get upstream branch for the local branch
+
+        @param local_branch: name fo the local branch
+        @type local_branch: C{str}
+        @return: upstream (remote/branch) or  '' if no upstream found
+        @rtype: C{str}
+
+        """
+        args = GitArgs('--format=%(upstream:short)')
+        if self.has_branch(local_branch, remote=False):
+            args.add('refs/heads/%s' % local_branch)
+        else:
+            raise GitRepositoryError("Branch %s doesn't exist!" % local_branch)
+
+        out = self._git_getoutput('for-each-ref', args.args)[0]
+
+        return out[0].strip()
+
 #{ Tags
 
     def create_tag(self, name, msg=None, commit=None, sign=False, keyid=None):
diff --git a/tests/test_GitRepository.py b/tests/test_GitRepository.py
index 10b448bd..aa20e755 100644
--- a/tests/test_GitRepository.py
+++ b/tests/test_GitRepository.py
@@ -177,6 +177,43 @@ def test_rename_branch():
     """
 
 
+def test_set_upstream_branch():
+    """
+    Set upstream branch master -> origin/master
+
+    >>> import os, shutil
+    >>> import gbp.git
+    >>> repo = gbp.git.GitRepository(repo_dir)
+    >>> os.makedirs(os.path.join(repo.git_dir, 'refs/remotes/origin'))
+    >>> shutil.copy(os.path.join(repo.git_dir, 'refs/heads/master'), \
+                    os.path.join(repo.git_dir, 'refs/remotes/origin/'))
+    >>> repo.set_upstream_branch('master', 'origin/master')
+    >>> repo.get_upstream_branch('master')
+    'origin/master'
+    >>> repo.set_upstream_branch('bla', 'origin/master')
+    Traceback (most recent call last):
+    GitRepositoryError: Branch bla doesn't exist!
+    >>> repo.set_upstream_branch('foo', 'origin/bla')
+    Traceback (most recent call last):
+    GitRepositoryError: Branch origin/bla doesn't exist!
+
+    """
+
+def test_get_upstream_branch():
+    """
+    Get info about upstream branches set in test_set_upstream_branch
+
+    >>> import gbp.git
+    >>> repo = gbp.git.GitRepository(repo_dir)
+    >>> repo.get_upstream_branch('master')
+    'origin/master'
+    >>> repo.get_upstream_branch('foo')
+    ''
+    >>> repo.get_upstream_branch('bla')
+    Traceback (most recent call last):
+    GitRepositoryError: Branch bla doesn't exist!
+    """
+
 def test_tag():
     """
     Create a tag named I{tag} and check it's existance