Add support for AWS_SESSION_TOKEN in postgres test

Signed-off-by: Prasad Ghangal <prasad.ghangal@gmail.com>

pkg/config/aws/aws.go

@@ -33,6 +33,8 @@ const (
 	SecretAccessKey = "AWS_SECRET_ACCESS_KEY"
 	// SessionToken represents AWS Session Key
 	SessionToken = "AWS_SESSION_TOKEN"
+	// Region represents AWS region
+	Region = "AWS_REGION"
 )
 
 // GetConfig returns a configuration to establish AWS connection, connected region name and the role to assume if it exists.

pkg/testing/db/rds-postgres.go

@@ -22,11 +22,13 @@ import (
 	"time"
 
 	log "github.com/sirupsen/logrus"
+
 	"k8s.io/api/core/v1"
+	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 	"k8s.io/client-go/kubernetes"
 
+	awsconfig "github.com/kanisterio/kanister/pkg/config/aws"
 	"github.com/kanisterio/kanister/pkg/testing/utils"
-	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 
 	// Initialize pq driver
 	_ "github.com/lib/pq"
@@ -42,6 +44,7 @@ type PostgresDB struct {
 	accessID        string
 	secretKey       string
 	region          string
+	sessionToken    string
 	securityGroupID string
 
 	PostgresData Data
@@ -65,17 +68,25 @@ func NewPostgresDB(cli kubernetes.Interface) (Database, error) {
 
 func (pdb *PostgresDB) GetConfig(ctx context.Context) error {
 	var ok bool
-	pdb.accessID, ok = os.LookupEnv("AWS_ACCESS_KEY_ID")
+
+	pdb.region, ok = os.LookupEnv(awsconfig.Region)
 	if !ok {
-		return fmt.Errorf("Env var AWS_ACCESS_KEY_ID is not set")
+		return fmt.Errorf("Env var %s is not set", awsconfig.Region)
 	}
-	pdb.secretKey, ok = os.LookupEnv("AWS_SECRET_ACCESS_KEY")
+
+	// If sessionToken is set, accessID and secretKey not required
+	pdb.sessionToken, ok = os.LookupEnv(awsconfig.SessionToken)
+	if ok {
+		return nil
+	}
+
+	pdb.accessID, ok = os.LookupEnv(awsconfig.AccessKeyID)
 	if !ok {
-		return fmt.Errorf("Env var AWS_SECRET_ACCESS_KEY is not set")
+		return fmt.Errorf("Env var %s is not set", awsconfig.AccessKeyID)
 	}
-	pdb.region, ok = os.LookupEnv("AWS_REGION")
+	pdb.secretKey, ok = os.LookupEnv(awsconfig.SecretAccessKey)
 	if !ok {
-		return fmt.Errorf("Env var AWS_REGION is not set")
+		return fmt.Errorf("Env var %s is not set", awsconfig.SecretAccessKey)
 	}
 	return nil
 }
@@ -94,7 +105,7 @@ func (pdb *PostgresDB) Install(ctx context.Context, nsName string) error {
 	}
 
 	// Create ec2 client
-	ec2, err := utils.NewEC2Client(ctx, pdb.accessID, pdb.secretKey, pdb.region)
+	ec2, err := utils.NewEC2Client(ctx, pdb.accessID, pdb.secretKey, pdb.region, pdb.sessionToken)
 	if err != nil {
 		return err
 	}
@@ -115,7 +126,7 @@ func (pdb *PostgresDB) Install(ctx context.Context, nsName string) error {
 	}
 
 	// Create rds client
-	rds, err := utils.NewRDSClient(ctx, pdb.accessID, pdb.secretKey, pdb.region)
+	rds, err := utils.NewRDSClient(ctx, pdb.accessID, pdb.secretKey, pdb.region, pdb.sessionToken)
 	if err != nil {
 		return err
 	}
@@ -197,7 +208,7 @@ func (pdb PostgresDB) CreateConfig(ctx context.Context, ns string) error {
 
 func (pdb PostgresDB) Remove(ctx context.Context, nsName string) error {
 	// Create rds client
-	rds, err := utils.NewRDSClient(ctx, pdb.accessID, pdb.secretKey, pdb.region)
+	rds, err := utils.NewRDSClient(ctx, pdb.accessID, pdb.secretKey, pdb.region, pdb.sessionToken)
 	if err != nil {
 		log.Errorf("Failed to create rds client: %s. You may need to delete RDS resources manually", err.Error())
 		return err
@@ -218,7 +229,7 @@ func (pdb PostgresDB) Remove(ctx context.Context, nsName string) error {
 	}
 
 	// Create ec2 client
-	ec2, err := utils.NewEC2Client(ctx, pdb.accessID, pdb.secretKey, pdb.region)
+	ec2, err := utils.NewEC2Client(ctx, pdb.accessID, pdb.secretKey, pdb.region, pdb.sessionToken)
 	if err != nil {
 		log.Errorf("Failed to ec2 rds client: %s. You may need to delete EC2 resources manually", err.Error())
 		return err

pkg/testing/integration_test.go

@@ -109,7 +109,7 @@ func (s *IntegrationSuite) TestRun(c *C) {
 		// Check config
 		err := t.database.GetConfig(ctx)
 		if err != nil {
-			log.Infof("Skipping %s. Reason: %s", name, err.Error())
+			log.Infof("Skipping integration test for %s. Reason: %s", name, err.Error())
 			s.tests[name].skip = true
 			continue
 		}

pkg/testing/utils/utils.go

@@ -30,11 +30,12 @@ type RDS struct {
 	Role string
 }
 
-func newAwsConfig(accessID, secretKey, region string) (*aws.Config, *session.Session, string, error) {
+func newAwsConfig(accessID, secretKey, region, sessionToken string) (*aws.Config, *session.Session, string, error) {
 	config := make(map[string]string)
 	config[awsconfig.ConfigRegion] = region
 	config[awsconfig.AccessKeyID] = accessID
 	config[awsconfig.SecretAccessKey] = secretKey
+	config[awsconfig.SessionToken] = sessionToken
 
 	awsConfig, region, role, err := awsconfig.GetConfig(config)
 	if err != nil {
@@ -53,8 +54,8 @@ func newAwsConfig(accessID, secretKey, region string) (*aws.Config, *session.Ses
 }
 
 // NewEC2Client returns ec2 client struct.
-func NewEC2Client(ctx context.Context, accessID, secretKey, region string) (*EC2, error) {
-	conf, s, role, err := newAwsConfig(accessID, secretKey, region)
+func NewEC2Client(ctx context.Context, accessID, secretKey, region, sessionToken string) (*EC2, error) {
+	conf, s, role, err := newAwsConfig(accessID, secretKey, region, sessionToken)
 	if err != nil {
 		return nil, err
 	}
@@ -99,8 +100,8 @@ func (e EC2) DeleteSecurityGroup(ctx context.Context, groupName string) (*ec2.De
 }
 
 // NewRDSClient returns ec2 client struct.
-func NewRDSClient(ctx context.Context, accessID, secretKey, region string) (*RDS, error) {
-	conf, s, role, err := newAwsConfig(accessID, secretKey, region)
+func NewRDSClient(ctx context.Context, accessID, secretKey, region, sessionToken string) (*RDS, error) {
+	conf, s, role, err := newAwsConfig(accessID, secretKey, region, sessionToken)
 	if err != nil {
 		return nil, err
 	}