1 files changed, 9 insertions, 3 deletions
diff --git a/aws/kms.go b/aws/kms.go
index 2961351a..67a5b44b 100644
--- a/aws/kms.go
+++ b/aws/kms.go
@@ -6,12 +6,18 @@ import (
 	"github.com/aws/aws-sdk-go/service/kms"
 )
 
-// KMS -
+// KMSAPI is a subset of kmsiface.KMSAPI
+type KMSAPI interface {
+	Encrypt(input *kms.EncryptInput) (*kms.EncryptOutput, error)
+	Decrypt(input *kms.DecryptInput) (*kms.DecryptOutput, error)
+}
+
+// KMS is an AWS KMS client
 type KMS struct {
-	Client *kms.KMS
+	Client KMSAPI
 }
 
-// NewKMS - Create new KMS client
+// NewKMS - Create new AWS KMS client using an SDKSession
 func NewKMS(option ClientOptions) *KMS {
 	client := kms.New(SDKSession())
 	return &KMS{