Trace test updates for PYTHON-318

aholmberg · aholmberg · commit 385baa1659df · 2015-10-14T10:18:23.000-05:00
diff --git a/tests/integration/standard/test_client_warnings.py b/tests/integration/standard/test_client_warnings.py
@@ -91,7 +91,7 @@ def test_warning_with_trace(self):
         future.result()
         self.assertEqual(len(future.warnings), 1)
         self.assertRegexpMatches(future.warnings[0], 'Batch.*exceeding.*')
-        self.assertIsNotNone(future._query_trace)
+        self.assertIsNotNone(future.get_query_trace())
 
     def test_warning_with_custom_payload(self):
         """
@@ -127,5 +127,5 @@ def test_warning_with_trace_and_custom_payload(self):
         future.result()
         self.assertEqual(len(future.warnings), 1)
         self.assertRegexpMatches(future.warnings[0], 'Batch.*exceeding.*')
-        self.assertIsNotNone(future._query_trace)
+        self.assertIsNotNone(future.get_query_trace())
         self.assertDictEqual(future.custom_payload, payload)
diff --git a/tests/integration/standard/test_cluster.py b/tests/integration/standard/test_cluster.py
@@ -438,11 +438,11 @@ def test_trace(self):
         session = cluster.connect()
 
         def check_trace(trace):
-            self.assertIsNot(None, trace.request_type)
-            self.assertIsNot(None, trace.duration)
-            self.assertIsNot(None, trace.started_at)
-            self.assertIsNot(None, trace.coordinator)
-            self.assertIsNot(None, trace.events)
+            self.assertIsNotNone(trace.request_type)
+            self.assertIsNotNone(trace.duration)
+            self.assertIsNotNone(trace.started_at)
+            self.assertIsNotNone(trace.coordinator)
+            self.assertIsNotNone(trace.events)
 
         result = session.execute( "SELECT * FROM system.local", trace=True)
         check_trace(result.get_query_trace())
diff --git a/tests/unit/test_resultset.py b/tests/unit/test_resultset.py
@@ -37,7 +37,8 @@ def test_iter_paged(self):
         response_future.result.side_effect = (ResultSet(Mock(), expected[-5:]), )  # ResultSet is iterable, so it must be protected in order to be returned whole by the Mock
         rs = ResultSet(response_future, expected[:5])
         itr = iter(rs)
-        type(response_future).has_more_pages = PropertyMock(side_effect=(True, False))  # after init to avoid side effects being consumed by init
+        # this is brittle, depends on internal impl details. Would like to find a better way
+        type(response_future).has_more_pages = PropertyMock(side_effect=(True, True, False))  # after init to avoid side effects being consumed by init
         self.assertListEqual(list(itr), expected)
 
     def test_list_non_paged(self):
@@ -54,7 +55,8 @@ def test_list_paged(self):
         response_future = Mock(has_more_pages=True)
         response_future.result.side_effect = (ResultSet(Mock(), expected[-5:]), )  # ResultSet is iterable, so it must be protected in order to be returned whole by the Mock
         rs = ResultSet(response_future, expected[:5])
-        type(response_future).has_more_pages = PropertyMock(side_effect=(True, True, False))  # one True for getitem check/warn, then True, False for two pages
+        # this is brittle, depends on internal impl details. Would like to find a better way
+        type(response_future).has_more_pages = PropertyMock(side_effect=(True, True, True, False))  # First two True are consumed on check entering list mode
         self.assertEqual(rs[9], expected[9])
         self.assertEqual(list(rs), expected)
 
@@ -119,7 +121,8 @@ def test_index_list_mode(self):
         response_future = Mock(has_more_pages=True)
         response_future.result.side_effect = (ResultSet(Mock(), expected[-5:]), )  # ResultSet is iterable, so it must be protected in order to be returned whole by the Mock
         rs = ResultSet(response_future, expected[:5])
-        type(response_future).has_more_pages = PropertyMock(side_effect=(True, True, False))  # First True is consumed on check entering list mode
+        # this is brittle, depends on internal impl details. Would like to find a better way
+        type(response_future).has_more_pages = PropertyMock(side_effect=(True, True, True, False))  # First two True are consumed on check entering list mode
         # index access before iteration causes list to be materialized
         self.assertEqual(rs[0], expected[0])
         self.assertEqual(rs[9], expected[9])
@@ -146,7 +149,7 @@ def test_eq(self):
         response_future = Mock(has_more_pages=True)
         response_future.result.side_effect = (ResultSet(Mock(), expected[-5:]), )  # ResultSet is iterable, so it must be protected in order to be returned whole by the Mock
         rs = ResultSet(response_future, expected[:5])
-        type(response_future).has_more_pages = PropertyMock(side_effect=(True, True, False))
+        type(response_future).has_more_pages = PropertyMock(side_effect=(True, True, True, False))
         # eq before iteration causes list to be materialized
         self.assertEqual(rs, expected)
 
