Add Bedrock contract tests.

aws-observability · Jul 26, 2024 · aa08642 · aa08642
1 parent 395a671
commit aa08642
Show file tree

Hide file tree

Showing 8 changed files with 1,136 additions and 11 deletions.
diff --git a/...c/test/java/software/amazon/opentelemetry/appsignals/test/awssdk/base/AwsSdkBaseTest.java b/...c/test/java/software/amazon/opentelemetry/appsignals/test/awssdk/base/AwsSdkBaseTest.java
diff --git a/...s/src/test/java/software/amazon/opentelemetry/appsignals/test/awssdk/v1/AwsSdkV1Test.java b/...s/src/test/java/software/amazon/opentelemetry/appsignals/test/awssdk/v1/AwsSdkV1Test.java
@@ -57,6 +57,25 @@ protected String getKinesisSpanNamePrefix() {
     return "Kinesis";
   }
 
+  protected String getBedrockSpanNamePrefix() {
+    return "Bedrock";
+  }
+
+  @Override
+  protected String getBedrockAgentSpanNamePrefix() {
+    return "AWSBedrockAgent";
+  }
+
+  @Override
+  protected String getBedrockRuntimeSpanNamePrefix() {
+    return "BedrockRuntime";
+  }
+
+  @Override
+  protected String getBedrockAgentRuntimeSpanNamePrefix() {
+    return "AWSBedrockAgentRuntime";
+  }
+
   protected String getS3RpcServiceName() {
     return "Amazon S3";
   }
@@ -75,6 +94,22 @@ protected String getKinesisRpcServiceName() {
     return "AmazonKinesis";
   }
 
