Create unittest for shortened(...)

diff --git a/test/test_text_helpers.py b/test/test_text_helpers.py
index 7bc4390..1641cf0 100644 (file)
--- a/test/test_text_helpers.py
+++ b/test/test_text_helpers.py
@@ -35,18 +35,64 @@ from text_helpers import *
 
 
 class TextHelpersTester(unittest.TestCase):
-    """
-    hard to test whether functions work correctly, can only check no raise
-    """
+    """ test functions in TextHelpers """
 
     def test_shortcuts_do_not_raise(self):
-        """ tests all shortcut functions """
+        """ tests whether color functions do not raise errors """
         for color in COLOR_BLACK, COLOR_RED, COLOR_GREEN, COLOR_YELLOW, \
                 COLOR_BLUE, COLOR_MAGENTA, COLOR_CYAN, COLOR_WHITE:
             exec('{0}("{0}")'.format(color))
         for style in 'normal', 'bold', 'underline', 'blink', 'reverse':
             exec('{0}("{0}")'.format(style))
 
+    def test_shortening(self):
+        """ test function shortened """
+
+        # test iterating over shortened works
+        orig = tuple(range(10))
+        for elem in shortened(orig):
+            pass
+
+        # test non-shortening
+        self.assertEqual(list(shortened(range(3))), list(range(3)))
+
+        # test shortening with indexable iter
+        short = tuple(shortened(tuple(range(100)), skip_elem=None))
+        expect = tuple(range(20)) + tuple(range(80,100))
+        self.assertEqual(short, expect)
+        self.assertEqual(len(short), 40)
+
+        # test strange cases: no head, no tail, no skip
+        short = tuple(shortened(orig, n_head=2, n_tail=0, skip_elem=None))
+        self.assertEqual(short, (0, 1))
+        short = tuple(shortened(orig, n_head=0, n_tail=2, skip_elem=None))
+        self.assertEqual(short, (8, 9))
+        short = tuple(shortened(orig, n_head=0, n_tail=0, skip_elem=None))
+        self.assertEqual(len(short), 0)
+        short = tuple(shortened(orig, n_head=0, n_tail=0, skip_elem="Hi!"))
+        self.assertEqual(short, ("Hi!",))
+
+        # test n_head + n_tail == n_elems
+        short = tuple(shortened(orig, n_head=5, n_tail=5))
+        self.assertEqual(short, orig)
+        short = tuple(shortened(orig, n_head=10, n_tail=0))
+        self.assertEqual(short, orig)
+        short = tuple(shortened(orig, n_head=9, n_tail=1))
+        self.assertEqual(short, orig)
+        short = tuple(shortened(orig, n_head=1, n_tail=9))
+        self.assertEqual(short, orig)
+        short = tuple(shortened(orig, n_head=0, n_tail=10))
+        self.assertEqual(short, orig)
+
+        # test shortening with non-indexable iter without len()
+        def my_iter():
+            for elem in range(100):
+                yield elem
+
+        short = tuple(shortened(my_iter(), n_elems=100, skip_elem=None))
+        expect = tuple(range(20)) + tuple(range(80,100))
+        self.assertEqual(short, expect)
+
 
 if __name__ == '__main__':
     unittest.main()