diff --git a/trait_documenter/module_trait_documenter.py b/trait_documenter/module_trait_documenter.py
index 2a0d183..9e04274 100644
--- a/trait_documenter/module_trait_documenter.py
+++ b/trait_documenter/module_trait_documenter.py
@@ -62,7 +62,7 @@ def add_directive_header(self, sig):
         if hasattr(self, 'get_sourcename'):
             sourcename = self.get_sourcename()
         else:
-            sourcename = u'<autodoc>'
+            sourcename = '<autodoc>'
         if not self.options.annotation:
             try:
                 definition = get_trait_definition(
diff --git a/trait_documenter/tests/test_class_trait_documenter.py b/trait_documenter/tests/test_class_trait_documenter.py
index 3b1d2cf..c6d6e39 100644
--- a/trait_documenter/tests/test_class_trait_documenter.py
+++ b/trait_documenter/tests/test_class_trait_documenter.py
@@ -38,39 +38,77 @@ def test_can_document_member(self):
     def test_import_object(self):
         # given
         documenter = ClassTraitDocumenter(mock.Mock(), 'test')
-        documenter.modname = u'trait_documenter.tests.test_file'
-        documenter.objpath = [u'Dummy', u'trait_1']
+        documenter.modname = 'trait_documenter.tests.test_file'
+        documenter.objpath = ['Dummy', 'trait_1']
 
         # when
-        documenter.import_object()
+        result = documenter.import_object()
 
         # then
-        self.assertEqual(documenter.object_name, u'trait_1')
+        self.assertTrue(result)
+        self.assertEqual(documenter.object_name, 'trait_1')
         self.assertTrue(documenter.object is None)
         self.assertEqual(documenter.parent, Dummy)
 
+    def test_import_object_with_error(self):
+        # given
+        documenter = ClassTraitDocumenter(mock.Mock(), 'test')
+        documenter.modname = 'invalid456767'
+        documenter.objpath = ['Dummy', 'trait_1']
+
+        # when
+        result = documenter.import_object()
+
+        # then
+        self.assertFalse(result)
+        documenter.env.note_reread.assert_called()
+
     def test_add_directive_header(self):
         # given
-        documenter = ClassTraitDocumenter(mock.Mock(), u'test')
+        documenter = ClassTraitDocumenter(mock.Mock(), 'test')
+        documenter.parent = Dummy
+        documenter.object_name = 'trait_2'
+        documenter.modname = 'trait_documenter.tests.test_file'
+        documenter.get_sourcename = mock.Mock(return_value='<autodoc>')
+        documenter.objpath = ['Dummy', 'trait_2']
+        documenter.add_line = mock.Mock()
+
+        # when
+        documenter.add_directive_header('')
+
+        # then
+        self.assertEqual(documenter.directive.warn.call_args_list, [])
+        expected = [
+            ('.. py:attribute:: Dummy.trait_2', '<autodoc>'),
+            ('   :noindex:', '<autodoc>'),
+            ('   :module: trait_documenter.tests.test_file', '<autodoc>'),
+            ("   :annotation: = Property(Float,depends_on='trait_1')", '<autodoc>')]  # noqa
+        calls = documenter.add_line.call_args_list
+        for index, line in enumerate(expected):
+            self.assertEqual(calls[index][0], line)
+
+    def test_add_directive_header_with_warning(self):
+        # given
+        documenter = ClassTraitDocumenter(mock.Mock(), 'test')
         documenter.parent = Dummy
-        documenter.object_name = u'trait_2'
-        documenter.modname = u'trait_documenter.tests.test_file'
+        documenter.object_name = 'trait_very_invalid'
+        documenter.modname = 'trait_documenter.tests.test_file'
         documenter.get_sourcename = mock.Mock(return_value='<autodoc>')
-        documenter.objpath = [u'Dummy', u'trait_2']
+        documenter.objpath = ['Dummy', 'trait_very_invalid']
         documenter.add_line = mock.Mock()
 
         # when
         documenter.add_directive_header('')
 
         # then
+        self.assertEqual(documenter.directive.warn.call_count, 1)
         expected = [
-            (u'.. py:attribute:: Dummy.trait_2', u'<autodoc>'),
-            (u'   :noindex:', u'<autodoc>'),
-            (u'   :module: trait_documenter.tests.test_file', u'<autodoc>'),
-            (u"   :annotation: = Property(Float,depends_on='trait_1')", u'<autodoc>')]  # noqa
+            ('.. py:attribute:: Dummy.trait_very_invalid', '<autodoc>'),
+            ('   :noindex:', '<autodoc>'),
+            ('   :module: trait_documenter.tests.test_file', '<autodoc>')]
         calls = documenter.add_line.call_args_list
