Allow callers to set log output

Signed-off-by: Jason Hall <jasonhall@redhat.com>
awslabs · Feb 14, 2022 · 276ba67 · 276ba67
1 parent 69c85dc
commit 276ba67
Show file tree

Hide file tree

Showing 3 changed files with 44 additions and 25 deletions.
diff --git a/ecr-login/cli/docker-credential-ecr-login/main.go b/ecr-login/cli/docker-credential-ecr-login/main.go
@@ -19,7 +19,6 @@ import (
 	"os"
 
 	ecr "github.com/awslabs/amazon-ecr-credential-helper/ecr-login"
-	"github.com/awslabs/amazon-ecr-credential-helper/ecr-login/api"
 	"github.com/awslabs/amazon-ecr-credential-helper/ecr-login/config"
 	"github.com/awslabs/amazon-ecr-credential-helper/ecr-login/version"
 	"github.com/docker/docker-credential-helpers/credentials"
@@ -42,5 +41,5 @@ func main() {
 	}
 
 	config.SetupLogger()
-	credentials.Serve(ecr.ECRHelper{ClientFactory: api.DefaultClientFactory{}})
+	credentials.Serve(ecr.NewECRHelper())
 }
diff --git a/ecr-login/ecr.go b/ecr-login/ecr.go
@@ -16,6 +16,7 @@ package ecr
 import (
 	"errors"
 	"fmt"
+	"io"
 
 	"github.com/sirupsen/logrus"
 
@@ -26,7 +27,34 @@ import (
 var notImplemented = errors.New("not implemented")
 
 type ECRHelper struct {
-	ClientFactory api.ClientFactory
+	clientFactory api.ClientFactory
+	logger        *logrus.Logger
+}
+
+type Option func(*ECRHelper)
+
+func WithClientFactory(clientFactory api.ClientFactory) Option {
+	return func(e *ECRHelper) {
+		e.clientFactory = clientFactory
+	}
+}
+
+func WithLogOutput(w io.Writer) Option {
+	return func(e *ECRHelper) {
+		e.logger.Out = w
+	}
+}
+
+func NewECRHelper(opts ...Option) *ECRHelper {
+	e := &ECRHelper{
+		clientFactory: api.DefaultClientFactory{},
+		logger:        logrus.StandardLogger(),
+	}
+	for _, o := range opts {
+		o(e)
+	}
+
+	return e
 }
 
 // ensure ECRHelper adheres to the credentials.Helper interface
@@ -45,7 +73,7 @@ func (ECRHelper) Delete(serverURL string) error {
 func (self ECRHelper) Get(serverURL string) (string, string, error) {
 	registry, err := api.ExtractRegistry(serverURL)
 	if err != nil {
-		logrus.
+		self.logger.
 			WithError(err).
 			WithField("serverURL", serverURL).
 			Error("Error parsing the serverURL")
@@ -54,30 +82,30 @@ func (self ECRHelper) Get(serverURL string) (string, string, error) {
 
 	var client api.Client
 	if registry.FIPS {
-		client, err = self.ClientFactory.NewClientWithFipsEndpoint(registry.Region)
+		client, err = self.clientFactory.NewClientWithFipsEndpoint(registry.Region)
 		if err != nil {
-			logrus.WithError(err).Error("Error resolving FIPS endpoint")
+			self.logger.WithError(err).Error("Error resolving FIPS endpoint")
 			return "", "", credentials.NewErrCredentialsNotFound()
 		}
 	} else {
-		client = self.ClientFactory.NewClientFromRegion(registry.Region)
+		client = self.clientFactory.NewClientFromRegion(registry.Region)
 	}
 
 	auth, err := client.GetCredentials(serverURL)
 	if err != nil {
-		logrus.WithError(err).Error("Error retrieving credentials")
+		self.logger.WithError(err).Error("Error retrieving credentials")
 		return "", "", credentials.NewErrCredentialsNotFound()
 	}
 	return auth.Username, auth.Password, nil
 }
 
 func (self ECRHelper) List() (map[string]string, error) {
-	logrus.Debug("Listing credentials")
-	client := self.ClientFactory.NewClientWithDefaults()
+	self.logger.Debug("Listing credentials")
+	client := self.clientFactory.NewClientWithDefaults()
 
 	auths, err := client.ListCredentials()
 	if err != nil {
-		logrus.WithError(err).Error("Error listing credentials")
+		self.logger.WithError(err).Error("Error listing credentials")
 		return nil, fmt.Errorf("ecr: could not list credentials: %v", err)
 	}
 

diff --git a/ecr-login/ecr_test.go b/ecr-login/ecr_test.go
@@ -19,7 +19,7 @@ import (
 	"testing"
 
 	ecr "github.com/awslabs/amazon-ecr-credential-helper/ecr-login/api"
-	"github.com/awslabs/amazon-ecr-credential-helper/ecr-login/mocks"
+	mock_api "github.com/awslabs/amazon-ecr-credential-helper/ecr-login/mocks"
 	"github.com/docker/docker-credential-helpers/credentials"
 	"github.com/stretchr/testify/assert"
 )
@@ -36,9 +36,7 @@ func TestGetSuccess(t *testing.T) {
 	factory := &mock_api.MockClientFactory{}
 	client := &mock_api.MockClient{}
 
-	helper := &ECRHelper{
-		ClientFactory: factory,
-	}
+	helper := NewECRHelper(WithClientFactory(factory))
 
 	factory.NewClientFromRegionFn = func(_ string) ecr.Client { return client }
 	client.GetCredentialsFn = func(serverURL string) (*ecr.Auth, error) {
@@ -62,9 +60,7 @@ func TestGetError(t *testing.T) {
 	factory := &mock_api.MockClientFactory{}
 	client := &mock_api.MockClient{}
 
-	helper := &ECRHelper{
-		ClientFactory: factory,
-	}
+	helper := NewECRHelper(WithClientFactory(factory))
 
 	factory.NewClientFromRegionFn = func(_ string) ecr.Client { return client }
 	client.GetCredentialsFn = func(serverURL string) (*ecr.Auth, error) {
@@ -78,7 +74,7 @@ func TestGetError(t *testing.T) {
 }
 
 func TestGetNoMatch(t *testing.T) {
-	helper := &ECRHelper{}
+	helper := NewECRHelper(WithClientFactory(nil))
 
 	username, password, err := helper.Get("not-ecr-server-url")
 	assert.True(t, credentials.IsErrCredentialsNotFound(err))
@@ -90,9 +86,7 @@ func TestListSuccess(t *testing.T) {
 	factory := &mock_api.MockClientFactory{}
 	client := &mock_api.MockClient{}
 
-	helper := &ECRHelper{
-		ClientFactory: factory,
-	}
+	helper := NewECRHelper(WithClientFactory(factory))
 
 	factory.NewClientWithDefaultsFn = func() ecr.Client { return client }
 	client.ListCredentialsFn = func() ([]*ecr.Auth, error) {
@@ -113,9 +107,7 @@ func TestListFailure(t *testing.T) {
 	factory := &mock_api.MockClientFactory{}
 	client := &mock_api.MockClient{}
 
-	helper := &ECRHelper{
-		ClientFactory: factory,
-	}
+	helper := NewECRHelper(WithClientFactory(factory))
 
 	factory.NewClientWithDefaultsFn = func() ecr.Client { return client }
 	client.ListCredentialsFn = func() ([]*ecr.Auth, error) {