Repository: hashicorp/nomad
File: scheduler/generic_sched_test.go (Go)

Code before:
	expectAllocNames := []string{}
	for i := 0; i < 10; i++ {
		expectAllocNames = append(expectAllocNames, fmt.Sprintf("%s.web[%d]", job.ID, i))
	}
	must.SliceContainsAll(t, expectAllocNames, allocNames)

	// Ensure different ports were used.
	used := make(map[int]map[string]struct{})
	for _, alloc := range out {
		for _, port := range alloc.AllocatedResources.Shared.Ports {
			nodeMap, ok := used[port.Value]
			if !ok {
				nodeMap = make(map[string]struct{})
				used[port.Value] = nodeMap
			}
			if _, ok := nodeMap[alloc.NodeID]; ok {
				t.Fatalf("Port collision on node %q %v", alloc.NodeID, port.Value)
			}
			nodeMap[alloc.NodeID] = struct{}{}
		}
	}

	h.AssertEvalStatus(t, structs.EvalStatusComplete)
}

func TestServiceSched_JobRegister_EphemeralDisk(t *testing.T) {
	ci.Parallel(t)

	t.Run("sticky ephemeral allocs in same node pool", func(t *testing.T) {

		h := tests.NewHarness(t)

		// Create some nodes
		for range 10 {
			node := mock.Node()
			must.NoError(t, h.State.UpsertNode(structs.MsgTypeTestSetup, h.NextIndex(), node))
		}

		// Create a job
		job := mock.Job()
		job.TaskGroups[0].EphemeralDisk.Sticky = true
		must.NoError(t, h.State.UpsertJob(structs.MsgTypeTestSetup, h.NextIndex(), nil, job))

		// Create a mock evaluation to register the job
		eval := &structs.Evaluation{
			Namespace:   structs.DefaultNamespace,
			ID:          uuid.Generate(),
			Priority:    job.Priority,
			TriggeredBy: structs.EvalTriggerJobRegister,
			JobID:       job.ID,
			Status:      structs.EvalStatusPending,

Review comment (suggestion):
These tests are pretty long... it might be nice to try to pull out some test setup code into a shared function.

Code after:
	expectAllocNames := []string{}
	for i := 0; i < 10; i++ {
		expectAllocNames = append(expectAllocNames, fmt.Sprintf("%s.web[%d]", job.ID, i))
	}
	must.SliceContainsAll(t, expectAllocNames, allocNames)

	// Ensure different ports were used.
	used := make(map[int]map[string]struct{})
	for _, alloc := range out {
		for _, port := range alloc.AllocatedResources.Shared.Ports {
			nodeMap, ok := used[port.Value]
			if !ok {
				nodeMap = make(map[string]struct{})
				used[port.Value] = nodeMap
			}
			if _, ok := nodeMap[alloc.NodeID]; ok {
				t.Fatalf("Port collision on node %q %v", alloc.NodeID, port.Value)
			}
			nodeMap[alloc.NodeID] = struct{}{}
		}
	}

	h.AssertEvalStatus(t, structs.EvalStatusComplete)
}

func TestServiceSched_JobRegister_EphemeralDisk(t *testing.T) {
	ci.Parallel(t)

	createEphemeralJob := func(t *testing.T, h *tests.Harness, sticky, migrate bool) *structs.Job {
		job := mock.Job()
		job.TaskGroups[0].EphemeralDisk.Sticky = sticky
		job.TaskGroups[0].EphemeralDisk.Migrate = migrate
		must.NoError(t, h.State.UpsertJob(structs.MsgTypeTestSetup, h.NextIndex(), nil, job))

		// Create a mock evaluation to register the job
		eval := &structs.Evaluation{
			Namespace:   structs.DefaultNamespace,
			ID:          uuid.Generate(),
			Priority:    job.Priority,
			TriggeredBy: structs.EvalTriggerJobRegister,
			JobID:       job.ID,
			Status:      structs.EvalStatusPending,
		}
		must.NoError(t, h.State.UpsertEvals(structs.MsgTypeTestSetup, h.NextIndex(), []*structs.Evaluation{eval}))

		// Process the evaluation
		must.NoError(t, h.Process(NewServiceScheduler, eval))

		return job
	}