Skip to content

Commit a9a4bd3

Browse files
authored
add tests for dynamodb and fix condition when checking if table exists (#1402)
1 parent 2dd7c78 commit a9a4bd3

File tree

2 files changed

+115
-18
lines changed

2 files changed

+115
-18
lines changed

cli/pkg/aws/dynamo_locking.go

Lines changed: 24 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -25,17 +25,23 @@ const (
2525
)
2626

2727
type DynamoDbLock struct {
28-
DynamoDb *dynamodb.Client
28+
DynamoDb DynamoDBClient
2929
}
3030

31-
func isResourceNotFoundExceptionError(err error) bool {
32-
if err != nil {
33-
var apiError smithy.APIError
34-
if errors.As(err, &apiError) {
35-
switch apiError.(type) {
36-
case *types.ResourceNotFoundException:
37-
return true
38-
}
31+
type DynamoDBClient interface {
32+
DescribeTable(ctx context.Context, params *dynamodb.DescribeTableInput, optFns ...func(*dynamodb.Options)) (*dynamodb.DescribeTableOutput, error)
33+
CreateTable(ctx context.Context, params *dynamodb.CreateTableInput, optFns ...func(*dynamodb.Options)) (*dynamodb.CreateTableOutput, error)
34+
UpdateItem(ctx context.Context, params *dynamodb.UpdateItemInput, optFns ...func(*dynamodb.Options)) (*dynamodb.UpdateItemOutput, error)
35+
DeleteItem(ctx context.Context, params *dynamodb.DeleteItemInput, optFns ...func(*dynamodb.Options)) (*dynamodb.DeleteItemOutput, error)
36+
GetItem(ctx context.Context, params *dynamodb.GetItemInput, optFns ...func(*dynamodb.Options)) (*dynamodb.GetItemOutput, error)
37+
}
38+
39+
func isTableNotFoundExceptionError(err error) bool {
40+
var apiError smithy.APIError
41+
if errors.As(err, &apiError) {
42+
switch apiError.(type) {
43+
case *types.TableNotFoundException:
44+
return true
3945
}
4046
}
4147
return false
@@ -49,7 +55,7 @@ func (dynamoDbLock *DynamoDbLock) waitUntilTableCreated(ctx context.Context) err
4955
cnt := 0
5056

5157
if err != nil {
52-
if !isResourceNotFoundExceptionError(err) {
58+
if !isTableNotFoundExceptionError(err) {
5359
return err
5460
}
5561
}
@@ -58,7 +64,7 @@ func (dynamoDbLock *DynamoDbLock) waitUntilTableCreated(ctx context.Context) err
5864
time.Sleep(TableCreationInterval)
5965
status, err = dynamoDbLock.DynamoDb.DescribeTable(ctx, input)
6066
if err != nil {
61-
if !isResourceNotFoundExceptionError(err) {
67+
if !isTableNotFoundExceptionError(err) {
6268
return err
6369
}
6470
}
@@ -78,15 +84,14 @@ func (dynamoDbLock *DynamoDbLock) createTableIfNotExists(ctx context.Context) er
7884
_, err := dynamoDbLock.DynamoDb.DescribeTable(ctx, &dynamodb.DescribeTableInput{
7985
TableName: aws.String(TABLE_NAME),
8086
})
81-
82-
if err != nil {
83-
if !isResourceNotFoundExceptionError(err) {
84-
return err
85-
}
87+
if err == nil { // Table exists
88+
return nil
89+
}
90+
if !isTableNotFoundExceptionError(err) {
91+
return err
8692
}
8793

8894
createtbl_input := &dynamodb.CreateTableInput{
89-
9095
AttributeDefinitions: []types.AttributeDefinition{
9196
{
9297
AttributeName: aws.String("PK"),
@@ -214,7 +219,8 @@ func (dynamoDbLock *DynamoDbLock) GetLock(lockId string) (*int, error) {
214219
}
215220

216221
type TransactionLock struct {
217-
TransactionID int `dynamodbav:"transaction_id"`
222+
TransactionID int `dynamodbav:"transaction_id"`
223+
Timeout string `dynamodbav:"timeout"`
218224
}
219225

220226
var t TransactionLock

cli/pkg/aws/dynamo_locking_test.go

Lines changed: 91 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,91 @@
1+
package aws
2+
3+
import (
4+
"context"
5+
"testing"
6+
7+
"github.com/aws/aws-sdk-go-v2/aws"
8+
"github.com/aws/aws-sdk-go-v2/service/dynamodb"
9+
"github.com/aws/aws-sdk-go-v2/service/dynamodb/types"
10+
)
11+
12+
type mockDynamoDbClient struct {
13+
table map[string]map[string]types.AttributeValue
14+
Options dynamodb.Options
15+
MockDescribeTable func(ctx context.Context, params dynamodb.DescribeTableInput, optFns ...func(*dynamodb.Options)) (*dynamodb.DescribeTableOutput, error)
16+
MockUpdateItem func(ctx context.Context, params *dynamodb.UpdateItemInput, optFns ...func(*dynamodb.Options)) (*dynamodb.UpdateItemOutput, error)
17+
MockGetItem func(ctx context.Context, params *dynamodb.GetItemInput, optFns ...func(*dynamodb.Options)) (*dynamodb.GetItemOutput, error)
18+
MockDeleteItem func(ctx context.Context, params *dynamodb.DeleteItemInput, optFns ...func(*dynamodb.Options)) (*dynamodb.DeleteItemOutput, error)
19+
}
20+
21+
func (m *mockDynamoDbClient) DescribeTable(ctx context.Context, params *dynamodb.DescribeTableInput, optFns ...func(*dynamodb.Options)) (*dynamodb.DescribeTableOutput, error) {
22+
if m.table == nil || m.table[aws.ToString(params.TableName)] == nil {
23+
return nil, &types.TableNotFoundException{}
24+
}
25+
if m.table[aws.ToString(params.TableName)] != nil {
26+
return &dynamodb.DescribeTableOutput{Table: &types.TableDescription{TableName: params.TableName}}, nil
27+
}
28+
return nil, nil
29+
}
30+
31+
func (m *mockDynamoDbClient) CreateTable(ctx context.Context, params *dynamodb.CreateTableInput, optFns ...func(*dynamodb.Options)) (*dynamodb.CreateTableOutput, error) {
32+
m.table[aws.ToString(params.TableName)] = make(map[string]types.AttributeValue)
33+
return nil, nil
34+
}
35+
36+
func (m *mockDynamoDbClient) UpdateItem(ctx context.Context, params *dynamodb.UpdateItemInput, optFns ...func(*dynamodb.Options)) (*dynamodb.UpdateItemOutput, error) {
37+
// TODO: Implement this
38+
return &dynamodb.UpdateItemOutput{}, nil
39+
}
40+
41+
func (m *mockDynamoDbClient) GetItem(ctx context.Context, params *dynamodb.GetItemInput, optFns ...func(*dynamodb.Options)) (*dynamodb.GetItemOutput, error) {
42+
return &dynamodb.GetItemOutput{
43+
Item: map[string]types.AttributeValue{
44+
"PK": &types.AttributeValueMemberS{Value: "LOCK"},
45+
"SK": &types.AttributeValueMemberS{Value: "RES#example-resource"},
46+
"transaction_id": &types.AttributeValueMemberN{Value: "123"},
47+
"timeout": &types.AttributeValueMemberS{Value: "2024-04-01T00:00:00Z"},
48+
},
49+
}, nil
50+
}
51+
52+
func (m *mockDynamoDbClient) DeleteItem(ctx context.Context, params *dynamodb.DeleteItemInput, optFns ...func(*dynamodb.Options)) (*dynamodb.DeleteItemOutput, error) {
53+
m.table[aws.ToString(params.TableName)][aws.ToString(&params.Key["SK"].(*types.AttributeValueMemberS).Value)] = nil
54+
return &dynamodb.DeleteItemOutput{}, nil
55+
}
56+
57+
func TestDynamoDbLock_Lock(t *testing.T) {
58+
client := mockDynamoDbClient{table: make(map[string]map[string]types.AttributeValue)}
59+
dynamodbLock := DynamoDbLock{
60+
DynamoDb: &client,
61+
}
62+
dynamodbLock.DynamoDb.CreateTable(context.Background(), &dynamodb.CreateTableInput{TableName: aws.String(TABLE_NAME)})
63+
64+
// Set up the input parameters for the Lock method
65+
transactionId := 123
66+
resource := "example-resource"
67+
68+
locked, err := dynamodbLock.Lock(transactionId, resource)
69+
if err != nil {
70+
t.Fatalf("Error: %v", err)
71+
}
72+
if !locked {
73+
t.Fatalf("Expected true, got %v", locked)
74+
}
75+
}
76+
func TestDynamoDbLock_GetLock(t *testing.T) {
77+
// Create a mock DynamoDB client
78+
client := mockDynamoDbClient{table: make(map[string]map[string]types.AttributeValue)}
79+
dynamodbLock := DynamoDbLock{
80+
DynamoDb: &client,
81+
}
82+
dynamodbLock.DynamoDb.CreateTable(context.Background(), &dynamodb.CreateTableInput{TableName: aws.String(TABLE_NAME)})
83+
84+
id, err := dynamodbLock.GetLock("example-resource")
85+
if err != nil {
86+
t.Fatalf("Error: %v", err)
87+
}
88+
if *id != 123 {
89+
t.Fatalf("Expected 123, got %v", id)
90+
}
91+
}

0 commit comments

Comments
 (0)