Making a few fixes based on PR feedback

GregBestland · GregBestland · commit 4ae307fdec4a · 2016-04-12T11:27:26.000-05:00
diff --git a/tests/integration/cqlengine/__init__.py b/tests/integration/cqlengine/__init__.py
@@ -90,9 +90,13 @@ def wrapped_function(*args, **kwargs):
             count.get_counter()
             # DeMonkey Patch our code
             cassandra.cqlengine.connection.execute = original_function
+            # Check to see if we have a pre-existing test case to work from.
+            if len(args) is 0:
+                test_case = unittest.TestCase("__init__")
+            else:
+                test_case = args[0]
             # Check to see if the count is what you expect
-            tc = unittest.TestCase("__init__")
-            tc.assertEquals(count.get_counter(), expected, "Expected number of cassandra.cqlengine.connection.execute calls doesn't match actual number invoked Expected: {0}, Invoked {1}".format(count.get_counter(), expected))
+            test_case.assertEqual(count.get_counter(), expected, msg="Expected number of cassandra.cqlengine.connection.execute calls doesn't match actual number invoked Expected: {0}, Invoked {1}".format(count.get_counter(), expected))
             return to_return
         # Name of the wrapped function must match the original or unittest will error out.
         wrapped_function.__name__ = fn.__name__
