From aced15ea273228d827bd972932c489c790601b62 Mon Sep 17 00:00:00 2001
From: Jasmine Dahilig <jasmine@dahilig.com>
Date: Thu, 23 Jan 2020 10:52:17 -0800
Subject: [PATCH] partial test for restore functionality

---
 client/allocrunner/alloc_runner_test.go | 65 +++++++++++++++++++++++++
 nomad/mock/mock.go                      | 65 ++++++++++++++++++++++++-
 2 files changed, 129 insertions(+), 1 deletion(-)

diff --git a/client/allocrunner/alloc_runner_test.go b/client/allocrunner/alloc_runner_test.go
index 90e5b5c85..0c49b833d 100644
--- a/client/allocrunner/alloc_runner_test.go
+++ b/client/allocrunner/alloc_runner_test.go
@@ -423,6 +423,8 @@ func TestAllocRunner_TaskLeader_StopRestoredTG(t *testing.T) {
 	// Wait for tasks to be stopped because leader is dead
 	testutil.WaitForResult(func() (bool, error) {
 		alloc := ar2.Alloc()
+		// TODO: this test does not test anything!!! alloc.TaskStates == map[]
+		t.Fatalf("%v", alloc.TaskStates)
 		for task, state := range alloc.TaskStates {
 			if state.State != structs.TaskStateDead {
 				return false, fmt.Errorf("Task %q should be dead: %v", task, state.State)
@@ -444,6 +446,69 @@ func TestAllocRunner_TaskLeader_StopRestoredTG(t *testing.T) {
 	}
 }
 
+func TestAllocRunner_Restore_LifecycleHooks(t *testing.T) {
+	t.Parallel()
+
+	alloc := mock.LifecycleAlloc()
+
+	conf, cleanup := testAllocRunnerConfig(t, alloc)
+	defer cleanup()
+
+	// Use a memory backed statedb
+	conf.StateDB = state.NewMemDB(conf.Logger)
+
+	ar, err := NewAllocRunner(conf)
+	require.NoError(t, err)
+
+	// Mimic client dies while init task running, and client restarts after init task finished
+	ar.tasks["init"].UpdateState(structs.TaskStateDead, structs.NewTaskEvent(structs.TaskTerminated))
+	ar.tasks["side"].UpdateState(structs.TaskStateRunning, structs.NewTaskEvent(structs.TaskStarted))
+
+	// Create a new AllocRunner to test RestoreState and Run
+	ar2, err := NewAllocRunner(conf)
+	require.NoError(t, err)
+	defer destroy(ar2)
+
+	if err := ar2.Restore(); err != nil {
+		t.Fatalf("error restoring state: %v", err)
+	}
+	ar2.Run()
+
+	// We want to see Restore resume execution with correct hook ordering:
+	// i.e. we should see the "web" main task go from pending to running
+	testutil.WaitForResult(func() (bool, error) {
+		alloc := ar2.Alloc()
+
+		// TODO: debug why this alloc has no TaskStates
+		t.Fatalf("%v", alloc.TaskStates)
+		for task, state := range alloc.TaskStates {
+			t.Fatalf("\n\n\n\t\tTASK %q state %v", task, state.State)
+			if state.State != structs.TaskStateDead {
+				return false, fmt.Errorf("Task %q should be dead: %v", task, state.State)
+			}
+		}
+		// TODO: check for these states specifically
+		//	require.Equal(t, structs.TaskStateDead, restoredAlloc.TaskStates["init"])
+		//	require.Equal(t, structs.TaskStateRunning, restoredAlloc.TaskStates["side"])
+		//	require.Equal(t, structs.TaskStateRunning, restoredAlloc.TaskStates["web"])
+
+		return true, nil
+	}, func(err error) {
+		t.Fatalf("err: %v", err)
+	})
+
+	// TODO: debug why this destroy fails
+	// Make sure it GCs properly
+	ar2.Destroy()
+
+	select {
+	case <-ar2.DestroyCh():
+		// exited as expected
+	case <-time.After(10 * time.Second):
+		t.Fatalf("timed out waiting for AR to GC")
+	}
+}
+
 func TestAllocRunner_Update_Semantics(t *testing.T) {
 	t.Parallel()
 	require := require.New(t)
diff --git a/nomad/mock/mock.go b/nomad/mock/mock.go
index cd8394d0c..83559719b 100644
--- a/nomad/mock/mock.go
+++ b/nomad/mock/mock.go
@@ -385,7 +385,7 @@ func LifecycleJob() *structs.Job {
 						Name:   "web",
 						Driver: "exec",
 						Config: map[string]interface{}{
-							"command": "/bin/date",
+							"command": "/bin/sleep 500",
 						},
 						LogConfig: structs.DefaultLogConfig(),
 						Resources: &structs.Resources{
@@ -440,6 +440,69 @@ func LifecycleJob() *structs.Job {
 	job.Canonicalize()
 	return job
 }
+func LifecycleAlloc() *structs.Allocation {
+	alloc := &structs.Allocation{
+		ID:        uuid.Generate(),
+		EvalID:    uuid.Generate(),
+		NodeID:    "12345678-abcd-efab-cdef-123456789abc",
+		Namespace: structs.DefaultNamespace,
+		TaskGroup: "web",
+
+		// TODO Remove once clientv2 gets merged
+		Resources: &structs.Resources{
+			CPU:      500,
+			MemoryMB: 256,
+		},
+		TaskResources: map[string]*structs.Resources{
+			"web": {
+				CPU:      1000,
+				MemoryMB: 256,
+			},
+			"init": {
+				CPU:      1000,
+				MemoryMB: 256,
+			},
+			"side": {
+				CPU:      1000,
+				MemoryMB: 256,
+			},
+		},
+
+		AllocatedResources: &structs.AllocatedResources{
+			Tasks: map[string]*structs.AllocatedTaskResources{
+				"web": {
+					Cpu: structs.AllocatedCpuResources{
+						CpuShares: 1000,
+					},
+					Memory: structs.AllocatedMemoryResources{
+						MemoryMB: 256,
+					},
+				},
+				"init": {
+					Cpu: structs.AllocatedCpuResources{
+						CpuShares: 1000,
+					},
+					Memory: structs.AllocatedMemoryResources{
+						MemoryMB: 256,
+					},
+				},
+				"side": {
+					Cpu: structs.AllocatedCpuResources{
+						CpuShares: 1000,
+					},
+					Memory: structs.AllocatedMemoryResources{
+						MemoryMB: 256,
+					},
+				},
+			},
+		},
+		Job:           LifecycleJob(),
+		DesiredStatus: structs.AllocDesiredStatusRun,
+		ClientStatus:  structs.AllocClientStatusPending,
+	}
+	alloc.JobID = alloc.Job.ID
+	return alloc
+}
 
 func MaxParallelJob() *structs.Job {
 	update := *structs.DefaultUpdateStrategy