-        for index, call in enumerate(calls):
-            self.assertEqual(tuple(call)[0], expected[index])
+        for index, line in enumerate(expected):
+            self.assertEqual(calls[index][0], line)
 
 
 if __name__ == '__main__':
diff --git a/trait_documenter/tests/test_file.py b/trait_documenter/tests/test_file.py
index 3c2aeaa..d0044cc 100644
--- a/trait_documenter/tests/test_file.py
+++ b/trait_documenter/tests/test_file.py
@@ -6,6 +6,8 @@
     low=0.2,
     high=34)
 
+def dummy_function():
+    pass
 
 class Dummy(HasTraits):
 
diff --git a/trait_documenter/tests/test_get_trait_definition.py b/trait_documenter/tests/test_get_trait_definition.py
index 557e268..9fbb440 100644
--- a/trait_documenter/tests/test_get_trait_definition.py
+++ b/trait_documenter/tests/test_get_trait_definition.py
@@ -2,8 +2,8 @@
 import unittest
 
 from trait_documenter.tests import test_file, test_file2
-from trait_documenter.tests.test_file import Dummy
-from trait_documenter.util import get_trait_definition
+from trait_documenter.tests.test_file import Dummy, dummy_function
+from trait_documenter.util import get_trait_definition, DefinitionError
 
 
 class TestGetTraitDefinition(unittest.TestCase):
@@ -19,6 +19,17 @@ def test_get_simple_module_trait_definition(self):
         # then
         self.assertEqual(definition, 'Float')
 
+    def test_get_inavlid_module_trait_definition(self):
+        # given
+        parent = test_file
+        object_name = 'invalid_trait'
+
+        # when/then
+        self.assertRaises(
+            DefinitionError,
+            get_trait_definition,
+            parent, object_name)
+
     def test_get_multi_line_module_trait_definition(self):
         # given
         parent = test_file
@@ -41,6 +52,16 @@ def test_get_simple_class_trait_definition(self):
         # then
         self.assertEqual(definition, 'Float')
 
+    def test_get_trait_definition_from_function(self):
+
+        object_name = 'trait_1'
+
+        # when/then
+        self.assertRaises(
+            DefinitionError,
+            get_trait_definition,
+            dummy_function, object_name)
+
     def test_get_multi_line_class_trait_definition(self):
         # given
         parent = Dummy
diff --git a/trait_documenter/tests/test_module_trait_documenter.py b/trait_documenter/tests/test_module_trait_documenter.py
index 6ded167..ed1fb58 100644
--- a/trait_documenter/tests/test_module_trait_documenter.py
+++ b/trait_documenter/tests/test_module_trait_documenter.py
@@ -2,7 +2,7 @@
 import unittest
 
 import mock
-from sphinx.ext.autodoc import ModuleDocumenter
+from sphinx.ext.autodoc import ModuleDocumenter, SUPPRESS
 
 from trait_documenter.module_trait_documenter import ModuleTraitDocumenter
 from trait_documenter.tests import test_file
@@ -13,48 +13,49 @@ class TestModuleTraitDocumenter(unittest.TestCase):
 
     def test_can_document_member(self):
         can_document_member = ModuleTraitDocumenter.can_document_member
-        parent = ModuleDocumenter(mock.Mock(), u'test')
+        parent = ModuleDocumenter(mock.Mock(), 'test')
 
         # modules
         parent.object = test_file
