Update testing module to new agent api

This commit is contained in:
Trask Stalnaker 2019-10-19 15:30:40 -07:00
parent 41185662df
commit 299c396eab
1 changed files with 12 additions and 8 deletions

View File

@ -3,13 +3,15 @@ package datadog.trace.agent.test.utils
import datadog.opentracing.DDSpan import datadog.opentracing.DDSpan
import datadog.trace.agent.decorator.BaseDecorator import datadog.trace.agent.decorator.BaseDecorator
import datadog.trace.agent.test.asserts.TraceAssert import datadog.trace.agent.test.asserts.TraceAssert
import datadog.trace.context.TraceScope import datadog.trace.instrumentation.api.AgentScope
import io.opentracing.Scope import datadog.trace.instrumentation.api.AgentSpan
import io.opentracing.util.GlobalTracer
import lombok.SneakyThrows import lombok.SneakyThrows
import java.util.concurrent.Callable import java.util.concurrent.Callable
import static datadog.trace.instrumentation.api.AgentTracer.activateSpan
import static datadog.trace.instrumentation.api.AgentTracer.startSpan
class TraceUtils { class TraceUtils {
private static final BaseDecorator DECORATOR = new BaseDecorator() { private static final BaseDecorator DECORATOR = new BaseDecorator() {
@ -28,17 +30,19 @@ class TraceUtils {
@SneakyThrows @SneakyThrows
static <T> T runUnderTrace(final String rootOperationName, final Callable<T> r) { static <T> T runUnderTrace(final String rootOperationName, final Callable<T> r) {
final Scope scope = GlobalTracer.get().buildSpan(rootOperationName).startActive(true) final AgentSpan span = startSpan(rootOperationName)
DECORATOR.afterStart(scope) DECORATOR.afterStart(span)
((TraceScope) scope).setAsyncPropagation(true)
AgentScope scope = activateSpan(span, true)
scope.setAsyncPropagation(true)
try { try {
return r.call() return r.call()
} catch (final Exception e) { } catch (final Exception e) {
DECORATOR.onError(scope, e) DECORATOR.onError(span, e)
throw e throw e
} finally { } finally {
DECORATOR.beforeFinish(scope) DECORATOR.beforeFinish(span)
scope.close() scope.close()
} }
} }