diff --git a/powerline/segments/vim.py b/powerline/segments/vim.py
index 16116639..7c3462e9 100644
--- a/powerline/segments/vim.py
+++ b/powerline/segments/vim.py
@@ -11,7 +11,7 @@ except ImportError:
 from powerline.bindings.vim import vim_get_func, getbufvar
 from powerline.theme import requires_segment_info
 from powerline.lib import add_divider_highlight_group
-from powerline.lib.vcs import guess
+from powerline.lib.vcs import guess, tree_status
 from powerline.lib.humanize_bytes import humanize_bytes
 from powerline.lib import wraps_saveargs as wraps
 from collections import defaultdict
@@ -325,7 +325,6 @@ def branch(pl, segment_info, status_colors=True):
 			branch = repo.branch()
 			scol = ['branch']
 			if status_colors:
-				from powerline.lib.vcs import tree_status
 				status = tree_status(repo, pl)
 				scol.insert(0, 'branch_dirty' if status and status.strip() else 'branch_clean')
 			return [{
diff --git a/tests/test_segments.py b/tests/test_segments.py
index f6a29952..befad178 100644
--- a/tests/test_segments.py
+++ b/tests/test_segments.py
@@ -77,17 +77,19 @@ class TestCommon(TestCase):
 		pl = Pl()
 		segment_info = {'getcwd': os.getcwd}
 		with replace_attr(common, 'guess', lambda path: Args(branch=lambda: os.path.basename(path), status=lambda: None, directory='/tmp/tests')):
-			self.assertEqual(common.branch(pl=pl, segment_info=segment_info, status_colors=False),
-					[{'highlight_group': ['branch'], 'contents': 'tests'}])
-			self.assertEqual(common.branch(pl=pl, segment_info=segment_info, status_colors=True),
-					[{'contents': 'tests', 'highlight_group': ['branch_clean', 'branch']}])
+			with replace_attr(common, 'tree_status', lambda repo, pl: None):
+				self.assertEqual(common.branch(pl=pl, segment_info=segment_info, status_colors=False),
+						[{'highlight_group': ['branch'], 'contents': 'tests'}])
+				self.assertEqual(common.branch(pl=pl, segment_info=segment_info, status_colors=True),
+						[{'contents': 'tests', 'highlight_group': ['branch_clean', 'branch']}])
 		with replace_attr(common, 'guess', lambda path: Args(branch=lambda: os.path.basename(path), status=lambda: 'D  ', directory='/tmp/tests')):
-			self.assertEqual(common.branch(pl=pl, segment_info=segment_info, status_colors=False),
-					[{'highlight_group': ['branch'], 'contents': 'tests'}])
-			self.assertEqual(common.branch(pl=pl, segment_info=segment_info, status_colors=True),
-					[{'contents': 'tests', 'highlight_group': ['branch_dirty', 'branch']}])
-			self.assertEqual(common.branch(pl=pl, segment_info=segment_info, status_colors=False),
-					[{'highlight_group': ['branch'], 'contents': 'tests'}])
+			with replace_attr(common, 'tree_status', lambda repo, pl: 'D '):
+				self.assertEqual(common.branch(pl=pl, segment_info=segment_info, status_colors=False),
+						[{'highlight_group': ['branch'], 'contents': 'tests'}])
+				self.assertEqual(common.branch(pl=pl, segment_info=segment_info, status_colors=True),
+						[{'contents': 'tests', 'highlight_group': ['branch_dirty', 'branch']}])
+				self.assertEqual(common.branch(pl=pl, segment_info=segment_info, status_colors=False),
+						[{'highlight_group': ['branch'], 'contents': 'tests'}])
 		with replace_attr(common, 'guess', lambda path: None):
 			self.assertEqual(common.branch(pl=pl, segment_info=segment_info, status_colors=False), None)
 
@@ -454,15 +456,17 @@ class TestVim(TestCase):
 		pl = Pl()
 		with vim_module._with('buffer', '/foo') as segment_info:
 			with replace_attr(vim, 'guess', lambda path: Args(branch=lambda: os.path.basename(path), status=lambda: None, directory=path)):
-				self.assertEqual(vim.branch(pl=pl, segment_info=segment_info, status_colors=False),
-						[{'divider_highlight_group': 'branch:divider', 'highlight_group': ['branch'], 'contents': 'foo'}])
-				self.assertEqual(vim.branch(pl=pl, segment_info=segment_info, status_colors=True),
-						[{'divider_highlight_group': 'branch:divider', 'highlight_group': ['branch_clean', 'branch'], 'contents': 'foo'}])
+				with replace_attr(vim, 'tree_status', lambda repo, pl: None):
+					self.assertEqual(vim.branch(pl=pl, segment_info=segment_info, status_colors=False),
+							[{'divider_highlight_group': 'branch:divider', 'highlight_group': ['branch'], 'contents': 'foo'}])
+					self.assertEqual(vim.branch(pl=pl, segment_info=segment_info, status_colors=True),
+							[{'divider_highlight_group': 'branch:divider', 'highlight_group': ['branch_clean', 'branch'], 'contents': 'foo'}])
 			with replace_attr(vim, 'guess', lambda path: Args(branch=lambda: os.path.basename(path), status=lambda: 'DU', directory=path)):
-				self.assertEqual(vim.branch(pl=pl, segment_info=segment_info, status_colors=False),
-						[{'divider_highlight_group': 'branch:divider', 'highlight_group': ['branch'], 'contents': 'foo'}])
-				self.assertEqual(vim.branch(pl=pl, segment_info=segment_info, status_colors=True),
-						[{'divider_highlight_group': 'branch:divider', 'highlight_group': ['branch_dirty', 'branch'], 'contents': 'foo'}])
+				with replace_attr(vim, 'tree_status', lambda repo, pl: 'DU'):
+					self.assertEqual(vim.branch(pl=pl, segment_info=segment_info, status_colors=False),
+							[{'divider_highlight_group': 'branch:divider', 'highlight_group': ['branch'], 'contents': 'foo'}])
+					self.assertEqual(vim.branch(pl=pl, segment_info=segment_info, status_colors=True),
+							[{'divider_highlight_group': 'branch:divider', 'highlight_group': ['branch_dirty', 'branch'], 'contents': 'foo'}])
 
 	def test_file_vcs_status(self):
 		pl = Pl()