@@ -42,7 +42,9 @@ def tearDown(self):
4242 super ().tearDown ()
4343 with self .disable_logging ():
4444 CassandraInstrumentor ().uninstrument ()
45-
45+ @property
46+ def _mocked_session (self ):
47+ return cassandra .cluster .Session (cluster = mock .Mock (), hosts = [])
4648 def test_instrument_uninstrument (self ):
4749 instrumentation = CassandraInstrumentor ()
4850 instrumentation .instrument ()
@@ -67,7 +69,7 @@ def test_instrumentor(
6769 ):
6870 mock_create_response_future .return_value = mock .Mock ()
6971 mock_session_init .return_value = None
70- mock_connect .return_value = cassandra . cluster . Session ()
72+ mock_connect .return_value = self . _mocked_session
7173
7274 CassandraInstrumentor ().instrument ()
7375
@@ -100,7 +102,7 @@ def test_custom_tracer_provider(
100102 ):
101103 mock_create_response_future .return_value = mock .Mock ()
102104 mock_session_init .return_value = None
103- mock_connect .return_value = cassandra . cluster . Session ()
105+ mock_connect .return_value = self . _mocked_session
104106
105107 resource = resources .Resource .create ({})
106108 result = self .create_tracer_provider (resource = resource )
@@ -124,7 +126,7 @@ def test_instrument_connection_no_op_tracer_provider(
124126 ):
125127 mock_create_response_future .return_value = mock .Mock ()
126128 mock_session_init .return_value = None
127- mock_connect .return_value = cassandra . cluster . Session ()
129+ mock_connect .return_value = self . _mocked_session
128130
129131 tracer_provider = trace_api .NoOpTracerProvider ()
130132 CassandraInstrumentor ().instrument (tracer_provider = tracer_provider )
0 commit comments