Skip to content

Commit 876fefc

Browse files
tabito-haratabito-hara
committed
Add an acceptance test for s3_vectors_configuration
1 parent bffb9f8 commit 876fefc

File tree

2 files changed

+199
-0
lines changed

2 files changed

+199
-0
lines changed

‎internal/service/bedrockagent/bedrockagent_test.go‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,7 @@ func TestAccBedrockAgent_serial(t *testing.T) {
2020
"OpenSearchBasic": testAccKnowledgeBase_OpenSearch_basic,
2121
"OpenSearchUpdate": testAccKnowledgeBase_OpenSearch_update,
2222
"OpenSearchSupplementalDataStorage": testAccKnowledgeBase_OpenSearch_supplementalDataStorage,
23+
"S3Vectors": testAccKnowledgeBase_S3Vectors,
2324
},
2425
"DataSource": {
2526
acctest.CtBasic: testAccDataSource_basic,

‎internal/service/bedrockagent/knowledge_base_test.go‎

Lines changed: 198 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -404,6 +404,69 @@ func testAccKnowledgeBase_OpenSearch_supplementalDataStorage(t *testing.T) {
404404
})
405405
}
406406

407+
func testAccKnowledgeBase_S3Vectors(t *testing.T) {
408+
ctx := acctest.Context(t)
409+
410+
var knowledgebase types.KnowledgeBase
411+
rName := sdkacctest.RandomWithPrefix(acctest.ResourcePrefix)
412+
resourceName := "aws_bedrockagent_knowledge_base.test"
413+
foundationModel := "amazon.titan-embed-text-v2:0"
414+
415+
resource.Test(t, resource.TestCase{
416+
PreCheck: func() {
417+
acctest.PreCheck(ctx, t)
418+
},
419+
ErrorCheck: acctest.ErrorCheck(t, names.BedrockAgentServiceID),
420+
ProtoV5ProviderFactories: acctest.ProtoV5ProviderFactories,
421+
CheckDestroy: testAccCheckKnowledgeBaseDestroy(ctx),
422+
Steps: []resource.TestStep{
423+
{
424+
Config: testAccKnowledgeBaseConfig_S3VectorsByIndexARN(rName, foundationModel),
425+
ConfigPlanChecks: resource.ConfigPlanChecks{
426+
PreApply: []plancheck.PlanCheck{
427+
plancheck.ExpectResourceAction(resourceName, plancheck.ResourceActionCreate),
428+
},
429+
},
430+
Check: resource.ComposeTestCheckFunc(
431+
testAccCheckKnowledgeBaseExists(ctx, resourceName, &knowledgebase),
432+
resource.TestCheckResourceAttrPair(resourceName, names.AttrRoleARN, "aws_iam_role.test", names.AttrARN),
433+
resource.TestCheckResourceAttr(resourceName, "knowledge_base_configuration.#", "1"),
434+
resource.TestCheckResourceAttr(resourceName, "knowledge_base_configuration.0.vector_knowledge_base_configuration.#", "1"),
435+
resource.TestCheckResourceAttr(resourceName, "knowledge_base_configuration.0.type", "VECTOR"),
436+
resource.TestCheckResourceAttr(resourceName, "storage_configuration.#", "1"),
437+
resource.TestCheckResourceAttr(resourceName, "storage_configuration.0.type", "S3_VECTORS"),
438+
resource.TestCheckResourceAttr(resourceName, "storage_configuration.0.s3_vectors_configuration.#", "1"),
439+
resource.TestCheckResourceAttrPair(resourceName, "storage_configuration.0.s3_vectors_configuration.0.index_arn", "aws_s3vectors_index.test", "index_arn"),
440+
),
441+
},
442+
{
443+
ResourceName: resourceName,
444+
ImportState: true,
445+
ImportStateVerify: true,
446+
},
447+
{
448+
Config: testAccKnowledgeBaseConfig_S3VectorsByIndexName(rName, foundationModel),
449+
ConfigPlanChecks: resource.ConfigPlanChecks{
450+
PreApply: []plancheck.PlanCheck{
451+
plancheck.ExpectResourceAction(resourceName, plancheck.ResourceActionDestroyBeforeCreate),
452+
},
453+
},
454+
Check: resource.ComposeTestCheckFunc(
455+
testAccCheckKnowledgeBaseExists(ctx, resourceName, &knowledgebase),
456+
resource.TestCheckResourceAttrPair(resourceName, names.AttrRoleARN, "aws_iam_role.test", names.AttrARN),
457+
resource.TestCheckResourceAttr(resourceName, "knowledge_base_configuration.#", "1"),
458+
resource.TestCheckResourceAttr(resourceName, "knowledge_base_configuration.0.vector_knowledge_base_configuration.#", "1"),
459+
resource.TestCheckResourceAttr(resourceName, "knowledge_base_configuration.0.type", "VECTOR"),
460+
resource.TestCheckResourceAttr(resourceName, "storage_configuration.#", "1"),
461+
resource.TestCheckResourceAttr(resourceName, "storage_configuration.0.type", "S3_VECTORS"),
462+
resource.TestCheckResourceAttr(resourceName, "storage_configuration.0.s3_vectors_configuration.#", "1"),
463+
resource.TestCheckResourceAttrPair(resourceName, "storage_configuration.0.s3_vectors_configuration.0.index_name", "aws_s3vectors_index.test", "index_name"),
464+
),
465+
},
466+
},
467+
})
468+
}
469+
407470
func testAccCheckKnowledgeBaseDestroy(ctx context.Context) resource.TestCheckFunc {
408471
return func(s *terraform.State) error {
409472
conn := acctest.Provider.Meta().(*conns.AWSClient).BedrockAgentClient(ctx)
@@ -946,3 +1009,138 @@ resource "aws_bedrockagent_knowledge_base" "test" {
9461009
}
9471010
`, rName, model))
9481011
}
1012+
1013+
func testAccKnowledgeBaseConfig_S3VectorsBase(rName string) string {
1014+
return fmt.Sprintf(`
1015+
data "aws_region" "current" {}
1016+
data "aws_partition" "current" {}
1017+
1018+
data "aws_iam_policy_document" "assume_role_bedrock" {
1019+
statement {
1020+
effect = "Allow"
1021+
principals {
1022+
type = "Service"
1023+
identifiers = ["bedrock.amazonaws.com"]
1024+
}
1025+
actions = ["sts:AssumeRole"]
1026+
}
1027+
}
1028+
1029+
data "aws_iam_policy_document" "bedrock" {
1030+
statement {
1031+
effect = "Allow"
1032+
actions = ["bedrock:InvokeModel"]
1033+
resources = ["*"]
1034+
}
1035+
statement {
1036+
effect = "Allow"
1037+
actions = ["s3:ListBucket", "s3:GetObject"]
1038+
resources = ["*"]
1039+
}
1040+
statement {
1041+
effect = "Allow"
1042+
actions = [
1043+
"s3vectors:GetIndex",
1044+
"s3vectors:QueryVectors",
1045+
"s3vectors:PutVectors",
1046+
"s3vectors:GetVectors",
1047+
"s3vectors:DeleteVectors"
1048+
]
1049+
resources = ["*"]
1050+
}
1051+
}
1052+
1053+
resource "aws_iam_role" "test" {
1054+
assume_role_policy = data.aws_iam_policy_document.assume_role_bedrock.json
1055+
name = %[1]q
1056+
}
1057+
1058+
resource "aws_iam_role_policy" "test" {
1059+
role = aws_iam_role.test.name
1060+
policy = data.aws_iam_policy_document.bedrock.json
1061+
}
1062+
1063+
resource "aws_s3vectors_vector_bucket" "test" {
1064+
vector_bucket_name = %[1]q
1065+
force_destroy = true
1066+
}
1067+
1068+
resource "aws_s3vectors_index" "test" {
1069+
index_name = %[1]q
1070+
vector_bucket_name = aws_s3vectors_vector_bucket.test.vector_bucket_name
1071+
1072+
data_type = "float32"
1073+
dimension = 256
1074+
distance_metric = "euclidean"
1075+
}
1076+
`, rName)
1077+
}
1078+
1079+
func testAccKnowledgeBaseConfig_S3VectorsByIndexARN(rName, model string) string {
1080+
return acctest.ConfigCompose(
1081+
testAccKnowledgeBaseConfig_S3VectorsBase(rName),
1082+
fmt.Sprintf(`
1083+
resource "aws_bedrockagent_knowledge_base" "test" {
1084+
depends_on = [
1085+
aws_iam_role_policy.test,
1086+
]
1087+
name = %[1]q
1088+
role_arn = aws_iam_role.test.arn
1089+
1090+
knowledge_base_configuration {
1091+
vector_knowledge_base_configuration {
1092+
embedding_model_arn = "arn:${data.aws_partition.current.partition}:bedrock:${data.aws_region.current.region}::foundation-model/%[2]s"
1093+
embedding_model_configuration {
1094+
bedrock_embedding_model_configuration {
1095+
dimensions = 256
1096+
embedding_data_type = "FLOAT32"
1097+
}
1098+
}
1099+
}
1100+
type = "VECTOR"
1101+
}
1102+
1103+
storage_configuration {
1104+
type = "S3_VECTORS"
1105+
s3_vectors_configuration {
1106+
index_arn = aws_s3vectors_index.test.index_arn
1107+
}
1108+
}
1109+
}
1110+
`, rName, model))
1111+
}
1112+
1113+
func testAccKnowledgeBaseConfig_S3VectorsByIndexName(rName, model string) string {
1114+
return acctest.ConfigCompose(
1115+
testAccKnowledgeBaseConfig_S3VectorsBase(rName),
1116+
fmt.Sprintf(`
1117+
resource "aws_bedrockagent_knowledge_base" "test" {
1118+
depends_on = [
1119+
aws_iam_role_policy.test,
1120+
]
1121+
name = %[1]q
1122+
role_arn = aws_iam_role.test.arn
1123+
1124+
knowledge_base_configuration {
1125+
vector_knowledge_base_configuration {
1126+
embedding_model_arn = "arn:${data.aws_partition.current.partition}:bedrock:${data.aws_region.current.region}::foundation-model/%[2]s"
1127+
embedding_model_configuration {
1128+
bedrock_embedding_model_configuration {
1129+
dimensions = 256
1130+
embedding_data_type = "FLOAT32"
1131+
}
1132+
}
1133+
}
1134+
type = "VECTOR"
1135+
}
1136+
1137+
storage_configuration {
1138+
type = "S3_VECTORS"
1139+
s3_vectors_configuration {
1140+
index_name = aws_s3vectors_index.test.index_name
1141+
vector_bucket_arn = aws_s3vectors_vector_bucket.test.vector_bucket_arn
1142+
}
1143+
}
1144+
}
1145+
`, rName, model))
1146+
}

0 commit comments

Comments
 (0)