reviewfix3

Author: Pawel Peksa

castai/resource_pod_mutation.go

@@ -1017,42 +1017,36 @@ func podMutationToState(mutation *patching_engine.PodMutation, d *schema.Resourc
 		}
 	}
 
-	// Object filter V2
 	if mutation.ObjectFilterV2 != nil {
 		if err := d.Set(FieldPodMutationObjectFilterV2, flattenObjectFilterV2(mutation.ObjectFilterV2)); err != nil {
 			return diag.FromErr(err)
 		}
 	}
 
-	// Labels
 	if mutation.Labels != nil {
 		if err := d.Set(FieldPodMutationLabels, *mutation.Labels); err != nil {
 			return diag.FromErr(err)
 		}
 	}
 
-	// Annotations
 	if mutation.Annotations != nil {
 		if err := d.Set(FieldPodMutationAnnotations, *mutation.Annotations); err != nil {
 			return diag.FromErr(err)
 		}
 	}
 
-	// Node selector
 	if mutation.NodeSelector != nil {
 		if err := d.Set(FieldPodMutationNodeSelector, flattenPodMutationNodeSelector(mutation.NodeSelector)); err != nil {
 			return diag.FromErr(err)
 		}
 	}
 
-	// Tolerations
 	if mutation.Tolerations != nil {
 		if err := d.Set(FieldPodMutationTolerations, flattenTolerations(*mutation.Tolerations)); err != nil {
 			return diag.FromErr(err)
 		}
 	}
 
-	// Affinity
 	if mutation.Affinity != nil {
 		if err := d.Set(FieldPodMutationAffinity, flattenAffinity(mutation.Affinity)); err != nil {
 			return diag.FromErr(err)
@@ -1070,28 +1064,24 @@ func podMutationToState(mutation *patching_engine.PodMutation, d *schema.Resourc
 		}
 	}
 
-	// Spot distribution percentage
 	if mutation.SpotDistributionPercentage != nil {
 		if err := d.Set(FieldPodMutationSpotDistributionPct, int(*mutation.SpotDistributionPercentage)); err != nil {
 			return diag.FromErr(err)
 		}
 	}
 
-	// Node templates
 	if mutation.NodeTemplatesToConsolidate != nil {
 		if err := d.Set(FieldPodMutationNodeTemplates, *mutation.NodeTemplatesToConsolidate); err != nil {
 			return diag.FromErr(err)
 		}
 	}
 
-	// Restart workloads
 	if mutation.RestartMatchingWorkloads != nil {
 		if err := d.Set(FieldPodMutationRestartWorkloads, *mutation.RestartMatchingWorkloads); err != nil {
 			return diag.FromErr(err)
 		}
 	}
 
-	// Patch
 	if mutation.Patch != nil && len(*mutation.Patch) > 0 {
 		patchJSON, err := json.Marshal(*mutation.Patch)
 		if err != nil {
@@ -1102,19 +1092,23 @@ func podMutationToState(mutation *patching_engine.PodMutation, d *schema.Resourc
 		}
 	}
 
-	// Distribution groups
 	if mutation.DistributionGroups != nil {
-		if err := d.Set(FieldPodMutationDistributionGroups, flattenDistributionGroups(*mutation.DistributionGroups)); err != nil {
+		dgs, err := flattenDistributionGroups(*mutation.DistributionGroups)
+		if err != nil {
+			return diag.FromErr(err)
+		}
+
+		if err := d.Set(FieldPodMutationDistributionGroups, dgs); err != nil {
 			return diag.FromErr(err)
 		}
 	}
 
-	// Computed fields
 	if mutation.Source != nil {
 		if err := d.Set(FieldPodMutationSource, string(*mutation.Source)); err != nil {
 			return diag.FromErr(err)
 		}
 	}
+	
 	return nil
 }
 
@@ -1259,7 +1253,7 @@ func flattenAffinity(affinity *patching_engine.Affinity) []map[string]interface{
 	}
 }
 
-func flattenDistributionGroups(groups []patching_engine.DistributionGroup) []map[string]interface{} {
+func flattenDistributionGroups(groups []patching_engine.DistributionGroup) ([]map[string]interface{}, error) {
 	result := make([]map[string]interface{}, 0, len(groups))
 	for _, g := range groups {
 		gMap := map[string]interface{}{
@@ -1292,14 +1286,17 @@ func flattenDistributionGroups(groups []patching_engine.DistributionGroup) []map
 			}
 			if g.Config.Patch != nil && len(*g.Config.Patch) > 0 {
 				patchJSON, err := json.Marshal(*g.Config.Patch)
-				if err == nil {
-					configMap[FieldPodMutationPatch] = string(patchJSON)
+				if err != nil {
+					return nil, fmt.Errorf("failed to marshal patch: %w", err)
 				}
+
+				configMap[FieldPodMutationPatch] = string(patchJSON)
 			}
 			gMap[FieldPodMutationDistributionGroupConfig] = []map[string]interface{}{configMap}
 		}
 
 		result = append(result, gMap)
 	}
-	return result
+
+	return result, nil
 }