github.com/diggerhq/digger/libs@v0.0.0-20240604170430-9d61cdf01cc5/locking/aws/dynamo_locking_test.go (about)

     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  }