Adds tests for contextmanager

tests/test_io.py

@@ -1,12 +1,21 @@
 # -*- coding: utf-8 -*-
 from __future__ import unicode_literals
-from tableprint import table, banner, dataframe, hrule
+from tableprint import table, banner, dataframe, hrule, TableContext
 from io import StringIO
 import numpy as np
 
 
-def test_table():
+def test_context():
+    """Tests the table context manager"""
+    output = StringIO()
+    with TableContext('ABC', style='round', width=5, out=output) as t:
+        t([1, 2, 3])
+        t([4, 5, 6])
+    assert output.getvalue() == '╭─────┬─────┬─────╮\n│  A  │  B  │  C  │\n├─────┼─────┼─────┤\n│    1│    2│    3│\n│    4│    5│    6│\n╰─────┴─────┴─────╯\n'
 
+
+def test_table():
+    """Tests the table function"""
     output = StringIO()
     table([[1, 2, 3], [4, 5, 6]], 'ABC', style='round', width=5, out=output)
     assert output.getvalue() == '╭─────┬─────┬─────╮\n│  A  │  B  │  C  │\n├─────┼─────┼─────┤\n│    1│    2│    3│\n│    4│    5│    6│\n╰─────┴─────┴─────╯\n'
@@ -17,7 +26,7 @@ def test_table():
 
 
 def test_frame():
-
+    """Tests the dataframe function"""
     # mock of a pandas DataFrame
     class DataFrame:
         def __init__(self, data, headers):
@@ -35,7 +44,7 @@ def test_frame():
 
 
 def test_banner():
-
+    """Tests the banner function"""
     output = StringIO()
     banner('hello world', style='clean', width=11, out=output)
     assert output.getvalue() == ' ─────────── \n hello world \n ─────────── \n'
@@ -46,7 +55,7 @@ def test_banner():
 
 
 def test_hrule():
-
+    """Tests the hrule function"""
     output = hrule(1, width=11)
     assert len(output) == 11
     assert '───────────'