diff --git a/moz-phab b/moz-phab
index d5b9316..1e9953a 100755
--- a/moz-phab
+++ b/moz-phab
@@ -569,6 +569,9 @@ class Repository(object):
         May be called multiple times.
         If an exception is raised this is NOT called (to avoid dataloss)."""
 
+    def finalize(self, commits):
+        """Update the history after node changed."""
+
     def set_args(self, args):
         self.args = args
 
@@ -1107,7 +1110,7 @@ class Git(Repository):
 
         # Store current branch (fails if HEAD in detached state)
         try:
-            self.current_head = self._get_current_head()
+            self.branch = self._get_current_head()
         except Exception:
             raise Error(
                 "Git failed to read the branch name.\n"
@@ -1127,7 +1130,35 @@ class Git(Repository):
 
     def cleanup(self):
         self.git(["gc", "--auto", "--quiet"])
-        self.checkout(self.current_head)
+        self.checkout(self.branch)
+
+    def _find_branches_to_rebase(self, commits):
+        """Create a list of branches to rebase."""
+        branches_to_rebase = dict()
+        for commit in commits:
+            if commit["node"] == commit["orig-node"]:
+                continue
+            branches = self.git_out(["branch", "--contains", commit["orig-node"]])
+            for branch in branches:
+                if branch.startswith("* ("):
+                    # Omit `* (detached from {SHA1})`
+                    continue
+
+                branch = branch.lstrip("* ")
+                # Rebase the branch to the last commit from the stack .
+                branches_to_rebase[branch] = [commit["node"], commit["orig-node"]]
+
+        return branches_to_rebase
+
+    def finalize(self, commits):
+        """Rebase all branches based on changed commits from the stack."""
+        branches_to_rebase = self._find_branches_to_rebase(commits)
+
+        for branch, nodes in branches_to_rebase.iteritems():
+            self.checkout(branch)
+            self._rebase(*nodes)
+
+        self.checkout(self.branch)
 
     def refresh_commit_stack(self, commits):
         """Update revset and names of the commits."""
@@ -1299,25 +1330,24 @@ class Git(Repository):
         for log_line in self._get_commits_info(*self.revset):
             if not log_line:
                 continue
+
             parents, tree_hash, node, desc = log_line.split("\n", 3)
             desc = desc.splitlines()
+
+            # Check if the commit is a child of the first one
             if not rev_list:
                 rev_list = self._git_get_children(node)
-
-            # Check if commit has multiple children, if so - raise an Error
-            if (
-                node[:12] != self.revset[1]
-                and len(self._get_direct_children(node, rev_list)) > 1
-            ):
+                first_node = node
+            elif not self._is_child(first_node, node, rev_list):
                 raise Error(
-                    "Multiple children found for commit %s, unable to continue"
-                    % node[:12]
+                    "Commit %s is not a child of %s, unable to continue"
+                    % (node[:12], first_node[:12])
                 )
 
             # Check if commit has multiple parents, if so - raise an Error
             # We may push the merging commit if it's the first one
             parents = parents.split(" ")
-            if node[:12] != self.revset[0] and len(parents) > 1:
+            if node[:12] != first_node[:12] and len(parents) > 1:
                 raise Error(
                     "Multiple parents found for commit %s, unable to continue"
                     % node[:12]
@@ -1327,6 +1357,7 @@ class Git(Repository):
                 {
                     "name": node[:12],
                     "node": node,
+                    "orig-node": node,
                     "title": desc[0],
                     "title-preview": desc[0],
                     "body": "\n".join(desc[1:]).rstrip(),
@@ -1334,19 +1365,10 @@ class Git(Repository):
                     "reviewers": [],
                     "rev-id": None,
                     "parent": parents[0],
-                    "tree_hash": tree_hash,
+                    "tree-hash": tree_hash,
                 }
             )
 
-        # Check if all commits are children of the first one
-        first_node = commits[0]["node"]
-        for commit in commits[1:]:
-            if not self._is_child(first_node, commit["node"], rev_list):
-                raise Error(
-                    "Commit %s is not a child of %s"
-                    % (commit["node"][:12], first_node[:12])
-                )
-
         return commits
 
     def checkout(self, node):
@@ -1406,7 +1428,7 @@ class Git(Repository):
 
         # Create a new commit with the updated body.
         new_parent_sha = self._commit_tree(
-            commit["parent"], commit["tree_hash"], updated_body
+            commit["parent"], commit["tree-hash"], updated_body
         )
 
         # Update commit info
@@ -1424,22 +1446,11 @@ class Git(Repository):
             # Update parent information and create a new commit
             c["parent"] = new_parent_sha
             new_parent_sha = self._commit_tree(
-                new_parent_sha, c["tree_hash"], "%s\n%s" % (c["title"], c["body"])
+                new_parent_sha, c["tree-hash"], "%s\n%s" % (c["title"], c["body"])
             )
             old_parent_sha = c["node"]
             c["node"] = new_parent_sha
 
-        # We've created a new linear history with detached commits.
-        # All branches based on the old SHA1 of the last commit need to be re-parented
-        # to its new SHA1.
-        branches = self.git_out(["branch", "--contains", old_parent_sha])
-        for branch in branches:
-            if branch.startswith("* ("):
-                # Omit `* (detached from {SHA1})`
-                continue
-            self.checkout(branch.lstrip("*").lstrip())
-            self._rebase(new_parent_sha, old_parent_sha)
-
     def rebase_commit(self, source_commit, dest_commit):
         self._rebase(dest_commit["node"], source_commit["node"])
 
@@ -2064,6 +2075,7 @@ def submit(repo, args):
     # final showing.
     repo.cleanup()
     repo.refresh_commit_stack(commits)
+    repo.finalize(commits)
 
     logger.warning("\nCompleted")
     show_commit_stack(repo, commits, show_rev_urls=True)
diff --git a/tests/conftest.py b/tests/conftest.py
new file mode 100644
index 0000000..011b306
--- /dev/null
+++ b/tests/conftest.py
@@ -0,0 +1,32 @@
+import imp
+import mock
+import os
+import pytest
+
+mozphab = imp.load_source(
+    "mozphab", os.path.join(os.path.dirname(__file__), os.path.pardir, "moz-phab")
+)
+
+
+@pytest.fixture()
+@mock.patch("mozphab.Git.git_out")
+@mock.patch("mozphab.Git._get_current_head")
+@mock.patch("mozphab.Config")
+@mock.patch("mozphab.os.path")
+@mock.patch("mozphab.which")
+@mock.patch("mozphab.Repository._phab_url")
+def git(
+    m_repository_phab_url,
+    m_which,
+    m_os_path,
+    m_config,
+    m_git_get_current_head,
+    m_git_git_out,
+):
+    m_os_path.join = os.path.join
+    m_os_path.exists.return_value = True
+    m_which.return_value = True
+    m_os_path.isfile.return_value = False
+    m_git_get_current_head.return_value = "branch"
+    m_git_git_out.return_value = ["user.email=email"]
+    return mozphab.Git("x")
diff --git a/tests/test_git.py b/tests/test_git.py
index f4d0f9c..42cff17 100644
--- a/tests/test_git.py
+++ b/tests/test_git.py
@@ -1,150 +1,159 @@
 import imp
 import mock
 import os
-import unittest
+import pytest
 
-review = imp.load_source(
-    "review", os.path.join(os.path.dirname(__file__), os.path.pardir, "moz-phab")
+mozphab = imp.load_source(
+    "mozphab", os.path.join(os.path.dirname(__file__), os.path.pardir, "moz-phab")
 )
 
-git_config = ["user.email=email"]
-
-
-@mock.patch("review.Git.git_out")
-@mock.patch("review.Git._get_current_head")
-@mock.patch("review.Config")
-@mock.patch("review.os.path")
-@mock.patch("review.which")
-@mock.patch("review.Repository._phab_url")
-class TestGit(unittest.TestCase):
-    def test_cherry(
-        self,
-        m_repository_phab_url,
-        m_which,
-        m_os_path,
-        m_config,
-        m_git_get_current_head,
-        m_git_git_out,
-    ):
-        m_os_path.join = os.path.join
-        m_os_path.exists.return_value = True
-        m_which.return_value = True
-        m_os_path.isfile.return_value = False
-        m_git_get_current_head.return_value = "branch"
-
-        m_git_git_out.side_effect = (git_config, review.CommandError, ["output"])
-        git = review.Git("x")
-        self.assertEqual(["output"], git._cherry(["cherry"], ["one", "two"]))
-        m_git_git_out.assert_has_calls(
-            [mock.call(["cherry", "one"]), mock.call(["cherry", "two"])]
-        )
-
-    @mock.patch("review.Git._cherry")
-    def test_first_unpublished(
-        self,
-        m_git_cherry,
-        m_repository_phab_url,
-        m_which,
-        m_os_path,
-        m_config,
-        m_git_get_current_head,
-        m_git_git_out,
-    ):
-        m_os_path.join = os.path.join
-        m_os_path.exists.return_value = True
-        m_which.return_value = True
-        m_os_path.isfile.return_value = False
-        m_git_get_current_head.return_value = "branch"
-
-        class Args:
-            def __init__(self, upstream=None, start_rev="(auto)"):
-                self.upstream = upstream
-                self.start_rev = start_rev
-
-        m_git_git_out.side_effect = (git_config, ["a", "b"], ["c"], ["d"])
-        m_git_cherry.side_effect = (["- sha1", "+ sha2"], [], None, [])
-        git = review.Git("x")
-        git.args = Args()
-        first = git._get_first_unpublished_node
-        self.assertEqual("sha2", first())
-        m_git_cherry.assert_called_with(["cherry", "--abbrev=12"], ["a", "b"])
-        self.assertIsNone(first())
-        with self.assertRaises(review.Error):
-            first()
-
-        git.args = Args(upstream=["upstream"])
+
+@mock.patch("mozphab.Git.git_out")
+def test_cherry(m_git_git_out, git):
+    m_git_git_out.side_effect = (mozphab.CommandError, ["output"])
+    assert git._cherry(["cherry"], ["one", "two"]) == ["output"]
+    m_git_git_out.assert_has_calls(
+        [mock.call(["cherry", "one"]), mock.call(["cherry", "two"])]
+    )
+
+
+@mock.patch("mozphab.Git.git_out")
+@mock.patch("mozphab.Git._cherry")
+def test_first_unpublished(m_git_cherry, m_git_git_out, git):
+    class Args:
+        def __init__(self, upstream=None, start_rev="(auto)"):
+            self.upstream = upstream
+            self.start_rev = start_rev
+
+    m_git_git_out.side_effect = (["a", "b"], ["c"], ["d"])
+    m_git_cherry.side_effect = (["- sha1", "+ sha2"], [], None, [])
+    git.args = Args()
+    first = git._get_first_unpublished_node
+    assert "sha2" == first()
+    m_git_cherry.assert_called_with(["cherry", "--abbrev=12"], ["a", "b"])
+    assert first() is None
+    with pytest.raises(mozphab.Error):
         first()
         m_git_cherry.assert_called_with(["cherry", "--abbrev=12", "upstream"], [])
 
-    def test_get_dirsect_children(
-        self,
-        m_repository_phab_url,
-        m_which,
-        m_os_path,
-        m_config,
-        m_git_get_current_head,
-        m_git_git_out,
-    ):
-        m_os_path.join = os.path.join
-        m_os_path.exists.return_value = True
-        m_which.return_value = True
-        m_os_path.isfile.return_value = False
-        m_git_get_current_head.return_value = "branch"
-        m_git_git_out.return_value = git_config
-        git = review.Git("x")
-
-        get_children = git._get_direct_children
-        rev_list = ["aaa bbb ccc", "bbb", "ccc ddd"]
-
-        self.assertEqual(["bbb", "ccc"], get_children("aaa", rev_list))
-        self.assertEqual([], get_children("bbb", rev_list))
-        self.assertEqual(["ddd"], get_children("ccc", rev_list))
-        self.assertEqual([], get_children("xxx", rev_list))
-
-
-    def test_is_child(
-        self,
-        m_repository_phab_url,
-        m_which,
-        m_os_path,
-        m_config,
-        m_git_get_current_head,
-        m_git_git_out,
-    ):
-        m_os_path.join = os.path.join
-        m_os_path.exists.return_value = True
-        m_which.return_value = True
-        m_os_path.isfile.return_value = False
-        m_git_get_current_head.return_value = "branch"
-        m_git_git_out.return_value = git_config
-        git = review.Git("x")
-
-        is_child = git._is_child
-        # * ccc
-        # * bbb
-        # * aaa
-        nodes = ["ccc", "bbb ccc", "aaa bbb"]
-        self.assertTrue(is_child("aaa", "bbb", nodes))
-        self.assertTrue(is_child("aaa", "ccc", nodes))
-        self.assertTrue(is_child("bbb", "ccc", nodes))
-        self.assertFalse(is_child("bbb", "aaa", nodes))
-        self.assertFalse(is_child("aaa", "aaa", nodes))
-        self.assertFalse(is_child("bbb", "bbb", nodes))
-        self.assertFalse(is_child("ccc", "ccc", nodes))
-
-        # * ddd
-        # | * ccc
-        # | | * eee
-        # | |/
-        # | * bbb
-        # |/
-        # * aaa
-        nodes = ["ddd", "ccc", "eee", "bbb ccc eee", "aaa bbb ddd"]
-        self.assertTrue(is_child("aaa", "bbb", nodes))
-        self.assertTrue(is_child("aaa", "ccc", nodes))
-        self.assertTrue(is_child("aaa", "ddd", nodes))
-        self.assertTrue(is_child("aaa", "eee", nodes))
-        self.assertTrue(is_child("bbb", "ccc", nodes))
-        self.assertTrue(is_child("bbb", "eee", nodes))
-        self.assertFalse(is_child("bbb", "ddd", nodes))
-        self.assertFalse(is_child("ccc", "ddd", nodes))
+    git.args = Args(upstream=["upstream"])
+    first()
+    m_git_cherry.assert_called_with(["cherry", "--abbrev=12", "upstream"], [])
+
+
+@mock.patch("mozphab.Git.git_out")
+def test_branches_to_rebase(m_git_git_out, git):
+    git_find = git._find_branches_to_rebase
+
+    # No branch returned - not a real case - we don't work without branches
+    m_git_git_out.return_value = []
+    assert dict() == git_find([{"orig-node": "_aaa", "node": "aaa"}])
+
+    # No amend, no branches to rebase
+    m_git_git_out.return_value = ["branch"]
+    assert dict() == git_find([{"orig-node": "aaa", "node": "aaa"}])
+
+    # One commit, one branch
+    m_git_git_out.return_value = ["branch"]
+    assert dict(branch=["aaa", "_aaa"]) == git_find(
+        [{"orig-node": "_aaa", "node": "aaa"}]
+    )
+
+    # Two commits one branch
+    m_git_git_out.return_value = ["branch"]
+    assert dict(branch=["bbb", "_bbb"]) == git_find(
+        [{"orig-node": "_aaa", "node": "aaa"}, {"orig-node": "_bbb", "node": "bbb"}]
+    )
+
+    # Two branches one commit
+    # ... (branch1)
+    # | ... (branch2)
+    # |/
+    # * aaa
+    # More realistic output from the git command
+    m_git_git_out.return_value = ["*  branch1", "  branch2"]
+    assert dict(branch1=["aaa", "_aaa"], branch2=["aaa", "_aaa"]) == git_find(
+        [{"orig-node": "_aaa", "node": "aaa"}]
+    )
+
+    # ... (branch1)
+    # | * bbb (branch2)
+    # |/
+    # * aaa
+    m_git_git_out.side_effect = (["branch1", "branch2"], ["branch2"])
+    assert dict(branch1=["aaa", "_aaa"], branch2=["bbb", "_bbb"]) == git_find(
+        [{"orig-node": "_aaa", "node": "aaa"}, {"orig-node": "_bbb", "node": "bbb"}]
+    )
+
+    # * ... (master)
+    # | * ... (feature1)
+    # | | * ... (feature2)
+    # | |/
+    # |/|
+    # | | * ddd (feature1_1)
+    # | |/
+    # | * ccc
+    # |/
+    # * bbb
+    # * aaa
+
+    m_git_git_out.side_effect = (
+        ["master", "feature1", "feature1_1", "feature2"],  # aaa
+        ["master", "feature1", "feature1_1", "feature2"],  # bbb
+        ["feature1", "feature1_1"],  # ccc
+        ["feature1_1"],  # ddd
+    )
+    assert dict(
+        master=["bbb", "_bbb"],
+        feature1=["ccc", "_ccc"],
+        feature2=["bbb", "_bbb"],
+        feature1_1=["ddd", "_ddd"],
+    ) == git_find(
+        [
+            {"orig-node": "_aaa", "node": "aaa"},
+            {"orig-node": "_bbb", "node": "bbb"},
+            {"orig-node": "_ccc", "node": "ccc"},
+            {"orig-node": "_ddd", "node": "ddd"},
+        ]
+    )
+
+
+def test_get_direct_children(git):
+    get_children = git._get_direct_children
+    rev_list = ["aaa bbb ccc", "bbb", "ccc ddd"]
+    assert ["bbb", "ccc"] == get_children("aaa", rev_list)
+    assert [] == get_children("bbb", rev_list)
+    assert ["ddd"] == get_children("ccc", rev_list)
+    assert [] == get_children("xxx", rev_list)
+
+
+def test_is_child(git):
+    is_child = git._is_child
+    # * ccc
+    # * bbb
+    # * aaa
+    nodes = ["ccc", "bbb ccc", "aaa bbb"]
+    assert is_child("aaa", "bbb", nodes)
+    assert is_child("aaa", "ccc", nodes)
+    assert is_child("bbb", "ccc", nodes)
+    assert not is_child("bbb", "aaa", nodes)
+    assert not is_child("aaa", "aaa", nodes)
+    assert not is_child("bbb", "bbb", nodes)
+    assert not is_child("ccc", "ccc", nodes)
+
+    # * ddd
+    # | * ccc
+    # | | * eee
+    # | |/
+    # | * bbb
+    # |/
+    # * aaa
+    nodes = ["ddd", "ccc", "eee", "bbb ccc eee", "aaa bbb ddd"]
+    assert is_child("aaa", "bbb", nodes)
+    assert is_child("aaa", "ccc", nodes)
+    assert is_child("aaa", "ddd", nodes)
+    assert is_child("aaa", "eee", nodes)
+    assert is_child("bbb", "ccc", nodes)
+    assert is_child("bbb", "eee", nodes)
+    assert not is_child("bbb", "ddd", nodes)
+    assert not is_child("ccc", "ddd", nodes)