open-nomad/scheduler/util_test.go

package scheduler

import (
	"fmt"
	"os"
	"testing"

	"github.com/hashicorp/nomad/nomad/mock"
	"github.com/hashicorp/nomad/nomad/state"
	"github.com/hashicorp/nomad/nomad/structs"
)

func TestMaterializeTaskGroups(t *testing.T) {
	job := mock.Job()
	index := materializeTaskGroups(job)
	if len(index) != 10 {
		t.Fatalf("Bad: %#v", index)
	}

	for i := 0; i < 10; i++ {
		name := fmt.Sprintf("my-job.web[%d]", i)
		tg, ok := index[name]
		if !ok {
			t.Fatalf("bad")
		}
		if tg != job.TaskGroups[0] {
			t.Fatalf("bad")
		}
	}
}

func TestIndexAllocs(t *testing.T) {
	allocs := []*structs.Allocation{
		&structs.Allocation{Name: "foo"},
		&structs.Allocation{Name: "foo"},
		&structs.Allocation{Name: "bar"},
	}
	index := indexAllocs(allocs)
	if len(index) != 2 {
		t.Fatalf("bad: %#v", index)
	}
	if len(index["foo"]) != 2 {
		t.Fatalf("bad: %#v", index)
	}
	if len(index["bar"]) != 1 {
		t.Fatalf("bad: %#v", index)
	}
}

func TestDiffAllocs(t *testing.T) {
	job := mock.Job()
	required := materializeTaskGroups(job)

	// The "old" job has a previous modify index
	oldJob := new(structs.Job)
	*oldJob = *job
	oldJob.ModifyIndex -= 1

	tainted := map[string]bool{
		"dead": true,
		"zip":  false,
	}

	allocs := []*structs.Allocation{
		// Update the 1st
		&structs.Allocation{
			ID:     mock.GenerateUUID(),
			NodeID: "zip",
			Name:   "my-job.web[0]",
			Job:    oldJob,
		},

		// Ignore the 2rd
		&structs.Allocation{
			ID:     mock.GenerateUUID(),
			NodeID: "zip",
			Name:   "my-job.web[1]",
			Job:    job,
		},

		// Evict 11th
		&structs.Allocation{
			ID:     mock.GenerateUUID(),
			NodeID: "zip",
			Name:   "my-job.web[10]",
		},

		// Migrate the 3rd
		&structs.Allocation{
			ID:     mock.GenerateUUID(),
			NodeID: "dead",
			Name:   "my-job.web[2]",
		},
	}
	existing := indexAllocs(allocs)

	place, update, migrate, evict, ignore := diffAllocs(job, tainted, required, existing)

	// We should update the first alloc
	if len(update) != 1 || update[0].ID != allocs[0].ID {
		t.Fatalf("bad: %#v", update)
	}

	// We should ignore the second alloc
	if len(ignore) != 1 || ignore[0].ID != allocs[1].ID {
		t.Fatalf("bad: %#v", ignore)
	}

	// We should evict the 3rd alloc
	if len(evict) != 1 || evict[0].ID != allocs[2].ID {
		t.Fatalf("bad: %#v", evict)
	}

	// We should migrate the 4rd alloc
	if len(migrate) != 1 || migrate[0].ID != allocs[3].ID {
		t.Fatalf("bad: %#v", migrate)
	}

	// We should place 7
	if len(place) != 7 {
		t.Fatalf("bad: %#v", place)
	}
}

func TestAddEvictsToPlan(t *testing.T) {
	allocs := []*structs.Allocation{
		&structs.Allocation{
			ID:     mock.GenerateUUID(),
			NodeID: "zip",
			Name:   "foo",
		},
		&structs.Allocation{
			ID:     mock.GenerateUUID(),
			NodeID: "zip",
			Name:   "foo",
		},
		&structs.Allocation{
			ID:     mock.GenerateUUID(),
			NodeID: "zip",
			Name:   "bar",
		},
	}
	plan := &structs.Plan{
		NodeEvict: make(map[string][]string),
	}
	index := indexAllocs(allocs)

	evict := []allocNameID{
		allocNameID{Name: "foo", ID: allocs[0].ID},
		allocNameID{Name: "bar", ID: allocs[2].ID},
	}
	addEvictsToPlan(plan, evict, index)

	nodeEvict := plan.NodeEvict["zip"]
	if len(nodeEvict) != 2 {
		t.Fatalf("bad: %#v %v", plan, nodeEvict)
	}
	if nodeEvict[0] != allocs[0].ID || nodeEvict[1] != allocs[2].ID {
		t.Fatalf("bad: %v", nodeEvict)
	}
}

func TestReadyNodesInDCs(t *testing.T) {
	state, err := state.NewStateStore(os.Stderr)
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	node1 := mock.Node()
	node2 := mock.Node()
	node2.Datacenter = "dc2"
	node3 := mock.Node()
	node3.Datacenter = "dc2"
	node3.Status = structs.NodeStatusDown

	noErr(t, state.RegisterNode(1000, node1))
	noErr(t, state.RegisterNode(1001, node2))
	noErr(t, state.RegisterNode(1002, node3))

	nodes, err := readyNodesInDCs(state, []string{"dc1", "dc2"})
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	if len(nodes) != 2 {
		t.Fatalf("err: %v")
	}
	if nodes[0].ID == node3.ID || nodes[1].ID == node3.ID {
		t.Fatalf("Bad: %#v", nodes)
	}
}

func TestRetryMax(t *testing.T) {
	calls := 0
	bad := func() (bool, error) {
		calls += 1
		return false, nil
	}
	err := retryMax(3, bad)
	if err == nil {
		t.Fatalf("should fail")
	}
	if calls != 3 {
		t.Fatalf("mis match")
	}

	calls = 0
	good := func() (bool, error) {
		calls += 1
		return true, nil
	}
	err = retryMax(3, good)
	if err != nil {
		t.Fatalf("err: %v")
	}
	if calls != 1 {
		t.Fatalf("mis match")
	}
}

func TestTaintedNodes(t *testing.T) {
	state, err := state.NewStateStore(os.Stderr)
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	node1 := mock.Node()
	node2 := mock.Node()
	node2.Datacenter = "dc2"
	node3 := mock.Node()
	node3.Datacenter = "dc2"
	node3.Status = structs.NodeStatusDown
	noErr(t, state.RegisterNode(1000, node1))
	noErr(t, state.RegisterNode(1001, node2))
	noErr(t, state.RegisterNode(1002, node3))

	allocs := []*structs.Allocation{
		&structs.Allocation{NodeID: node1.ID},
		&structs.Allocation{NodeID: node2.ID},
		&structs.Allocation{NodeID: node3.ID},
		&structs.Allocation{NodeID: "blah"},
	}
	tainted, err := taintedNodes(state, allocs)
	if err != nil {
		t.Fatalf("err: %v", err)
	}

	if len(tainted) != 4 {
		t.Fatalf("err: %v")
	}
	if tainted[node1.ID] || tainted[node2.ID] {
		t.Fatalf("Bad: %v", tainted)
	}
	if !tainted[node3.ID] || !tainted["blah"] {
		t.Fatalf("Bad: %v", tainted)
	}
}