101 lines
2.9 KiB
Go
101 lines
2.9 KiB
Go
package testdynamo
|
|
|
|
import (
|
|
"context"
|
|
"os"
|
|
"testing"
|
|
|
|
"github.com/aws/aws-sdk-go-v2/aws"
|
|
"github.com/aws/aws-sdk-go-v2/config"
|
|
"github.com/aws/aws-sdk-go-v2/credentials"
|
|
"github.com/aws/aws-sdk-go-v2/feature/dynamodb/attributevalue"
|
|
"github.com/aws/aws-sdk-go-v2/service/dynamodb"
|
|
"github.com/aws/aws-sdk-go-v2/service/dynamodb/types"
|
|
"github.com/stretchr/testify/assert"
|
|
)
|
|
|
|
type TestData struct {
|
|
TableName string
|
|
Index []string
|
|
Data []map[string]interface{}
|
|
}
|
|
|
|
func SetupTestTable(t *testing.T, testData []TestData) *dynamodb.Client {
|
|
t.Helper()
|
|
ctx := context.Background()
|
|
|
|
cfg, err := config.LoadDefaultConfig(ctx,
|
|
config.WithRegion("ap-southeast-2"),
|
|
config.WithCredentialsProvider(credentials.NewStaticCredentialsProvider("abc", "123", "")))
|
|
assert.NoError(t, err)
|
|
|
|
testDynamoURL, ok := os.LookupEnv("TEST_DYNAMO_URL")
|
|
if !ok {
|
|
testDynamoURL = "http://localhost:4566"
|
|
}
|
|
|
|
dynamoClient := dynamodb.NewFromConfig(cfg,
|
|
dynamodb.WithEndpointResolver(dynamodb.EndpointResolverFromURL(testDynamoURL)))
|
|
|
|
for _, table := range testData {
|
|
tableInput := &dynamodb.CreateTableInput{
|
|
TableName: aws.String(table.TableName),
|
|
KeySchema: []types.KeySchemaElement{
|
|
{AttributeName: aws.String("pk"), KeyType: types.KeyTypeHash},
|
|
{AttributeName: aws.String("sk"), KeyType: types.KeyTypeRange},
|
|
},
|
|
AttributeDefinitions: []types.AttributeDefinition{
|
|
{AttributeName: aws.String("pk"), AttributeType: types.ScalarAttributeTypeS},
|
|
{AttributeName: aws.String("sk"), AttributeType: types.ScalarAttributeTypeS},
|
|
},
|
|
ProvisionedThroughput: &types.ProvisionedThroughput{
|
|
ReadCapacityUnits: aws.Int64(100),
|
|
WriteCapacityUnits: aws.Int64(100),
|
|
},
|
|
}
|
|
for _, index := range table.Index {
|
|
tableInput.AttributeDefinitions = append(tableInput.AttributeDefinitions, types.AttributeDefinition{
|
|
AttributeName: aws.String(index),
|
|
AttributeType: types.ScalarAttributeTypeS,
|
|
})
|
|
tableInput.GlobalSecondaryIndexes = append(tableInput.GlobalSecondaryIndexes, types.GlobalSecondaryIndex{
|
|
IndexName: aws.String(index + "-index"),
|
|
KeySchema: []types.KeySchemaElement{
|
|
{AttributeName: aws.String(index), KeyType: types.KeyTypeHash},
|
|
},
|
|
Projection: &types.Projection{
|
|
ProjectionType: types.ProjectionTypeAll,
|
|
},
|
|
ProvisionedThroughput: &types.ProvisionedThroughput{
|
|
ReadCapacityUnits: aws.Int64(100),
|
|
WriteCapacityUnits: aws.Int64(100),
|
|
},
|
|
})
|
|
}
|
|
|
|
_, err = dynamoClient.CreateTable(ctx, tableInput)
|
|
assert.NoError(t, err)
|
|
|
|
for _, item := range table.Data {
|
|
m, err := attributevalue.MarshalMap(item)
|
|
assert.NoError(t, err)
|
|
|
|
_, err = dynamoClient.PutItem(ctx, &dynamodb.PutItemInput{
|
|
TableName: aws.String(table.TableName),
|
|
Item: m,
|
|
})
|
|
assert.NoError(t, err)
|
|
}
|
|
}
|
|
|
|
t.Cleanup(func() {
|
|
for _, table := range testData {
|
|
dynamoClient.DeleteTable(ctx, &dynamodb.DeleteTableInput{
|
|
TableName: aws.String(table.TableName),
|
|
})
|
|
}
|
|
})
|
|
|
|
return dynamoClient
|
|
}
|