@@ -12,18 +12,17 @@ import io.grpc.Server
1212import io.grpc.inprocess.InProcessChannelBuilder
1313import io.grpc.inprocess.InProcessServerBuilder
1414import io.grpc.stub.StreamObserver
15- import org.mockito.internal.util.MockUtil
1615
16+ import java.lang.reflect.Method
1717import java.util.concurrent.CopyOnWriteArrayList
1818import java.util.concurrent.Executors
1919import java.util.concurrent.TimeUnit
2020import java.util.concurrent.atomic.AtomicReference
2121
2222import static datadog.trace.api.config.TraceInstrumentationConfig.*
23- import org.mockito.Mockito
2423
2524abstract class GrpcCodeOriginTest extends VersionedNamingTestBase {
26- private CodeOriginRecorder codeOriginRecorder
25+ def codeOriginRecorder
2726
2827 @Override
2928 final String service () {
@@ -157,8 +156,7 @@ abstract class GrpcCodeOriginTest extends VersionedNamingTestBase {
157156 }. flatten(). sort()
158157
159158
160- def invocations = MockUtil . getInvocationContainer(codeOriginRecorder)
161- assert invocations. invocations. stream(). anyMatch { it. method. name == " captureCodeOrigin" }
159+ assert codeOriginRecorder. invoked
162160 assertTraces(2 ) {
163161 trace((hasClientMessageSpans() ? clientMessageCount * serverMessageCount : 0 ) + 1 ) {
164162 span {
@@ -253,7 +251,20 @@ abstract class GrpcCodeOriginTest extends VersionedNamingTestBase {
253251
254252 void codeOriginSetup () {
255253 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+ }
257268 DebuggerContext . initCodeOrigin(codeOriginRecorder)
258269 }
259270}
0 commit comments