r/aws_budgets_budget: Add PlannedBudgetLimits support

internal/service/budgets/budget.go

@@ -2,9 +2,6 @@ package budgets
 
 import (
 	"fmt"
-	"log"
-	"strings"
-
 	"github.com/aws/aws-sdk-go/aws"
 	"github.com/aws/aws-sdk-go/aws/arn"
 	"github.com/aws/aws-sdk-go/service/budgets"
@@ -17,6 +14,8 @@ import (
 	"github.com/hashicorp/terraform-provider-aws/internal/tfresource"
 	"github.com/hashicorp/terraform-provider-aws/internal/verify"
 	"github.com/shopspring/decimal"
+	"log"
+	"strings"
 )
 
 func ResourceBudget() *schema.Resource {
@@ -143,12 +142,39 @@ func ResourceBudget() *schema.Resource {
 			},
 			"limit_amount": {
 				Type:             schema.TypeString,
-				Required:         true,
+				Optional:         true,
+				Computed:         true,
 				DiffSuppressFunc: suppressEquivalentBudgetLimitAmount,
+				ConflictsWith:    []string{"planned_limit"},
 			},
 			"limit_unit": {
-				Type:     schema.TypeString,
-				Required: true,
+				Type:          schema.TypeString,
+				Optional:      true,
+				Computed:      true,
+				ConflictsWith: []string{"planned_limit"},
+			},
+			"planned_limit": {
+				Type:     schema.TypeSet,
+				Optional: true,
+				Elem: &schema.Resource{
+					Schema: map[string]*schema.Schema{
+						"start_time": {
+							Type:         schema.TypeString,
+							Required:     true,
+							ValidateFunc: ValidTimePeriodTimestamp,
+						},
+						"amount": {
+							Type:             schema.TypeString,
+							Required:         true,
+							DiffSuppressFunc: suppressEquivalentBudgetLimitAmount,
+						},
+						"unit": {
+							Type:     schema.TypeString,
+							Required: true,
+						},
+					},
+				},
+				ConflictsWith: []string{"limit_amount", "limit_unit"},
 			},
 			"name": {
 				Type:          schema.TypeString,
@@ -312,6 +338,7 @@ func resourceBudgetRead(d *schema.ResourceData, meta interface{}) error {
 		d.Set("limit_unit", budget.BudgetLimit.Unit)
 	}
 
+	d.Set("planned_limit", convertPlannedBudgetLimitsToSet(budget.PlannedBudgetLimits))
 	d.Set("name", budget.BudgetName)
 	d.Set("name_prefix", create.NamePrefixFromName(aws.StringValue(budget.BudgetName)))
 
@@ -556,13 +583,58 @@ func convertCostFiltersToStringMap(costFilters map[string][]*string) map[string]
 	return convertedCostFilters
 }
 
+func convertPlannedBudgetLimitsToSet(plannedBudgetLimits map[string]*budgets.Spend) []interface{} {
+	if plannedBudgetLimits == nil {
+		return nil
+	}
+
+	convertedPlannedBudgetLimits := make([]interface{}, len(plannedBudgetLimits))
+	i := 0
+
+	for k, v := range plannedBudgetLimits {
+		if v == nil {
+			return nil
+		}
+
+		startTime, err := TimePeriodSecondsToString(k)
+		if err != nil {
+			return nil
+		}
+
+		convertedPlannedBudgetLimit := make(map[string]string)
+		convertedPlannedBudgetLimit["start_time"] = startTime
+		convertedPlannedBudgetLimit["amount"] = *v.Amount
+		convertedPlannedBudgetLimit["unit"] = *v.Unit
+
+		convertedPlannedBudgetLimits[i] = convertedPlannedBudgetLimit
+		i++
+	}
+
+	return convertedPlannedBudgetLimits
+}
+
 func expandBudgetUnmarshal(d *schema.ResourceData) (*budgets.Budget, error) {
 	budgetName := d.Get("name").(string)
 	budgetType := d.Get("budget_type").(string)
-	budgetLimitAmount := d.Get("limit_amount").(string)
-	budgetLimitUnit := d.Get("limit_unit").(string)
 	budgetTimeUnit := d.Get("time_unit").(string)
 	budgetCostFilters := make(map[string][]*string)
+	var budgetLimit *budgets.Spend
+	var plannedBudgetLimits map[string]*budgets.Spend
+
+	if plannedBudgetLimitsRaw, ok := d.GetOk("planned_limit"); ok {
+		plannedBudgetLimitsRaw := plannedBudgetLimitsRaw.(*schema.Set).List()
+
+		var err error
+		plannedBudgetLimits, err = expandPlannedBudgetLimitsUnmarshal(plannedBudgetLimitsRaw)
+		if err != nil {
+			return nil, err
+		}
+	} else {
+		budgetLimit = &budgets.Spend{
+			Amount: aws.String(d.Get("limit_amount").(string)),
+			Unit:   aws.String(d.Get("limit_unit").(string)),
+		}
+	}
 
 	if costFilter, ok := d.GetOk("cost_filter"); ok {
 		for _, v := range costFilter.(*schema.Set).List() {
@@ -592,16 +664,14 @@ func expandBudgetUnmarshal(d *schema.ResourceData) (*budgets.Budget, error) {
 	}
 
 	budget := &budgets.Budget{
-		BudgetName: aws.String(budgetName),
-		BudgetType: aws.String(budgetType),
-		BudgetLimit: &budgets.Spend{
-			Amount: aws.String(budgetLimitAmount),
-			Unit:   aws.String(budgetLimitUnit),
-		},
+		BudgetName:          aws.String(budgetName),
+		BudgetType:          aws.String(budgetType),
+		PlannedBudgetLimits: plannedBudgetLimits,
 		TimePeriod: &budgets.TimePeriod{
 			End:   budgetTimePeriodEnd,
 			Start: budgetTimePeriodStart,
 		},
+		BudgetLimit: budgetLimit,
 		TimeUnit:    aws.String(budgetTimeUnit),
 		CostFilters: budgetCostFilters,
 	}
@@ -657,6 +727,29 @@ func expandCostTypes(tfMap map[string]interface{}) *budgets.CostTypes {
 	return apiObject
 }
 
+func expandPlannedBudgetLimitsUnmarshal(plannedBudgetLimitsRaw []interface{}) (map[string]*budgets.Spend, error) {
+	plannedBudgetLimits := make(map[string]*budgets.Spend, len(plannedBudgetLimitsRaw))
+
+	for _, plannedBudgetLimit := range plannedBudgetLimitsRaw {
+		plannedBudgetLimit := plannedBudgetLimit.(map[string]interface{})
+
+		key, err := TimePeriodSecondsFromString(plannedBudgetLimit["start_time"].(string))
+		if err != nil {
+			return nil, err
+		}
+
+		amount := plannedBudgetLimit["amount"].(string)
+		unit := plannedBudgetLimit["unit"].(string)
+
+		plannedBudgetLimits[key] = &budgets.Spend{
+			Amount: aws.String(amount),
+			Unit:   aws.String(unit),
+		}
+	}
+
+	return plannedBudgetLimits, nil
+}
+
 func expandBudgetNotificationsUnmarshal(notificationsRaw []interface{}) ([]*budgets.Notification, [][]*budgets.Subscriber) {
 
 	notifications := make([]*budgets.Notification, len(notificationsRaw))

internal/service/budgets/budget_test.go

@@ -343,6 +343,55 @@ func TestAccBudgetsBudget_notifications(t *testing.T) {
 	})
 }
 
+func TestAccBudgetsBudget_plannedLimits(t *testing.T) {
+	var budget budgets.Budget
+	rName := sdkacctest.RandomWithPrefix(acctest.ResourcePrefix)
+	resourceName := "aws_budgets_budget.test"
+	now := time.Now()
+	config1, testCheckFuncs1 := generateStartTimes(resourceName, "100.0", now)
+	config2, testCheckFuncs2 := generateStartTimes(resourceName, "200.0", now)
+
+	resource.ParallelTest(t, resource.TestCase{
+		PreCheck:          func() { acctest.PreCheck(t); acctest.PreCheckPartitionHasService(budgets.EndpointsID, t) },
+		ErrorCheck:        acctest.ErrorCheck(t, budgets.EndpointsID),
+		ProviderFactories: acctest.ProviderFactories,
+		CheckDestroy:      testAccBudgetDestroy,
+		Steps: []resource.TestStep{
+			{
+				Config: testAccBudgetConfig_plannedLimits(rName, config1),
+				Check: resource.ComposeTestCheckFunc(
+					append(
+						testCheckFuncs1,
+						testAccBudgetExists(resourceName, &budget),
+						resource.TestCheckResourceAttr(resourceName, "name", rName),
+						resource.TestCheckResourceAttr(resourceName, "budget_type", "COST"),
+						resource.TestCheckResourceAttr(resourceName, "time_unit", "MONTHLY"),
+						resource.TestCheckResourceAttr(resourceName, "planned_limit.#", "12"),
+					)...,
+				),
+			},
+			{
+				ResourceName:      resourceName,
+				ImportState:       true,
+				ImportStateVerify: true,
+			},
+			{
+				Config: testAccBudgetConfig_plannedLimitsUpdated(rName, config2),
+				Check: resource.ComposeTestCheckFunc(
+					append(
+						testCheckFuncs2,
+						testAccBudgetExists(resourceName, &budget),
+						resource.TestCheckResourceAttr(resourceName, "name", rName),
+						resource.TestCheckResourceAttr(resourceName, "budget_type", "COST"),
+						resource.TestCheckResourceAttr(resourceName, "time_unit", "MONTHLY"),
+						resource.TestCheckResourceAttr(resourceName, "planned_limit.#", "12"),
+					)...,
+				),
+			},
+		},
+	})
+}
+
 func testAccBudgetExists(resourceName string, v *budgets.Budget) resource.TestCheckFunc {
 	return func(s *terraform.State) error {
 		rs, ok := s.RootModule().Resources[resourceName]
@@ -555,3 +604,58 @@ resource "aws_budgets_budget" "test" {
 }
 `, rName, emailAddress1)
 }
+
+func testAccBudgetConfig_plannedLimits(rName, config string) string {
+	return fmt.Sprintf(`
+resource "aws_budgets_budget" "test" {
+  name         = %[1]q
+  budget_type  = "COST"
+  time_unit    = "MONTHLY"
+  %[2]s
+}
+`, rName, config)
+}
+
+func testAccBudgetConfig_plannedLimitsUpdated(rName, config string) string {
+	return fmt.Sprintf(`
+resource "aws_budgets_budget" "test" {
+  name         = %[1]q
+  budget_type  = "COST"
+  time_unit    = "MONTHLY"
+  %[2]s
+}
+`, rName, config)
+}
+
+func generateStartTimes(resourceName, amount string, now time.Time) (string, []resource.TestCheckFunc) {
+	startTimes := make([]time.Time, 12)
+
+	year, month, _ := now.Date()
+	startTimes[0] = time.Date(year, month, 1, 0, 0, 0, 0, time.UTC)
+
+	for i := 1; i < len(startTimes); i++ {
+		startTimes[i] = startTimes[i-1].AddDate(0, 1, 0)
+	}
+
+	configBuilder := strings.Builder{}
+	for i := 0; i < len(startTimes); i++ {
+		configBuilder.WriteString(fmt.Sprintf(`
+planned_limit {
+  start_time = %[1]q
+  amount     = %[2]q
+  unit       = "USD"
+}
+`, tfbudgets.TimePeriodTimestampToString(&startTimes[i]), amount))
+	}
+
+	testCheckFuncs := make([]resource.TestCheckFunc, len(startTimes))
+	for i := 0; i < len(startTimes); i++ {
+		testCheckFuncs[i] = resource.TestCheckTypeSetElemNestedAttrs(resourceName, "planned_limit.*", map[string]string{
+			"start_time": tfbudgets.TimePeriodTimestampToString(&startTimes[i]),
+			"amount":     amount,
+			"unit":       "USD",
+		})
+	}
+
+	return configBuilder.String(), testCheckFuncs
+}

internal/service/budgets/time.go

@@ -2,6 +2,7 @@ package budgets
 
 import (
 	"fmt"
+	"strconv"
 	"time"
 
 	"github.com/aws/aws-sdk-go/aws"
@@ -33,6 +34,31 @@ func TimePeriodTimestampToString(ts *time.Time) string {
 	return aws.TimeValue(ts).Format(timePeriodLayout)
 }
 
+func TimePeriodSecondsFromString(s string) (string, error) {
+	if s == "" {
+		return "", nil
+	}
+
+	ts, err := time.Parse(timePeriodLayout, s)
+
+	if err != nil {
+		return "", err
+	}
+
+	return strconv.FormatInt(aws.Time(ts).Unix(), 10), nil
+}
+
+func TimePeriodSecondsToString(s string) (string, error) {
+	startTime, err := strconv.ParseInt(s, 10, 64)
+	if err != nil {
+		return "", err
+	}
+
+	startTime = startTime * 1000
+
+	return aws.SecondsTimeValue(&startTime).UTC().Format(timePeriodLayout), nil
+}
+
 func ValidTimePeriodTimestamp(v interface{}, k string) (ws []string, errors []error) {
 	_, err := time.Parse(timePeriodLayout, v.(string))
 

internal/service/budgets/time_test.go

@@ -0,0 +1,27 @@
+package budgets
+
+import "testing"
+
+func TestTimePeriodSecondsFromString(t *testing.T) {
+	seconds, err := TimePeriodSecondsFromString("2020-03-01_00:00")
+	if err != nil {
+		t.Errorf("unexpected error: %s", err)
+	}
+
+	want := "1583020800"
+	if seconds != want {
+		t.Errorf("got %s, expected %s", seconds, want)
+	}
+}
+
+func TestTimePeriodSecondsToString(t *testing.T) {
+	ts, err := TimePeriodSecondsToString("1583020800")
+	if err != nil {
+		t.Errorf("unexpected error: %s", err)
+	}
+
+	want := "2020-03-01_00:00"
+	if ts != want {
+		t.Errorf("got %s, expected %s", ts, want)
+	}
+}