From f0ef1b0cee22ec7b43135121607243787007a21f Mon Sep 17 00:00:00 2001
From: Cooper Filby <cooper.filby@gmail.com>
Date: Mon, 27 Jan 2020 14:28:31 -0600
Subject: [PATCH] Add more ctx methods

---
 scan.go                 | 21 +++++++++++++++++++++
 transact_write_items.go | 33 +++++++++++++++++++++++++++++++++
 2 files changed, 54 insertions(+)

diff --git a/scan.go b/scan.go
index 8044ab4..4338030 100644
--- a/scan.go
+++ b/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")
+}
\ No newline at end of file
diff --git a/transact_write_items.go b/transact_write_items.go
index 945c302..679c51c 100644
--- a/transact_write_items.go
+++ b/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")
 }
\ No newline at end of file