@@ -43,6 +43,10 @@ def tearDown(self):
4343 with self .disable_logging ():
4444 CassandraInstrumentor ().uninstrument ()
4545
46+ @property
47+ def _mocked_session (self ):
48+ return cassandra .cluster .Session (cluster = mock .Mock (), hosts = [])
49+
4650 def test_instrument_uninstrument (self ):
4751 instrumentation = CassandraInstrumentor ()
4852 instrumentation .instrument ()
@@ -67,7 +71,7 @@ def test_instrumentor(
6771 ):
6872 mock_create_response_future .return_value = mock .Mock ()
6973 mock_session_init .return_value = None
70- mock_connect .return_value = cassandra . cluster . Session ()
74+ mock_connect .return_value = self . _mocked_session
7175
7276 CassandraInstrumentor ().instrument ()
7377
@@ -100,7 +104,7 @@ def test_custom_tracer_provider(
100104 ):
101105 mock_create_response_future .return_value = mock .Mock ()
102106 mock_session_init .return_value = None
103- mock_connect .return_value = cassandra . cluster . Session ()
107+ mock_connect .return_value = self . _mocked_session
104108
105109 resource = resources .Resource .create ({})
106110 result = self .create_tracer_provider (resource = resource )
@@ -124,7 +128,7 @@ def test_instrument_connection_no_op_tracer_provider(
124128 ):
125129 mock_create_response_future .return_value = mock .Mock ()
126130 mock_session_init .return_value = None
127- mock_connect .return_value = cassandra . cluster . Session ()
131+ mock_connect .return_value = self . _mocked_session
128132
129133 tracer_provider = trace_api .NoOpTracerProvider ()
130134 CassandraInstrumentor ().instrument (tracer_provider = tracer_provider )
0 commit comments