fix unit tests using non-existent mock function

aholmberg · aholmberg · commit 81c5aef702df · 2016-06-29T14:44:18.000-05:00
PYTHON-591
diff --git a/tests/unit/test_host_connection_pool.py b/tests/unit/test_host_connection_pool.py
@@ -165,7 +165,8 @@ def test_spawn_when_at_max(self):
     def test_return_defunct_connection(self):
         host = Mock(spec=Host, address='ip1')
         session = self.make_session()
-        conn = NonCallableMagicMock(spec=Connection, in_flight=0, is_defunct=False, is_closed=False, max_request_id=100)
+        conn = NonCallableMagicMock(spec=Connection, in_flight=0, is_defunct=False, is_closed=False,
+                                    max_request_id=100, signaled_error=False)
         session.cluster.connection_factory.return_value = conn
 
         pool = HostConnectionPool(host, HostDistance.LOCAL, session)
@@ -177,14 +178,14 @@ def test_return_defunct_connection(self):
         pool.return_connection(conn)
 
         # the connection should be closed a new creation scheduled
-        conn.close.assert_called_once()
-        session.submit.assert_called_once()
+        self.assertTrue(session.submit.call_args)
         self.assertFalse(pool.is_shutdown)
 
     def test_return_defunct_connection_on_down_host(self):
         host = Mock(spec=Host, address='ip1')
         session = self.make_session()
-        conn = NonCallableMagicMock(spec=Connection, in_flight=0, is_defunct=False, is_closed=False, max_request_id=100, signaled_error=False)
+        conn = NonCallableMagicMock(spec=Connection, in_flight=0, is_defunct=False, is_closed=False,
+                                    max_request_id=100, signaled_error=False)
         session.cluster.connection_factory.return_value = conn
 
         pool = HostConnectionPool(host, HostDistance.LOCAL, session)
@@ -196,15 +197,15 @@ def test_return_defunct_connection_on_down_host(self):
         pool.return_connection(conn)
 
         # the connection should be closed a new creation scheduled
-        session.cluster.signal_connection_failure.assert_called_once()
-        conn.close.assert_called_once()
+        self.assertTrue(session.cluster.signal_connection_failure.call_args)
+        self.assertTrue(conn.close.call_args)
         self.assertFalse(session.submit.called)
         self.assertTrue(pool.is_shutdown)
 
     def test_return_closed_connection(self):
         host = Mock(spec=Host, address='ip1')
         session = self.make_session()
-        conn = NonCallableMagicMock(spec=Connection, in_flight=0, is_defunct=False, is_closed=True, max_request_id=100)
+        conn = NonCallableMagicMock(spec=Connection, in_flight=0, is_defunct=False, is_closed=True, max_request_id=100, signaled_error=False)
         session.cluster.connection_factory.return_value = conn
 
         pool = HostConnectionPool(host, HostDistance.LOCAL, session)
@@ -216,7 +217,7 @@ def test_return_closed_connection(self):
         pool.return_connection(conn)
 
         # a new creation should be scheduled
-        session.submit.assert_called_once()
+        self.assertTrue(session.submit.call_args)
         self.assertFalse(pool.is_shutdown)
 
     def test_host_instantiations(self):
diff --git a/tests/unit/test_response_future.py b/tests/unit/test_response_future.py
@@ -480,7 +480,7 @@ def test_prepared_query_not_found(self):
         result = Mock(spec=PreparedQueryNotFound, info='a' * 16)
         rf._set_result(result)
 
-        session.submit.assert_called_once()
+        self.assertTrue(session.submit.call_args)
         args, kwargs = session.submit.call_args
         self.assertEqual(rf._reprepare, args[-2])
         self.assertIsInstance(args[-1], PrepareMessage)
