func: update integration handler tests

Author: Juanadelacuesta

policy/handler.go

@@ -172,20 +172,6 @@ func NewPolicyHandler(config HandlerConfig) (*Handler, error) {
 	return h, nil
 }
 
-func (h *Handler) configureHorizontalPolicy() error {
-
-	err := h.loadCheckRunners()
-	if err != nil {
-		return fmt.Errorf("failed to load check handlers: %w", err)
-	}
-
-	h.minCount = h.policy.Min
-	h.maxCount = h.policy.Max
-
-	h.calculateNewCount = h.calculateHorizontalNewCount
-	return nil
-}
-
 func (h *Handler) loadCheckRunners() error {
 	runners := []*checkRunner{}
 
@@ -211,6 +197,7 @@ func (h *Handler) loadCheckRunners() error {
 		runners = append(runners, runner)
 	}
 
+	// Do the update as a single operation to avoid partial updates.
 	h.checkRunners = runners
 	return nil
 }
@@ -397,24 +384,22 @@ func (h *Handler) updateHandler(updatedPolicy *sdk.ScalingPolicy) {
 
 	switch updatedPolicy.Type {
 	case sdk.ScalingPolicyTypeCluster, sdk.ScalingPolicyTypeHorizontal:
-		h.log.Debug("updating check handlers", "old_checks", len(h.policy.Checks),
-			"new_checks", len(updatedPolicy.Checks))
-		err := h.configureHorizontalPolicy()
+		err := h.updateHorizontalPolicy(updatedPolicy)
 		if err != nil {
-			h.errChn <- fmt.Errorf("unable to update horizontal policy: %w", err)
+			h.errChn <- fmt.Errorf("unable to update horizontal policy %w", err)
 			return
 		}
 
 	default:
-		err := h.configureVerticalPolicy()
+		err := h.updateVerticalPolicy(updatedPolicy)
 		if err != nil {
-			h.errChn <- fmt.Errorf("unable to update vertical policy: %w", err)
+			h.errChn <- fmt.Errorf("unable to update vertical policy %w", err)
+			return
 		}
 	}
 
-	h.log.Debug("policy successfully updated")
 	h.policy = updatedPolicy
-
+	h.log.Debug("policy successfully updated")
 }
 
 // applyMutators applies the mutators registered with the handler in order and

policy/handler_oss.go

@@ -8,10 +8,28 @@ package policy
 
 import (
 	"context"
+	"fmt"
 
 	"github.com/hashicorp/nomad-autoscaler/sdk"
 )
 
+func (h *Handler) configureHorizontalPolicy() error {
+	return h.updateHorizontalPolicy(h.policy)
+}
+
+func (h *Handler) updateHorizontalPolicy(up *sdk.ScalingPolicy) error {
+	err := h.loadCheckRunners()
+	if err != nil {
+		return fmt.Errorf("failed to load check handlers: %w", err)
+	}
+
+	h.minCount = up.Min
+	h.maxCount = up.Max
+
+	h.calculateNewCount = h.calculateHorizontalNewCount
+	return nil
+}
+
 func (h *Handler) configureVerticalPolicy() error {
 
 	h.calculateNewCount = func(ctx context.Context, currentCount int64) (sdk.ScalingAction, error) {
@@ -22,3 +40,7 @@ func (h *Handler) configureVerticalPolicy() error {
 	}
 	return nil
 }
+
+func (h *Handler) updateVerticalPolicy(up *sdk.ScalingPolicy) error {
+	return h.configureVerticalPolicy()
+}

policy/handler_test.go

@@ -467,7 +467,7 @@ func TestHandler_Run_ScalingNotNeeded_Integration(t *testing.T) {
 		pm:               mdg,
 	}
 
-	must.NoError(t, handler.loadCheckRunners())
+	must.NoError(t, handler.configureHorizontalPolicy())
 
 	go handler.Run(ctx)
 	time.Sleep(30 * time.Millisecond)