func: add test for meta info

go.mod

@@ -26,6 +26,7 @@ require (
 	github.com/mitchellh/go-homedir v1.1.0
 	github.com/prometheus/client_golang v1.12.2
 	github.com/prometheus/common v0.32.1
+	github.com/shoenig/test v0.6.3
 	github.com/stretchr/testify v1.8.1
 	google.golang.org/api v0.80.0
 	google.golang.org/grpc v1.46.0

go.sum

@@ -474,6 +474,7 @@ github.com/sean-/seed v0.0.0-20170313163322-e2103e2c3529/go.mod h1:DxrIzT+xaE7yg
 github.com/sergi/go-diff v1.0.0 h1:Kpca3qRNrduNnOQeazBd0ysaKrUJiIuISHxogkT9RPQ=
 github.com/sergi/go-diff v1.0.0/go.mod h1:0CfEIISq7TuYL3j771MWULgwwjU+GofnZX9QAmXWZgo=
 github.com/shoenig/test v0.6.3 h1:GVXWJFk9PiOjN0KoJ7VrJGH6uLPnqxR7/fe3HUPfE0c=
+github.com/shoenig/test v0.6.3/go.mod h1:byHiCGXqrVaflBLAMq/srcZIHynQPQgeyvkvXnjqq0k=
 github.com/sirupsen/logrus v1.2.0/go.mod h1:LxeOpSwHxABJmUn/MG1IvRgCAasNZTLOkJPxbbu5VWo=
 github.com/sirupsen/logrus v1.4.2/go.mod h1:tLMulIdttU9McNUspp0xgXVQah82FyeX6MwdIuYE2rE=
 github.com/sirupsen/logrus v1.6.0/go.mod h1:7uNnSEd1DgxDLC74fIahvMZmmYsHGZGEOFrfsX/uA88=

sdk/helper/scaleutils/cluster.go

@@ -19,6 +19,13 @@ import (
 	"github.com/hashicorp/nomad/api"
 )
 
+type nodeDrainer interface {
+	UpdateDrainOpts(nodeID string, opts *api.DrainOptions,
+		q *api.WriteOptions) (*api.NodeDrainUpdateResponse, error)
+	MonitorDrain(ctx context.Context, nodeID string, index uint64,
+		ignoreSys bool) <-chan *api.MonitorMessage
+}
+
 // ClusterScaleUtils provides common functionality when performing horizontal
 // cluster scaling evaluations and actions.
 type ClusterScaleUtils struct {
@@ -35,6 +42,8 @@ type ClusterScaleUtils struct {
 	// ClusterNodeIDLookupFunc is the callback function used to translate a
 	// Nomad nodes ID to the remote resource ID used by the target platform.
 	ClusterNodeIDLookupFunc ClusterNodeIDLookupFunc
+
+	drainer nodeDrainer
 }
 
 // NewClusterScaleUtils instantiates a new ClusterScaleUtils object for use.
@@ -56,6 +65,7 @@ func NewClusterScaleUtils(cfg *api.Config, log hclog.Logger) (*ClusterScaleUtils
 		log:       log,
 		client:    client,
 		curNodeID: id,
+		drainer:   client.Nodes(),
 	}, nil
 }
 

sdk/helper/scaleutils/node_drain.go

@@ -16,6 +16,13 @@ import (
 	"github.com/hashicorp/nomad/api"
 )
 
+/* type nodeDrainer interface {
+	UpdateDrainOpts(nodeID string, opts api.DrainOptions,
+		q api.WriteOptions) (api.NodeDrainUpdateResponse, error)
+
+	MonitorDrain(ctx context.Context, nodeID string, index uint64, ignoreSys bool) <-chan *api.MonitorMessage
+} */
+
 const (
 	defaultNodeDrainDeadline    = 15 * time.Minute
 	defaultNodeIgnoreSystemJobs = false
@@ -128,28 +135,29 @@ func (c *ClusterScaleUtils) drainNode(ctx context.Context, nodeID string, spec *
 		DrainSpec:    spec,
 		MarkEligible: false,
 		Meta: map[string]string{
-			"autoscaler": "true",
-			"timestamp":  time.Now().String(),
+			"DrainedBy": "Autoscaler",
 		},
 	}
+
 	// Update the drain on the node.
-	resp, err := c.client.Nodes().UpdateDrainOpts(nodeID, opts, nil)
+	resp, err := c.drainer.UpdateDrainOpts(nodeID, opts, nil)
 	if err != nil {
-		return fmt.Errorf("failed to drain node: %v", err)
+		return fmt.Errorf("failed to drain node: %w", err)
 	}
 
 	// Monitor the drain so we output the log messages. An error here indicates
 	// the drain failed to complete successfully.
 	if err := c.monitorNodeDrain(ctx, nodeID, resp.LastIndex, spec.IgnoreSystemJobs); err != nil {
-		return fmt.Errorf("context done while monitoring node drain: %v", err)
+		return fmt.Errorf("context done while monitoring node drain: %w", err)
 	}
 	return nil
 }
 
 // monitorNodeDrain follows the drain of a node, logging the messages we
 // receive to their appropriate level.
 func (c *ClusterScaleUtils) monitorNodeDrain(ctx context.Context, nodeID string, index uint64, ignoreSys bool) error {
-	for msg := range c.client.Nodes().MonitorDrain(ctx, nodeID, index, ignoreSys) {
+	for msg := range c.drainer.MonitorDrain(ctx, nodeID, index, ignoreSys) {
+
 		switch msg.Level {
 		case api.MonitorMsgLevelInfo:
 			c.log.Info("received node drain message", "node_id", nodeID, "msg", msg.Message)

sdk/helper/scaleutils/node_drain_test.go

@@ -4,16 +4,38 @@
 package scaleutils
 
 import (
+	"context"
 	"errors"
 	"testing"
 	"time"
 
+	hclog "github.com/hashicorp/go-hclog"
 	multierror "github.com/hashicorp/go-multierror"
 	errHelper "github.com/hashicorp/nomad-autoscaler/sdk/helper/error"
 	"github.com/hashicorp/nomad/api"
+	"github.com/shoenig/test/must"
 	"github.com/stretchr/testify/assert"
 )
 
+type mockDrainer struct {
+	drain func(nodeID string, opts *api.DrainOptions,
+		q *api.WriteOptions) (*api.NodeDrainUpdateResponse, error)
+	monitor func(ctx context.Context, nodeID string,
+		index uint64, ignoreSys bool) <-chan *api.MonitorMessage
+	monitorCalled bool
+}
+
+func (md *mockDrainer) UpdateDrainOpts(nodeID string, opts *api.DrainOptions,
+	q *api.WriteOptions) (*api.NodeDrainUpdateResponse, error) {
+	return md.drain(nodeID, opts, q)
+}
+
+func (md *mockDrainer) MonitorDrain(ctx context.Context, nodeID string,
+	index uint64, ignoreSys bool) <-chan *api.MonitorMessage {
+	md.monitorCalled = true
+	return md.monitor(ctx, nodeID, index, ignoreSys)
+}
+
 func TestNewClusterScaleUtils_drainSpec(t *testing.T) {
 	testCases := []struct {
 		inputCfg            map[string]string
@@ -91,3 +113,41 @@ func TestNewClusterScaleUtils_drainSpec(t *testing.T) {
 		})
 	}
 }
+
+func Test_DrainNode(t *testing.T) {
+	testNodeID := "nodeID"
+	testLogger := hclog.New(&hclog.LoggerOptions{
+		Level: hclog.LevelFromString("ERROR"),
+	})
+
+	md := &mockDrainer{
+		drain: func(nodeID string, opts *api.DrainOptions,
+			_ *api.WriteOptions) (*api.NodeDrainUpdateResponse, error) {
+
+			must.StrContains(t, opts.Meta["DrainedBy"], "Autoscaler")
+			return &api.NodeDrainUpdateResponse{}, nil
+		},
+		monitor: func(ctx context.Context, nodeID string,
+			index uint64, ignoreSys bool) <-chan *api.MonitorMessage {
+			outCh := make(chan *api.MonitorMessage, 1)
+			go func() {
+				outCh <- &api.MonitorMessage{
+					Level:   api.MonitorMsgLevelNormal,
+					Message: "test message",
+				}
+				close(outCh)
+			}()
+			return outCh
+		},
+	}
+
+	cu := &ClusterScaleUtils{
+		log:     testLogger,
+		drainer: md,
+	}
+
+	ctx := context.Background()
+	err := cu.drainNode(ctx, testNodeID, &api.DrainSpec{})
+	must.NoError(t, err)
+	must.True(t, md.monitorCalled)
+}