codewars · ggorlen · Jun 12, 2020 · Jun 12, 2020 · Jun 12, 2020 · Jun 12, 2020
diff --git a/README.md b/README.md
@@ -16,7 +16,7 @@ def example_tests():
 
 ### Using Other Assertions
 
-Any function that raises an `AssertionError` can be used instead of `codewars_test` assertions:
+Any function that raises an `AssertionError` can be used alongside built-in `codewars_test` assertions. Decorate the function using `codewars_test.make_assertion` before calling it. This ensures that test output is correctly formatted as expected by the runner. In the following example, the tests are intended to fail in order to show the custom output.
 
 ```python
 import numpy as np
@@ -25,26 +25,29 @@ import codewars_test as test
 
 @test.describe('Example Tests')
 def test_custom_assertions():
+    np.testing.assert_equal = test.make_assertion(np.testing.assert_equal)
+    pd.testing.assert_frame_equal = test.make_assertion(pd.testing.assert_frame_equal)
+
+    @test.make_assertion
+    def custom_assert_eq(actual, expected, msg=None):
+        if actual != expected:
+            default_msg = f'`{actual}` did not equal expected `{expected}`'
+            raise AssertionError(default_msg if msg is None else msg)
 
     @test.it('Test something in numpy')
     def test_numpy_assertion():
         actual = np.reshape(range(16), [4, 4])
         expected = np.reshape(range(16, 0, -1), [4, 4])
-        np.testing.assert_equal(expected, actual)
+        np.testing.assert_equal(actual, expected)
 
     @test.it('Test something in pandas')
     def test_pandas_assertion():
         actual = pd.DataFrame({'foo': [1, 2, 3]})
         expected = pd.DataFrame({'foo': [1, 42, 3]})
-        pd.testing.assert_frame_equal(expected, actual)
+        pd.testing.assert_frame_equal(actual, expected)
 
     @test.it('Test something using a custom assertion')
     def test_custom_assertion():
-        def custom_assert_eq(actual, expected, msg=None):
-            if actual != expected:
-                default_msg = f'`{actual}` did not equal expected `{expected}`'
-                raise AssertionError(default_msg if msg is None else msg)
-
         actual = 2
         expected = 1
         custom_assert_eq(actual, expected)

diff --git a/codewars_test/test_framework.py b/codewars_test/test_framework.py
@@ -1,4 +1,9 @@
 from __future__ import print_function
+import functools
+import sys
+from multiprocessing import Process
+from timeit import default_timer
+from traceback import format_exception
 
 
 class AssertException(Exception):
@@ -85,60 +90,73 @@ def assert_approx_equals(
     expect(abs((actual - expected) / div) < margin, message, allow_raise)
 
 
-'''
-Usage:
-@describe('describe text')
-def describe1():
-    @it('it text')
-    def it1():
-        # some test cases...
-'''
+def make_assertion(func):
+    '''
+    Wraps an assertion function to emit pass/failure stdout prints.
+    The function should raise an AssertionError to cause a failure.
+
+    @test.make_assertion
+    def custom_assert_eq(actual, expected, msg=None):
+        if actual != expected:
+            default_msg = f'`{actual}` did not equal expected `{expected}`'
+            raise AssertionError(default_msg if msg is None else msg)
+
+    # or decorate with a normal function call:
+    custom_assert_eq = make_assertion(custom_assert_eq)
+    '''
+    @functools.wraps(func)
+    def wrapper(*args, **kwargs):
+        try:
+            func(*args, **kwargs)
+            pass_()
+        except AssertionError as e:
+            fail(str(e))
+    return wrapper
 
 
 def _timed_block_factory(opening_text):
-    from timeit import default_timer as timer
-    from traceback import format_exception
-    from sys import exc_info
-
     def _timed_block_decorator(s, before=None, after=None):
         display(opening_text, s)
 
         def wrapper(func):
             if callable(before):
                 before()
-            time = timer()
+            time = default_timer()
             try:
                 func()
-            except AssertionError as e:
-                display('FAILED', str(e))
             except Exception:
                 fail('Unexpected exception raised')
-                tb_str = ''.join(format_exception(*exc_info()))
+                tb_str = ''.join(format_exception(*sys.exc_info()))
                 display('ERROR', tb_str)
-            display('COMPLETEDIN', '{:.2f}'.format((timer() - time) * 1000))
+            display('COMPLETEDIN', '{:.2f}'.format((default_timer() - time) * 1000))
             if callable(after):
                 after()
         return wrapper
     return _timed_block_decorator
 
 
-describe = _timed_block_factory('DESCRIBE')
-it = _timed_block_factory('IT')
-
-
 '''
-Timeout utility
 Usage:
-@timeout(sec)
-def some_tests():
-    any code block...
-Note: Timeout value can be a float.
+@describe('describe text')
+def describe1():
+    @it('it text')
+    def it1():
+        # some test cases...
 '''
+describe = _timed_block_factory('DESCRIBE')
+it = _timed_block_factory('IT')
 
 
 def timeout(sec):
+    '''
+    Timeout utility
+    Usage:
+    @timeout(sec)
+    def some_tests():
+        any code block...
+    Note: Timeout value can be a float.
+    '''
     def wrapper(func):
-        from multiprocessing import Process
         msg = 'Should not throw any exceptions inside timeout'
 
         def wrapped():
@@ -151,3 +169,4 @@ def wrapped():
             process.terminate()
             process.join()
     return wrapper
+