diff --git a/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/decorator/OrmClientDecorator.java b/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/decorator/OrmClientDecorator.java index e590058a15..529671ed0b 100644 --- a/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/decorator/OrmClientDecorator.java +++ b/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/decorator/OrmClientDecorator.java @@ -5,9 +5,9 @@ import io.opentracing.Span; public abstract class OrmClientDecorator extends DatabaseClientDecorator { - public abstract String entityName(final ENTITY entity); + public abstract String entityName(final Object entity); - public Span onOperation(final Span span, final ENTITY entity) { + public Span onOperation(final Span span, final Object entity) { assert span != null; if (entity != null) { diff --git a/dd-java-agent/agent-tooling/src/test/groovy/datadog/trace/agent/decorator/OrmClientDecoratorTest.groovy b/dd-java-agent/agent-tooling/src/test/groovy/datadog/trace/agent/decorator/OrmClientDecoratorTest.groovy new file mode 100644 index 0000000000..81367bdfb2 --- /dev/null +++ b/dd-java-agent/agent-tooling/src/test/groovy/datadog/trace/agent/decorator/OrmClientDecoratorTest.groovy @@ -0,0 +1,75 @@ +package datadog.trace.agent.decorator + +import datadog.trace.api.DDTags +import io.opentracing.Span + +class OrmClientDecoratorTest extends DatabaseClientDecoratorTest { + + def span = Mock(Span) + + def "test onOperation #testName"() { + setup: + decorator = newDecorator({ e -> entityName }) + + when: + decorator.onOperation(span, entity) + + then: + if (isSet) { + 1 * span.setTag(DDTags.RESOURCE_NAME, entityName) + } + 0 * _ + + where: + testName | entity | entityName || isSet + "null entity" | null | "name" || false + "null entityName" | "not null" | null || false + "name set" | "not null" | "name" || true + } + + + def newDecorator(name) { + return new OrmClientDecorator() { + @Override + String entityName(Object entity) { + return name.call(entity) + } + + @Override + protected String dbType() { + return "test-db" + } + + @Override + protected String dbUser(Object o) { + return "test-user" + } + + @Override + protected String dbInstance(Object o) { + return "test-user" + } + + @Override + protected String service() { + return "test-service" + } + + @Override + protected String[] instrumentationNames() { + return ["test1"] + } + + @Override + protected String spanType() { + return "test-type" + } + + @Override + protected String component() { + return "test-component" + } + } + } + +}