-        self.assertFalse(can_document_member(Dummy, u'Dummy', True, parent))
-        self.assertFalse(can_document_member(Dummy, u'Dummy', False, parent))
+        self.assertFalse(can_document_member(Dummy, 'Dummy', True, parent))
+        self.assertFalse(can_document_member(Dummy, 'Dummy', False, parent))
         self.assertFalse(can_document_member(
-            test_file.module_trait, u'module_trait', False, parent))
+            test_file.module_trait, 'module_trait', False, parent))
         self.assertTrue(can_document_member(
-            test_file.module_trait, u'module_trait', True, parent))
+            test_file.module_trait, 'module_trait', True, parent))
 
         # class
         parent.object = Dummy
         self.assertFalse(
             can_document_member(
-                Dummy.class_traits()[u'trait_1'], u'trait_1', False, parent))
+                Dummy.class_traits()['trait_1'], 'trait_1', False, parent))
 
     def test_import_object(self):
         # given
-        documenter = ModuleTraitDocumenter(mock.Mock(), u'test')
+        documenter = ModuleTraitDocumenter(mock.Mock(), 'test')
         documenter.env.config = mock.Mock(autodoc_mock_imports=[])
-        documenter.modname = u'trait_documenter.tests.test_file'
-        documenter.fullname = u'trait_documenter.tests.test_file.long_module_trait'  # noqa
-        documenter.objpath = [u'long_module_trait']
+        documenter.modname = 'trait_documenter.tests.test_file'
+        documenter.fullname = 'trait_documenter.tests.test_file.long_module_trait'  # noqa
+        documenter.objpath = ['long_module_trait']
 
         # when
-        documenter.import_object()
+        result = documenter.import_object()
 
         # then
-        self.assertEqual(documenter.object_name, u'long_module_trait')
+        self.assertTrue(result)
+        self.assertEqual(documenter.object_name, 'long_module_trait')
         self.assertTrue(documenter.object is not None)
         self.assertEqual(documenter.parent, test_file)
 
     def test_add_directive_header(self):
         # given
-        documenter = ModuleTraitDocumenter(mock.Mock(), u'test')
+        documenter = ModuleTraitDocumenter(mock.Mock(), 'test')
         documenter.parent = test_file
         documenter.options = mock.Mock(annotation=False)
-        documenter.modname = u'trait_documenter.tests.test_file'
+        documenter.modname = 'trait_documenter.tests.test_file'
         documenter.get_sourcename = mock.Mock(return_value='<autodoc>')
-        documenter.object_name = u'long_module_trait'
-        documenter.objpath = [u'long_module_trait']
+        documenter.object_name = 'long_module_trait'
+        documenter.objpath = ['long_module_trait']
         documenter.add_line = mock.Mock()
 
         # when
@@ -62,13 +63,86 @@ def test_add_directive_header(self):
 
         # then
         expected = [
-            (u'.. py:data:: long_module_trait', u'<autodoc>'),
-            (u'   :noindex:', u'<autodoc>'),
-            (u'   :module: trait_documenter.tests.test_file', u'<autodoc>'),
-            (u'   :annotation: = Range(low=0.2,high=34)', u'<autodoc>')]  # noqa
+            ('.. py:data:: long_module_trait', '<autodoc>'),
+            ('   :noindex:', '<autodoc>'),
+            ('   :module: trait_documenter.tests.test_file', '<autodoc>'),
+            ('   :annotation: = Range(low=0.2,high=34)', '<autodoc>')]
         calls = documenter.add_line.call_args_list
