diff --git a/psh/commands.py b/psh/commands.py
index 89cb0bf..6bf481d 100644
--- a/psh/commands.py
+++ b/psh/commands.py
@@ -1,6 +1,6 @@
 from io import StringIO
 
-from tree import TreeNode
+from psh.tree import TreeNode
 
 
 class BaseCommand(object):
diff --git a/psh/example_cmd.py b/psh/example_cmd.py
index 90c496d..e91de76 100644
--- a/psh/example_cmd.py
+++ b/psh/example_cmd.py
@@ -1,6 +1,6 @@
 from psh.commands import BaseCommand, register_cmd
 
-from tree import TreeNode
+from psh.tree import TreeNode
 
 
 @register_cmd("example")
diff --git a/psh/raw_commands.py b/psh/raw_commands.py
index d509511..e6787c1 100644
--- a/psh/raw_commands.py
+++ b/psh/raw_commands.py
@@ -2,8 +2,7 @@ import shlex
 
 from psh.formatters import Printer
 from psh.commands import BaseCommand
-
-from tree import TreeNode
+from psh.tree import TreeNode
 
 
 class RawCommand(BaseCommand):
diff --git a/test/utils.py b/test/utils.py
index be7696c..5228c55 100644
--- a/test/utils.py
+++ b/test/utils.py
@@ -14,7 +14,7 @@ class TestFormatter(BaseCommand):
     def call(self):
         input_generator = self.get_input_generator()
         for line in input_generator:
-            self.buffer.write(line.decode('utf-8'))
+            self.buffer.write(line.data.decode('utf-8'))
         return None
 
     def get_data(self):