From f80601a93ea25f643ab7a00ba9a58a0da37a99b6 Mon Sep 17 00:00:00 2001
From: Will McGugan <willmcgugan@gmail.com>
Date: Sat, 9 Jan 2021 14:10:07 +0000
Subject: [PATCH] tree tests

---
 rich/segment.py    | 11 ----------
 tests/test_tree.py | 53 ++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 53 insertions(+), 11 deletions(-)
 create mode 100644 tests/test_tree.py

diff --git a/rich/segment.py b/rich/segment.py
index ae4a99a6..db17b0fb 100644
--- a/rich/segment.py
+++ b/rich/segment.py
@@ -40,17 +40,6 @@ class Segment(NamedTuple):
         """Get cell length of segment."""
         return 0 if self.is_control else cell_len(self.text)
 
-    def with_style(self, style: Style) -> "Segment":
-        """Get a copy of the segment with a new style.
-
-        Args:
-            style (Style): New style to set.
-
-        Returns:
-            Segment: New segment instance.
-        """
-        return Segment(self.text, style, self.is_control)
-
     @classmethod
     def control(cls, text: str, style: Optional[Style] = None) -> "Segment":
         """Create a Segment with control codes.
diff --git a/tests/test_tree.py b/tests/test_tree.py
new file mode 100644
index 00000000..58573b58
--- /dev/null
+++ b/tests/test_tree.py
@@ -0,0 +1,53 @@
+from rich.console import Console
+from rich.tree import Tree
+
+
+def test_render_single_node():
+    tree = Tree("foo")
+    console = Console(color_system=None, width=20)
+    console.begin_capture()
+    console.print(tree)
+    assert console.end_capture() == "foo                 \n"
+
+
+def test_render_single_branch():
+    tree = Tree("foo")
+    tree.add("bar")
+    console = Console(color_system=None, width=20)
+    console.begin_capture()
+    console.print(tree)
+    result = console.end_capture()
+    print(repr(result))
+    expected = "foo                 \n└── bar             \n"
+    assert result == expected
+
+
+def test_render_double_branch():
+    tree = Tree("foo")
+    tree.add("bar")
+    tree.add("baz")
+    console = Console(color_system=None, width=20)
+    console.begin_capture()
+    console.print(tree)
+    result = console.end_capture()
+    print(repr(result))
+    expected = "foo                 \n├── bar             \n└── baz             \n"
+    assert result == expected
+
+
+def test_render_ascii():
+    tree = Tree("foo")
+    tree.add("bar")
+    tree.add("baz")
+
+    class AsciiConsole(Console):
+        @property
+        def encoding(self):
+            return "ascii"
+
+    console = AsciiConsole(color_system=None, width=20)
+    console.begin_capture()
+    console.print(tree)
+    result = console.end_capture()
+    expected = "foo                 \n+-- bar             \n`-- baz             \n"
+    assert result == expected
\ No newline at end of file