Merge pull request #23 from shipt/add-more-ctx-methods

Add more ctx methods

README.md

@@ -136,15 +136,20 @@ DescribeTable(*dynamodb.DescribeTableInput) (*dynamodb.DescribeTableOutput, erro
 GetItem(*dynamodb.GetItemInput) (*dynamodb.GetItemOutput, error)
 GetItemWithContext(aws.Context, *dynamodb.GetItemInput, ...request.Option) (*dynamodb.GetItemOutput, error)
 PutItem(*dynamodb.PutItemInput) (*dynamodb.PutItemOutput, error)
+PutItemWithContext(aws.Context, *dynamodb.PutItemInput, ...request.Option) (*dynamodb.PutItemOutput, error)
 UpdateItem(*dynamodb.UpdateItemInput) (*dynamodb.UpdateItemOutput, error)
 UpdateItemWithContext(aws.Context, *dynamodb.UpdateItemInput, ...request.Option) (*dynamodb.UpdateItemOutput, error)
 DeleteItem(*dynamodb.DeleteItemInput) (*dynamodb.DeleteItemOutput, error)
+DeleteItemWithContext(aws.Context, *dynamodb.DeleteItemInput, ...request.Option) (*dynamodb.DeleteItemOutput, error)
 BatchGetItem(*dynamodb.BatchGetItemInput) (*dynamodb.BatchGetItemOutput, error)
 BatchGetItemWithContext(aws.Context, *dynamodb.BatchGetItemInput, ...request.Option) (*dynamodb.BatchGetItemOutput, error)
 BatchWriteItem(*dynamodb.BatchWriteItemInput) (*dynamodb.BatchWriteItemOutput, error)
+BatchWriteItemWithContext(aws.Context, *dynamodb.BatchWriteItemInput, ...request.Option) (*dynamodb.BatchWriteItemOutput, error)
 WaitUntilTableExists(*dynamodb.DescribeTableInput) error
 Scan(input *dynamodb.ScanInput) (*dynamodb.ScanOutput, error)
+ScanWithContext(aws.Context, *dynamodb.ScanInput, ...request.Option) (*dynamodb.ScanOutput, error)
 Query(input *dynamodb.QueryInput) (*dynamodb.QueryOutput, error)
+QueryWithContext(aws.Context, *dynamodb.QueryInput, request.Option) (*dynamodb.QueryOutput, error)
 ```
 ## Contributions
 

scan.go

@@ -2,6 +2,8 @@ package dynamock
 
 import (
 	"fmt"
+	"github.com/aws/aws-sdk-go/aws"
+	"github.com/aws/aws-sdk-go/aws/request"
 
 	"github.com/aws/aws-sdk-go/service/dynamodb"
 )
@@ -37,3 +39,22 @@ func (e *MockDynamoDB) Scan(input *dynamodb.ScanInput) (*dynamodb.ScanOutput, er
 
 	return nil, fmt.Errorf("Scan Table Expectation Not Found")
 }
+
+func (e *MockDynamoDB) ScanWithContext(ctx aws.Context, input *dynamodb.ScanInput, opts ...request.Option) (*dynamodb.ScanOutput, error) {
+	if len(e.dynaMock.ScanExpect) > 0 {
+		x := e.dynaMock.ScanExpect[0] //get first element of expectation
+
+		if x.table != nil {
+			if *x.table != *input.TableName {
+				return nil, fmt.Errorf("Expect table %s but found table %s", *x.table, *input.TableName)
+			}
+		}
+
+		// delete first element of expectation
+		e.dynaMock.ScanExpect = append(e.dynaMock.ScanExpect[:0], e.dynaMock.ScanExpect[1:]...)
+
+		return x.output, nil
+	}
+
+	return nil, fmt.Errorf("Scan Table With Context Expectation Not Found")
+}

transact_write_items.go

@@ -2,6 +2,8 @@ package dynamock
 
 import (
 	"fmt"
+	"github.com/aws/aws-sdk-go/aws"
+	"github.com/aws/aws-sdk-go/aws/request"
 	"github.com/aws/aws-sdk-go/service/dynamodb"
 )
 
@@ -47,4 +49,35 @@ func (e *MockDynamoDB) TransactWriteItems(input *dynamodb.TransactWriteItemsInpu
 	}
 
 	return nil, fmt.Errorf("Transact Write Items Table Expectation Not Found")
+}
+
+func (e *MockDynamoDB) TransactWriteItemsWithContext(ctx aws.Context, input *dynamodb.TransactWriteItemsInput, opts ...request.Option) (*dynamodb.TransactWriteItemsOutput, error) {
+	if len(e.dynaMock.TransactWriteItemsExpect) > 0 {
+		x := e.dynaMock.TransactWriteItemsExpect[0] //get first element of expectation
+
+		foundTable := false
+
+		if x.table != nil {
+			for _, item := range input.TransactItems {
+				if (item.Update != nil && x.table == item.Update.TableName) ||
+					(item.Put != nil && x.table == item.Put.TableName) ||
+					(item.Delete != nil && x.table == item.Delete.TableName) {
+					foundTable = true
+				}
+			}
+
+			if foundTable == false {
+				return nil, fmt.Errorf("Expect table %s not found", *x.table)
+			}
+		}
+
+		// delete first element of expectation
+
+		e.dynaMock.TransactWriteItemsExpect = append(e.dynaMock.TransactWriteItemsExpect[:0],
+			e.dynaMock.TransactWriteItemsExpect[1:]...)
+
+		return x.output, nil
+	}
+
+	return nil, fmt.Errorf("Transact Write Items Table With Context Expectation Not Found")
 }