adjust tests to use addClientAttributeExtractor and addServerAttribut…

…eExtractor
open-telemetry · Nov 14, 2022 · 27ae3bb · 27ae3bb
1 parent 76b337b
commit 27ae3bb
Showing 1 changed file with 30 additions and 0 deletions.
diff --git a/...n/grpc-1.6/library/src/test/java/io/opentelemetry/instrumentation/grpc/v1_6/GrpcTest.java b/...n/grpc-1.6/library/src/test/java/io/opentelemetry/instrumentation/grpc/v1_6/GrpcTest.java
@@ -35,6 +35,7 @@ class GrpcTest extends AbstractGrpcTest {
   static final InstrumentationExtension testing = LibraryInstrumentationExtension.create();
 
   private static final AttributeKey<String> CUSTOM_KEY = AttributeKey.stringKey("customKey");
+  private static final AttributeKey<String> CUSTOM_KEY2 = AttributeKey.stringKey("customKey2");
 
   private static final Metadata.Key<String> CUSTOM_METADATA_KEY =
       Metadata.Key.of("customMetadataKey", Metadata.ASCII_STRING_MARSHALLER);
@@ -76,6 +77,7 @@ public void sayHello(
             .intercept(
                 GrpcTelemetry.builder(testing.getOpenTelemetry())
                     .addAttributeExtractor(new CustomAttributesExtractor())
+                    .addServerAttributeExtractor(new CustomAttributesExtractorV2("serverSideValue"))
                     .build()
                     .newServerInterceptor())
             .build()
@@ -87,6 +89,7 @@ public void sayHello(
                 .intercept(
                     GrpcTelemetry.builder(testing.getOpenTelemetry())
                         .addAttributeExtractor(new CustomAttributesExtractor())
+                        .addClientAttributeExtractor(new CustomAttributesExtractorV2("clientSideValue"))
                         .build()
                         .newClientInterceptor()));
 
@@ -117,11 +120,13 @@ public void sayHello(
                         span.hasName("example.Greeter/SayHello")
                             .hasKind(SpanKind.CLIENT)
                             .hasParent(trace.getSpan(0))
+                            .hasAttribute(CUSTOM_KEY2, "clientSideValue")
                             .hasAttribute(CUSTOM_KEY, "customValue"),
                     span ->
                         span.hasName("example.Greeter/SayHello")
                             .hasKind(SpanKind.SERVER)
                             .hasParent(trace.getSpan(1))
+                            .hasAttribute(CUSTOM_KEY2, "serverSideValue")
                             .hasAttribute(CUSTOM_KEY, "customValue")));
   }
 
@@ -149,4 +154,29 @@ public void onEnd(
       }
     }
   }
+
+  private static class CustomAttributesExtractorV2
+      implements AttributesExtractor<GrpcRequest, Status> {
+
+    private final String valueOfKey2;
+
+    public CustomAttributesExtractorV2(String valueOfKey2) {
+      this.valueOfKey2 = valueOfKey2;
+    }
+
+    @Override
+    public void onStart(
+        AttributesBuilder attributes, Context parentContext, GrpcRequest grpcRequest) {
+
+      attributes.put(CUSTOM_KEY2, valueOfKey2);
+    }
+
+    @Override
+    public void onEnd(
+        AttributesBuilder attributes,
+        Context context,
+        GrpcRequest grpcRequest,
+        @Nullable Status status,
+        @Nullable Throwable error) {}
+  }
 }