fix test for dt.timer

flaggo · Apr 23, 2018 · 12e43e4 · 12e43e4
1 parent 9e9a6a5
commit 12e43e4
Showing 1 changed file with 10 additions and 2 deletions.
diff --git a/tests/test_dt.py b/tests/test_dt.py
@@ -12,12 +12,20 @@ def test_context_manager(self):
         assert timeit.elapsed is not None
 
     def test_decorator(self):
-        import sys
-        timeit = timer(print_func=print)
+        timeit = timer()
 
         @timeit
         def foo():
             os.getcwd()
 
         foo()
         assert timeit.elapsed is not None
+
+    def test_print_func(self):
+        import sys
+        timeit = timer(print_func=sys.stdout.write)
+
+        with timeit:
+            os.getcwd()
+
+        assert timeit.elapsed is not None