+  protected String getBedrockRpcServiceName() {
+    return "AmazonBedrock";
+  }
+
+  protected String getBedrockAgentRpcServiceName() {
+    return "AWSBedrockAgent";
+  }
+
+  protected String getBedrockRuntimeRpcServiceName() {
+    return "AmazonBedrockRuntime";
+  }
+
+  protected String getBedrockAgentRuntimeRpcServiceName() {
+    return "AWSBedrockAgentRuntime";
+  }
+
   @Test
   void testS3CreateBucket() throws Exception {
     doTestS3CreateBucket();
@@ -165,4 +200,34 @@ void testKinsesisError() throws Exception {
   void testKinesisFault() throws Exception {
     doTestKinesisFault();
   }
+
+  @Test
+  void testBedrockGetKnowlesgeBase() {
+    doTestBedrockGetKnowlesgeBase();
+  }
+
+  @Test
+  void testBedrockGetAgent() {
+    doTestBedrockGetAgent();
+  }
+
+  @Test
+  void testBedrockFault() throws Exception {
+    doTestBedrockFault();
+  }
+
+  @Test
+  void testBedrockRuntimeError() throws Exception {
+    doTestBedrockRuntimeError();
+  }
+
+  @Test
+  void testBedrockGetguardrail() throws Exception {
+    doTestBedrockGetGuardrail();
+  }
+
+  @Test
+  void testBedrockAgentRuntimeGetAgentMemory() throws Exception {
+    doTestBedrockAgentRuntimeGetAgentMemory();
+  }
 }
diff --git a/...s/src/test/java/software/amazon/opentelemetry/appsignals/test/awssdk/v2/AwsSdkV2Test.java b/...s/src/test/java/software/amazon/opentelemetry/appsignals/test/awssdk/v2/AwsSdkV2Test.java
@@ -56,6 +56,25 @@ protected String getKinesisSpanNamePrefix() {
     return "Kinesis";
   }
 
+  protected String getBedrockSpanNamePrefix() {
+    return "Bedrock";
+  }
+
+  @Override
+  protected String getBedrockAgentSpanNamePrefix() {
+    return "BedrockAgent";
+  }
+
+  @Override
+  protected String getBedrockRuntimeSpanNamePrefix() {
+    return "BedrockRuntime";
+  }
+
+  @Override
+  protected String getBedrockAgentRuntimeSpanNamePrefix() {
+    return "BedrockAgentRuntime";
+  }
+
   @Override
   protected String getS3RpcServiceName() {
     return "S3";
@@ -75,6 +94,26 @@ protected String getKinesisRpcServiceName() {
     return "Kinesis";
   }
 
+  @Override
+  protected String getBedrockRpcServiceName() {
+    return "Bedrock";
+  }
+
+  @Override
+  protected String getBedrockAgentRpcServiceName() {
+    return "BedrockAgent";
+  }
+
+  @Override
+  protected String getBedrockRuntimeRpcServiceName() {
+    return "BedrockRuntime";
+  }
+
+  @Override
+  protected String getBedrockAgentRuntimeRpcServiceName() {
+    return "BedrockAgentRuntime";
+  }
+
   @Test
   void testS3CreateBucket() throws Exception {
     doTestS3CreateBucket();
@@ -164,4 +203,34 @@ void testKinesisError() throws Exception {
   void testKinesisFault() throws Exception {
     doTestKinesisFault();
   }
+
+  @Test
+  void testBedrockGetKnowlesgeBase() {
+    doTestBedrockGetKnowlesgeBase();
+  }
+
+  @Test
+  void testBedrockGetAgent() {
+    doTestBedrockGetAgent();
+  }
+
+  @Test
+  void testBedrockFault() throws Exception {
+    doTestBedrockFault();
+  }
+
+  @Test
+  void testBedrockRuntimeError() throws Exception {
+    doTestBedrockRuntimeError();
+  }
+
+  @Test
+  void testBedrockGetguardrail() throws Exception {
+    doTestBedrockGetGuardrail();
+  }
+
+  @Test
+  void testBedrockAgentRuntimeGetAgentMemory() throws Exception {
+    doTestBedrockAgentRuntimeGetAgentMemory();
+  }
 }
diff --git a/...ava/software/amazon/opentelemetry/appsignals/test/utils/SemanticConventionsConstants.java b/...ava/software/amazon/opentelemetry/appsignals/test/utils/SemanticConventionsConstants.java
@@ -58,6 +58,11 @@ public class SemanticConventionsConstants {
   public static final String AWS_QUEUE_URL = "aws.queue.url";
   public static final String AWS_QUEUE_NAME = "aws.queue.name";
   public static final String AWS_STREAM_NAME = "aws.stream.name";
+  public static final String AWS_KNOWLEDGEBASE_ID = "aws.bedrock.knowledge_base.id";
+  public static final String AWS_DATASOURCE_ID = "aws.bedrock.data_source.id";
+  public static final String AWS_AGENT_ID = "aws.bedrock.agent.id";
+  public static final String AWS_GUARDRAIL_ID = "aws.bedrock.guardrail.id";
+  public static final String AWS_BEDROCK_RUNTIME_MODEL_ID = "gen_ai.request.model";
 
   // kafka
   public static final String MESSAGING_CLIENT_ID = "messaging.client_id";

diff --git a/appsignals-tests/images/aws-sdk/aws-sdk-v1/build.gradle.kts b/appsignals-tests/images/aws-sdk/aws-sdk-v1/build.gradle.kts
@@ -27,13 +27,18 @@ dependencies {
   implementation("com.sparkjava:spark-core")
   implementation("com.google.code.gson:gson:2.10.1")
   implementation("org.slf4j:slf4j-simple")
-  implementation(enforcedPlatform("com.amazonaws:aws-java-sdk-bom:1.12.514"))
+  implementation(enforcedPlatform("com.amazonaws:aws-java-sdk-bom:1.12.761"))
   implementation("com.amazonaws:aws-java-sdk-s3")
   implementation("com.amazonaws:aws-java-sdk-dynamodb")
   implementation("com.amazonaws:aws-java-sdk-sqs")
   implementation("com.amazonaws:aws-java-sdk-kinesis")
+  implementation("com.amazonaws:aws-java-sdk-bedrock:1.12.761")
+  implementation("com.amazonaws:aws-java-sdk-bedrockagent:1.12.761")
+  implementation("com.amazonaws:aws-java-sdk-bedrockruntime:1.12.761")
+  implementation("com.amazonaws:aws-java-sdk-bedrockagentruntime:1.12.761")
   implementation("commons-logging:commons-logging")
   implementation("com.linecorp.armeria:armeria")
+  implementation("com.fasterxml.jackson.core:jackson-databind:2.13.3")
 }
 
 java {