From d4cb9f4ac739503948adb3fa7c314e6072d92bbe Mon Sep 17 00:00:00 2001
From: Mariano Cano <mariano@smallstep.com>
Date: Wed, 20 May 2020 12:43:14 -0700
Subject: [PATCH] Define an interface for kms operations.

This interface will be used for unit testing.
---
 kms/awskms/awskms.go | 16 +++++++++++++---
 kms/awskms/signer.go |  4 ++--
 2 files changed, 15 insertions(+), 5 deletions(-)

diff --git a/kms/awskms/awskms.go b/kms/awskms/awskms.go
index aea157b2..d287043d 100644
--- a/kms/awskms/awskms.go
+++ b/kms/awskms/awskms.go
@@ -7,19 +7,29 @@ import (
 	"strings"
 	"time"
 
-	"github.com/smallstep/certificates/kms/uri"
-
+	"github.com/aws/aws-sdk-go/aws"
+	"github.com/aws/aws-sdk-go/aws/request"
 	"github.com/aws/aws-sdk-go/aws/session"
 	"github.com/aws/aws-sdk-go/service/kms"
 	"github.com/pkg/errors"
 	"github.com/smallstep/certificates/kms/apiv1"
+	"github.com/smallstep/certificates/kms/uri"
 	"github.com/smallstep/cli/crypto/pemutil"
 )
 
 // KMS implements a KMS using AWS Key Management Service.
 type KMS struct {
 	session *session.Session
-	service *kms.KMS
+	service KeyManagementClient
+}
+
+// KeyManagementClient defines the methods on KeyManagementClient that this
+// package will use. This interface will be used for unit testing.
+type KeyManagementClient interface {
+	GetPublicKeyWithContext(ctx aws.Context, input *kms.GetPublicKeyInput, opts ...request.Option) (*kms.GetPublicKeyOutput, error)
+	CreateKeyWithContext(ctx aws.Context, input *kms.CreateKeyInput, opts ...request.Option) (*kms.CreateKeyOutput, error)
+	CreateAliasWithContext(ctx aws.Context, input *kms.CreateAliasInput, opts ...request.Option) (*kms.CreateAliasOutput, error)
+	SignWithContext(ctx aws.Context, input *kms.SignInput, opts ...request.Option) (*kms.SignOutput, error)
 }
 
 // customerMasterKeySpecMapping is a mapping between the step signature algorithm,
diff --git a/kms/awskms/signer.go b/kms/awskms/signer.go
index 194eeae4..aa1eb26c 100644
--- a/kms/awskms/signer.go
+++ b/kms/awskms/signer.go
@@ -12,13 +12,13 @@ import (
 )
 
 type Signer struct {
-	service   *kms.KMS
+	service   KeyManagementClient
 	keyID     string
 	publicKey crypto.PublicKey
 }
 
 // NewSigner creates a new signer using a key in the AWS KMS.
-func NewSigner(svc *kms.KMS, signingKey string) (*Signer, error) {
+func NewSigner(svc KeyManagementClient, signingKey string) (*Signer, error) {
 	keyID, err := parseKeyID(signingKey)
 	if err != nil {
 		return nil, err