@@ -43,6 +43,10 @@ def tearDown(self):
43
43
with self .disable_logging ():
44
44
CassandraInstrumentor ().uninstrument ()
45
45
46
+ @property
47
+ def _mocked_session (self ):
48
+ return cassandra .cluster .Session (cluster = mock .Mock (), hosts = [])
49
+
46
50
def test_instrument_uninstrument (self ):
47
51
instrumentation = CassandraInstrumentor ()
48
52
instrumentation .instrument ()
@@ -67,7 +71,7 @@ def test_instrumentor(
67
71
):
68
72
mock_create_response_future .return_value = mock .Mock ()
69
73
mock_session_init .return_value = None
70
- mock_connect .return_value = cassandra . cluster . Session ()
74
+ mock_connect .return_value = self . _mocked_session
71
75
72
76
CassandraInstrumentor ().instrument ()
73
77
@@ -100,7 +104,7 @@ def test_custom_tracer_provider(
100
104
):
101
105
mock_create_response_future .return_value = mock .Mock ()
102
106
mock_session_init .return_value = None
103
- mock_connect .return_value = cassandra . cluster . Session ()
107
+ mock_connect .return_value = self . _mocked_session
104
108
105
109
resource = resources .Resource .create ({})
106
110
result = self .create_tracer_provider (resource = resource )
@@ -124,7 +128,7 @@ def test_instrument_connection_no_op_tracer_provider(
124
128
):
125
129
mock_create_response_future .return_value = mock .Mock ()
126
130
mock_session_init .return_value = None
127
- mock_connect .return_value = cassandra . cluster . Session ()
131
+ mock_connect .return_value = self . _mocked_session
128
132
129
133
tracer_provider = trace_api .NoOpTracerProvider ()
130
134
CassandraInstrumentor ().instrument (tracer_provider = tracer_provider )
0 commit comments