feat: Allow enabling predictive scaling on workload optimization scaling policy (#524)

Author: zrudzionis

castai/resource_workload_scaling_policy.go

@@ -36,6 +36,7 @@ const (
 	FieldLimitStrategyType                         = "type"
 	FieldLimitStrategyMultiplier                   = "multiplier"
 	FieldConfidence                                = "confidence"
+	FieldPredictiveScaling                         = "predictive_scaling"
 	FieldConfidenceThreshold                       = "threshold"
 	DeprecatedFieldApplyThreshold                  = "apply_threshold"
 	FieldApplyThresholdStrategy                    = "apply_threshold_strategy"
@@ -257,6 +258,16 @@ It can be either:
 					},
 				},
 			},
+			FieldPredictiveScaling: {
+				Type:     schema.TypeList,
+				Optional: true,
+				MaxItems: 1,
+				Elem: &schema.Resource{
+					Schema: map[string]*schema.Schema{
+						FieldCPU: getPredictiveScalingResourceSchema(),
+					},
+				},
+			},
 		},
 		Timeouts: &schema.ResourceTimeout{
 			Create: schema.DefaultTimeout(15 * time.Second),
@@ -267,6 +278,24 @@ It can be either:
 	}
 }
 
+func getPredictiveScalingResourceSchema() *schema.Schema {
+	return &schema.Schema{
+		Type:        schema.TypeList,
+		Optional:    true,
+		MaxItems:    1,
+		Description: "Defines predictive scaling resource configuration.",
+		Elem: &schema.Resource{
+			Schema: map[string]*schema.Schema{
+				FieldEnabled: {
+					Type:        schema.TypeBool,
+					Required:    true,
+					Description: "Defines if predictive scaling is enabled for resource.",
+				},
+			},
+		},
+	}
+}
+
 func k8sLabelExpressionsSchema() *schema.Schema {
 	return &schema.Schema{
 		Type:        schema.TypeList,
@@ -513,6 +542,8 @@ func resourceWorkloadScalingPolicyCreate(ctx context.Context, d *schema.Resource
 
 	req.RecommendationPolicies.AntiAffinity = toAntiAffinity(toSection(d, "anti_affinity"))
 
+	req.RecommendationPolicies.PredictiveScaling = toPredictiveScaling(toSection(d, FieldPredictiveScaling))
+
 	ar, err := toAssignmentRules(toSection(d, FieldAssignmentRules))
 	if err != nil {
 		return diag.FromErr(err)
@@ -593,6 +624,9 @@ func resourceWorkloadScalingPolicyRead(ctx context.Context, d *schema.ResourceDa
 	if err := d.Set("anti_affinity", toAntiAffinityMap(sp.RecommendationPolicies.AntiAffinity)); err != nil {
 		return diag.FromErr(fmt.Errorf("setting anti-affinity: %w", err))
 	}
+	if err := d.Set(FieldPredictiveScaling, toPredictiveScalingMap(sp.RecommendationPolicies.PredictiveScaling)); err != nil {
+		return diag.FromErr(fmt.Errorf("setting predictive scaling: %w", err))
+	}
 
 	if err := d.Set(FieldAssignmentRules, toAssignmentRulesMap(getResourceFrom(d, FieldAssignmentRules), sp.AssignmentRules)); err != nil {
 		return diag.FromErr(fmt.Errorf("setting assignment rules: %w", err))
@@ -621,6 +655,7 @@ func resourceWorkloadScalingPolicyUpdate(ctx context.Context, d *schema.Resource
 		"anti_affinity",
 		FieldConfidence,
 		FieldAssignmentRules,
+		FieldPredictiveScaling,
 	) {
 		tflog.Info(ctx, "scaling policy up to date")
 		return nil
@@ -646,14 +681,15 @@ func resourceWorkloadScalingPolicyUpdate(ctx context.Context, d *schema.Resource
 		ApplyType:       sdk.WorkloadoptimizationV1ApplyType(d.Get("apply_type").(string)),
 		AssignmentRules: ar,
 		RecommendationPolicies: sdk.WorkloadoptimizationV1RecommendationPolicies{
-			ManagementOption: sdk.WorkloadoptimizationV1ManagementOption(d.Get("management_option").(string)),
-			Cpu:              cpu,
-			Memory:           memory,
-			Startup:          toStartup(toSection(d, "startup")),
-			Downscaling:      toDownscaling(toSection(d, "downscaling")),
-			MemoryEvent:      toMemoryEvent(toSection(d, "memory_event")),
-			AntiAffinity:     toAntiAffinity(toSection(d, "anti_affinity")),
-			Confidence:       toConfidence(toSection(d, FieldConfidence)),
+			ManagementOption:  sdk.WorkloadoptimizationV1ManagementOption(d.Get("management_option").(string)),
+			Cpu:               cpu,
+			Memory:            memory,
+			Startup:           toStartup(toSection(d, "startup")),
+			Downscaling:       toDownscaling(toSection(d, "downscaling")),
+			MemoryEvent:       toMemoryEvent(toSection(d, "memory_event")),
+			AntiAffinity:      toAntiAffinity(toSection(d, "anti_affinity")),
+			Confidence:        toConfidence(toSection(d, FieldConfidence)),
+			PredictiveScaling: toPredictiveScaling(toSection(d, FieldPredictiveScaling)),
 		},
 	}
 
@@ -1161,6 +1197,59 @@ func toAntiAffinityMap(s *sdk.WorkloadoptimizationV1AntiAffinitySettings) []map[
 	return []map[string]any{m}
 }
 
+func toPredictiveScalingMap(s *sdk.WorkloadoptimizationV1PredictiveScalingSettings) []map[string]any {
+	if s == nil || s.Cpu == nil {
+		return nil
+	}
+
+	return []map[string]any{
+		{
+			FieldCPU: toPredictiveScalingResourceMap(s.Cpu),
+		},
+	}
+}
+
+func toPredictiveScaling(m map[string]any) *sdk.WorkloadoptimizationV1PredictiveScalingSettings {
+	if len(m) == 0 {
+		return nil
+	}
+
+	cpuResource := toPredictiveScalingResource(getFirstElem(m, FieldCPU))
+	if cpuResource == nil {
+		return nil
+	}
+
+	return &sdk.WorkloadoptimizationV1PredictiveScalingSettings{
+		Cpu: cpuResource,
+	}
+}
+
+func toPredictiveScalingResourceMap(s *sdk.WorkloadoptimizationV1PredictiveScaling) []map[string]any {
+	if s == nil {
+		return nil
+	}
+
+	return []map[string]any{
+		{
+			FieldEnabled: s.Enabled,
+		},
+	}
+}
+
+func toPredictiveScalingResource(m map[string]any) *sdk.WorkloadoptimizationV1PredictiveScaling {
+	if len(m) == 0 {
+		return nil
+	}
+
+	r := &sdk.WorkloadoptimizationV1PredictiveScaling{}
+
+	if v, ok := m[FieldEnabled].(bool); ok {
+		r.Enabled = v
+	}
+
+	return r
+}
+
 func getWorkloadScalingPolicyByName(ctx context.Context, client sdk.ClientWithResponsesInterface, clusterID, name string) (*sdk.WorkloadoptimizationV1WorkloadScalingPolicy, error) {
 	list, err := client.WorkloadOptimizationAPIListWorkloadScalingPoliciesWithResponse(ctx, clusterID)
 	if checkErr := sdk.CheckOKResponse(list, err); checkErr != nil {

castai/resource_workload_scaling_policy_test.go

@@ -12,6 +12,8 @@ import (
 	"github.com/hashicorp/terraform-plugin-sdk/v2/helper/resource"
 	"github.com/hashicorp/terraform-plugin-sdk/v2/terraform"
 	"github.com/stretchr/testify/require"
+
+	"github.com/castai/terraform-provider-castai/castai/sdk"
 )
 
 func TestAccResourceWorkloadScalingPolicy(t *testing.T) {
@@ -105,6 +107,7 @@ func TestAccResourceWorkloadScalingPolicy(t *testing.T) {
 					resource.TestCheckResourceAttr(resourceName, "assignment_rules.0.rules.1.workload.0.gvk.0", "DaemonSet"),
 					resource.TestCheckResourceAttr(resourceName, "assignment_rules.0.rules.1.workload.0.labels_expressions.0.key", "helm.sh/chart"),
 					resource.TestCheckResourceAttr(resourceName, "assignment_rules.0.rules.1.workload.0.labels_expressions.0.operator", "DoesNotExist"),
+					resource.TestCheckResourceAttr(resourceName, "predictive_scaling.0.cpu.0.enabled", "true"),
 				),
 			},
 		},
@@ -213,6 +216,11 @@ func scalingPolicyConfigUpdated(clusterName, projectID, name string) string {
 				}
 			}
 		}
+		predictive_scaling {
+			cpu {
+				enabled = true
+			}
+		}
 		cpu {
 			function 		= "QUANTILE"
 			overhead 		= 0.15
@@ -392,3 +400,36 @@ func Test_validateResourcePolicy(t *testing.T) {
 		})
 	}
 }
+
+func Test_toPredictiveScaling(t *testing.T) {
+	tests := map[string]struct {
+		args map[string]any
+		exp  *sdk.WorkloadoptimizationV1PredictiveScalingSettings
+	}{
+		"should return predictive scaling settings": {
+			args: map[string]any{
+				FieldCPU: []any{
+					map[string]any{
+						FieldEnabled: true,
+					},
+				},
+			},
+			exp: &sdk.WorkloadoptimizationV1PredictiveScalingSettings{
+				Cpu: &sdk.WorkloadoptimizationV1PredictiveScaling{
+					Enabled: true,
+				},
+			},
+		},
+		"should return nil on empty map": {
+			args: map[string]any{},
+			exp:  nil,
+		},
+	}
+	for name, tt := range tests {
+		t.Run(name, func(t *testing.T) {
+			r := require.New(t)
+			got := toPredictiveScaling(tt.args)
+			r.Equal(tt.exp, got)
+		})
+	}
+}