clean up error handling and don't assume "aws" if we can't find a partition

Author: Michael Kania

cmd/main.go

@@ -453,11 +453,12 @@ func checkExistingAWSProfile(profileName string, config *vault.Config) error {
 	return nil
 }
 
-func getPartition(region string) string {
-	if partition, ok := endpoints.PartitionForRegion(endpoints.DefaultPartitions(), region); ok {
-		return partition.ID()
+func getPartition(region string) (string, error) {
+	partition, ok := endpoints.PartitionForRegion(endpoints.DefaultPartitions(), region)
+	if !ok {
+		return "", fmt.Errorf("Error finding partition for region: %s", region)
 	}
-	return "aws"
+	return partition.ID(), nil
 }
 
 func main() {
@@ -473,10 +474,15 @@ func main() {
 	validate = validator.New()
 
 	// initialize things
+	partition, err := getPartition(options.AwsRegion)
+	if err != nil {
+		log.Fatal(err)
+	}
+
 	profile := vault.Profile{
 		Name: options.AwsProfile,
 		RoleARN: fmt.Sprintf("arn:%s:iam::%d:role/%s",
-			getPartition(options.AwsRegion),
+			partition,
 			options.AwsAccountID,
 			options.Role),
 		Region: options.AwsRegion,

cmd/main_test.go

@@ -63,10 +63,14 @@ func TestGenerateQrCode(t *testing.T) {
 }
 
 func TestGetPartition(t *testing.T) {
-	commPartition := getPartition("us-west-2")
+	commPartition, err := getPartition("us-west-2")
 	assert.Equal(t, commPartition, "aws")
-	govPartition := getPartition("us-gov-west-1")
+	assert.NoError(t, err)
+
+	govPartition, err := getPartition("us-gov-west-1")
 	assert.Equal(t, govPartition, "aws-us-gov")
-	unknownPartition := getPartition("aws-under-the-sea")
-	assert.Equal(t, unknownPartition, "aws")
+	assert.NoError(t, err)
+
+	_, err = getPartition("aws-under-the-sea")
+	assert.Error(t, err)
 }