-        for index, call in enumerate(calls):
-            self.assertEqual(tuple(call)[0], expected[index])
+        for index, line in enumerate(expected):
+            self.assertEqual(calls[index][0], line)
+
+    def test_add_directive_header_with_annotation(self):
+        # given
+        documenter = ModuleTraitDocumenter(mock.Mock(), 'test')
+        documenter.parent = test_file
+        documenter.options = mock.Mock(annotation='my annotation')
+        documenter.modname = 'trait_documenter.tests.test_file'
+        documenter.get_sourcename = mock.Mock(return_value='<autodoc>')
+        documenter.object_name = 'long_module_trait'
+        documenter.objpath = ['long_module_trait']
+        documenter.add_line = mock.Mock()
+
+        # when
+        documenter.add_directive_header('')
+
+        # then
+        self.assertEqual(documenter.directive.warn.call_args_list, [])
+        expected = [
+            ('.. py:data:: long_module_trait', '<autodoc>'),
+            ('   :noindex:', '<autodoc>'),
+            ('   :module: trait_documenter.tests.test_file', '<autodoc>'),
+            ('   :annotation: my annotation', '<autodoc>')]
+        calls = documenter.add_line.call_args_list
+        for index, line in enumerate(expected):
+            self.assertEqual(calls[index][0], line)
+
+    def test_add_directive_header_with_suppress(self):
+        # given
+        documenter = ModuleTraitDocumenter(mock.Mock(), 'test')
+        documenter.parent = test_file
+        documenter.options = mock.Mock(annotation=SUPPRESS)
+        documenter.modname = 'trait_documenter.tests.test_file'
+        documenter.get_sourcename = mock.Mock(return_value='<autodoc>')
+        documenter.object_name = 'long_module_trait'
+        documenter.objpath = ['long_module_trait']
+        documenter.add_line = mock.Mock()
+
+        # when
+        documenter.add_directive_header('')
+
+        # then
+        self.assertEqual(documenter.directive.warn.call_args_list, [])
+        expected = [
+            ('.. py:data:: long_module_trait', '<autodoc>'),
+            ('   :noindex:', '<autodoc>'),
+            ('   :module: trait_documenter.tests.test_file', '<autodoc>')]
+        calls = documenter.add_line.call_args_list
+        for index, line in enumerate(expected):
+            self.assertEqual(calls[index][0], line)
+
+    def test_add_directive_header_with_warning(self):
+        # given
+        documenter = ModuleTraitDocumenter(mock.Mock(), 'test')
+        documenter.parent = test_file
+        documenter.options = mock.Mock(annotation=False)
+        documenter.modname = 'trait_documenter.tests.test_file'
+        documenter.get_sourcename = mock.Mock(return_value='<autodoc>')
+        documenter.object_name = 'long_invalid_trait'
+        documenter.objpath = ['long_module_trait']
+        documenter.add_line = mock.Mock()
+
+        # when
+        documenter.add_directive_header('')
+
+        # the
+        self.assertEqual(documenter.directive.warn.call_count, 1)
+        expected = [
+            ('.. py:data:: long_module_trait', '<autodoc>'),
+            ('   :noindex:', '<autodoc>'),
+            ('   :module: trait_documenter.tests.test_file', '<autodoc>')]
+        calls = documenter.add_line.call_args_list
+        for index, line in enumerate(expected):
+            self.assertEqual(calls[index][0], line)
 
 
 if __name__ == '__main__':
diff --git a/trait_documenter/tests/test_sphinx_setup.py b/trait_documenter/tests/test_sphinx_setup.py
new file mode 100644
index 0000000..f83976d
--- /dev/null
+++ b/trait_documenter/tests/test_sphinx_setup.py
@@ -0,0 +1,21 @@
+import unittest
+
+import mock
+
+from trait_documenter import setup
+from trait_documenter.class_trait_documenter import ClassTraitDocumenter
+from trait_documenter.module_trait_documenter import ModuleTraitDocumenter
+
+
+class TestSphinxSetup(unittest.TestCase):
+
+    def test_setup(self):
+        app = mock.Mock()
+        setup(app)
+        expected = app.add_autodocumenter.call_args_list
+        calls = [
+            mock.call(ModuleTraitDocumenter), mock.call(ClassTraitDocumenter)]
+        self.assertEqual(expected, calls)
+
+if __name__ == '__main__':
+    unittest.main()
diff --git a/trait_documenter/util.py b/trait_documenter/util.py
index 28d9fa5..adb19a6 100644
--- a/trait_documenter/util.py
+++ b/trait_documenter/util.py
@@ -1,7 +1,8 @@
 import ast
 import inspect
 import collections
-from _ast import ClassDef, Assign, Name, If
+from _ast import ClassDef, Assign, Name
+
 
 class DefinitionError(Exception):
     pass
@@ -47,12 +48,12 @@ def get_trait_definition(parent, trait_name):
             if target is not None:
                 targets[node.col_offset].append((node, target))
 
-    # keep the assignment with the smallest column offset
-    assignments = targets[min(targets)]
-    if len(assignments) == 0:
+    if len(targets) == 0:
         message = 'Could not find trait definition of {0} in {1}'
         raise DefinitionError(message.format(trait_name, parent))
     else:
+        # keep the assignment with the smallest column offset
+        assignments = targets[min(targets)]
         # we always get the last assignment in the file
         node, name = assignments[-1]