Add more ctx methods

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")
 }