From 30bc6617cb522fd537600485a23538873276d30d Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Cau=C3=AA=20Braz?= <cauealvesbraz@gmail.com>
Date: Sat, 12 Nov 2022 21:48:08 -0300
Subject: [PATCH] feat(ecr): add get and put lifecycle policy functions

---
 modules/aws/ecr.go      | 44 ++++++++++++++++++++++++++++++++++++
 modules/aws/ecr_test.go | 49 +++++++++++++++++++++++++++++++++++++++++
 2 files changed, 93 insertions(+)

diff --git a/modules/aws/ecr.go b/modules/aws/ecr.go
index d002a12e3..cb2f9f8cc 100644
--- a/modules/aws/ecr.go
+++ b/modules/aws/ecr.go
@@ -6,6 +6,7 @@ import (
 	"github.com/aws/aws-sdk-go/aws"
 	"github.com/aws/aws-sdk-go/service/ecr"
 	"github.com/gruntwork-io/go-commons/errors"
+	"github.com/gruntwork-io/terratest/modules/logger"
 	"github.com/gruntwork-io/terratest/modules/testing"
 	"github.com/stretchr/testify/require"
 )
@@ -97,3 +98,46 @@ func NewECRClientE(t testing.TestingT, region string) (*ecr.ECR, error) {
 	}
 	return ecr.New(sess), nil
 }
+
+// GetECRRepoLifecyclePolicy gets the policies for the given ECR repository.
+// This will fail the test and stop execution if there is an error.
+func GetECRRepoLifecyclePolicy(t testing.TestingT, region string, repo *ecr.Repository) string {
+	policy, err := GetECRRepoLifecyclePolicyE(t, region, repo)
+	require.NoError(t, err)
+	return policy
+}
+
+// GetECRRepoLifecyclePolicyE gets the policies for the given ECR repository.
+func GetECRRepoLifecyclePolicyE(t testing.TestingT, region string, repo *ecr.Repository) (string, error) {
+	client := NewECRClient(t, region)
+	resp, err := client.GetLifecyclePolicy(&ecr.GetLifecyclePolicyInput{RepositoryName: repo.RepositoryName})
+	if err != nil {
+		return "", err
+	}
+	return *resp.LifecyclePolicyText, nil
+}
+
+// PutECRRepoLifecyclePolicy puts the given policy for the given ECR repository.
+// This will fail the test and stop execution if there is an error.
+func PutECRRepoLifecyclePolicy(t testing.TestingT, region string, repo *ecr.Repository, policy string) {
+	err := PutECRRepoLifecyclePolicyE(t, region, repo, policy)
+	require.NoError(t, err)
+}
+
+// PutEcrRepoLifecyclePolicy puts the given policy for the given ECR repository.
+func PutECRRepoLifecyclePolicyE(t testing.TestingT, region string, repo *ecr.Repository, policy string) error {
+	logger.Logf(t, "Applying policy for repository %s in %s", *repo.RepositoryName, region)
+
+	client, err := NewECRClientE(t, region)
+	if err != nil {
+		return err
+	}
+
+	input := &ecr.PutLifecyclePolicyInput{
+		RepositoryName:      repo.RepositoryName,
+		LifecyclePolicyText: aws.String(policy),
+	}
+
+	_, err = client.PutLifecyclePolicy(input)
+	return err
+}
diff --git a/modules/aws/ecr_test.go b/modules/aws/ecr_test.go
index b21652192..30a4c342a 100644
--- a/modules/aws/ecr_test.go
+++ b/modules/aws/ecr_test.go
@@ -26,3 +26,52 @@ func TestEcrRepo(t *testing.T) {
 	require.NoError(t, err)
 	assert.Equal(t, ecrRepoName, aws.StringValue(repo2.RepositoryName))
 }
+
+func TestGetEcrRepoLifecyclePolicyError(t *testing.T) {
+	t.Parallel()
+
+	region := GetRandomStableRegion(t, nil, nil)
+	ecrRepoName := fmt.Sprintf("terratest%s", strings.ToLower(random.UniqueId()))
+	repo1, err := CreateECRRepoE(t, region, ecrRepoName)
+	defer DeleteECRRepo(t, region, repo1)
+	require.NoError(t, err)
+
+	assert.Equal(t, ecrRepoName, aws.StringValue(repo1.RepositoryName))
+
+	_, err = GetECRRepoLifecyclePolicyE(t, region, repo1)
+	require.Error(t, err)
+}
+
+func TestCanSetECRRepoLifecyclePolicyWithSingleRule(t *testing.T) {
+	t.Parallel()
+
+	region := GetRandomStableRegion(t, nil, nil)
+	ecrRepoName := fmt.Sprintf("terratest%s", strings.ToLower(random.UniqueId()))
+	repo1, err := CreateECRRepoE(t, region, ecrRepoName)
+	defer DeleteECRRepo(t, region, repo1)
+	require.NoError(t, err)
+
+	lifecyclePolicy := `{
+		"rules": [
+			{
+				"rulePriority": 1,
+				"description": "Expire images older than 14 days",
+				"selection": {
+					"tagStatus": "untagged",
+					"countType": "sinceImagePushed",
+					"countUnit": "days",
+					"countNumber": 14
+				},
+				"action": {
+					"type": "expire"
+				}
+			}
+		]
+	}`
+
+	err = PutECRRepoLifecyclePolicyE(t, region, repo1, lifecyclePolicy)
+	require.NoError(t, err)
+
+	policy := GetECRRepoLifecyclePolicy(t, region, repo1)
+	assert.JSONEq(t, lifecyclePolicy, policy)
+}