@@ -12,18 +12,17 @@ import io.grpc.Server
12
12
import io.grpc.inprocess.InProcessChannelBuilder
13
13
import io.grpc.inprocess.InProcessServerBuilder
14
14
import io.grpc.stub.StreamObserver
15
- import org.mockito.internal.util.MockUtil
16
15
16
+ import java.lang.reflect.Method
17
17
import java.util.concurrent.CopyOnWriteArrayList
18
18
import java.util.concurrent.Executors
19
19
import java.util.concurrent.TimeUnit
20
20
import java.util.concurrent.atomic.AtomicReference
21
21
22
22
import static datadog.trace.api.config.TraceInstrumentationConfig.*
23
- import org.mockito.Mockito
24
23
25
24
abstract class GrpcCodeOriginTest extends VersionedNamingTestBase {
26
- private CodeOriginRecorder codeOriginRecorder
25
+ def codeOriginRecorder
27
26
28
27
@Override
29
28
final String service () {
@@ -157,8 +156,7 @@ abstract class GrpcCodeOriginTest extends VersionedNamingTestBase {
157
156
}. flatten(). sort()
158
157
159
158
160
- def invocations = MockUtil . getInvocationContainer(codeOriginRecorder)
161
- assert invocations. invocations. stream(). anyMatch { it. method. name == " captureCodeOrigin" }
159
+ assert codeOriginRecorder. invoked
162
160
assertTraces(2 ) {
163
161
trace((hasClientMessageSpans() ? clientMessageCount * serverMessageCount : 0 ) + 1 ) {
164
162
span {
@@ -253,7 +251,20 @@ abstract class GrpcCodeOriginTest extends VersionedNamingTestBase {
253
251
254
252
void codeOriginSetup () {
255
253
injectSysConfig(CODE_ORIGIN_FOR_SPANS_ENABLED , " true" , true )
256
- codeOriginRecorder = Mockito . mock(CodeOriginRecorder )
254
+ codeOriginRecorder = new CodeOriginRecorder () {
255
+ def invoked = false
256
+ @Override
257
+ String captureCodeOrigin (boolean entry ) {
258
+ invoked = true
259
+ return " done"
260
+ }
261
+
262
+ @Override
263
+ String captureCodeOrigin (Method method , boolean entry ) {
264
+ invoked = true
265
+ return " done"
266
+ }
267
+ }
257
268
DebuggerContext . initCodeOrigin(codeOriginRecorder)
258
269
}
259
270
}
0 commit comments