From 2cc492c95e06a74b8564146c230d26dc024977f2 Mon Sep 17 00:00:00 2001
From: Michael Schurter <schmichael@hashicorp.com>
Date: Tue, 25 Apr 2017 11:41:03 -0700
Subject: [PATCH] Test pre-0.6 script check upgrade path

---
 client/alloc_runner_test.go | 79 +++++++++++++++++++++++++++++++++++++
 client/task_runner.go       | 10 ++++-
 client/task_runner_test.go  |  2 +-
 3 files changed, 88 insertions(+), 3 deletions(-)

diff --git a/client/alloc_runner_test.go b/client/alloc_runner_test.go
index f1bc828b8..99b03f659 100644
--- a/client/alloc_runner_test.go
+++ b/client/alloc_runner_test.go
@@ -5,6 +5,7 @@ import (
 	"io/ioutil"
 	"os"
 	"path/filepath"
+	"strings"
 	"testing"
 	"text/template"
 	"time"
@@ -476,6 +477,84 @@ func TestAllocRunner_SaveRestoreState_TerminalAlloc(t *testing.T) {
 	})
 }
 
+// TestAllocRunner_SaveRestoreState_Upgrade asserts that pre-0.6 exec tasks are
+// restarted on upgrade.
+func TestAllocRunner_SaveRestoreState_Upgrade(t *testing.T) {
+	alloc := mock.Alloc()
+	task := alloc.Job.TaskGroups[0].Tasks[0]
+	task.Driver = "mock_driver"
+	task.Config = map[string]interface{}{
+		"exit_code": "0",
+		"run_for":   "10s",
+	}
+
+	upd, ar := testAllocRunnerFromAlloc(alloc, false)
+	// Hack in old version to cause an upgrade on RestoreState
+	origConfig := ar.config.Copy()
+	ar.config.Version = "0.5.6"
+	go ar.Run()
+
+	// Snapshot state
+	testutil.WaitForResult(func() (bool, error) {
+		return len(ar.tasks) == 1, nil
+	}, func(err error) {
+		t.Fatalf("task never started: %v", err)
+	})
+
+	err := ar.SaveState()
+	if err != nil {
+		t.Fatalf("err: %v", err)
+	}
+
+	// Create a new alloc runner
+	l2 := prefixedTestLogger("----- ar2:  ")
+	ar2 := NewAllocRunner(l2, origConfig, upd.Update,
+		&structs.Allocation{ID: ar.alloc.ID}, ar.vaultClient,
+		ar.consulClient)
+	err = ar2.RestoreState()
+	if err != nil {
+		t.Fatalf("err: %v", err)
+	}
+	go ar2.Run()
+
+	testutil.WaitForResult(func() (bool, error) {
+		if len(ar2.tasks) != 1 {
+			return false, fmt.Errorf("Incorrect number of tasks")
+		}
+
+		if upd.Count < 3 {
+			return false, nil
+		}
+
+		for _, ev := range ar2.alloc.TaskStates["web"].Events {
+			if strings.HasSuffix(ev.RestartReason, pre06ScriptCheckReason) {
+				return true, nil
+			}
+		}
+		return false, fmt.Errorf("no restart with proper reason found")
+	}, func(err error) {
+		t.Fatalf("err: %v\nAllocs: %#v\nWeb State: %#v", err, upd.Allocs, ar2.alloc.TaskStates["web"])
+	})
+
+	// Destroy and wait
+	ar2.Destroy()
+	start := time.Now()
+
+	testutil.WaitForResult(func() (bool, error) {
+		alloc := ar2.Alloc()
+		if alloc.ClientStatus != structs.AllocClientStatusComplete {
+			return false, fmt.Errorf("Bad client status; got %v; want %v", alloc.ClientStatus, structs.AllocClientStatusComplete)
+		}
+		return true, nil
+	}, func(err error) {
+		t.Fatalf("err: %v %#v %#v", err, upd.Allocs[0], ar.alloc.TaskStates)
+	})
+
+	if time.Since(start) > time.Duration(testutil.TestMultiplier()*5)*time.Second {
+		t.Fatalf("took too long to terminate")
+	}
+}
+
 // Ensure pre-#2132 state files containing the Context struct are properly
 // migrated to the new format.
 //
diff --git a/client/task_runner.go b/client/task_runner.go
index b88bc6a7d..f3715f338 100644
--- a/client/task_runner.go
+++ b/client/task_runner.go
@@ -298,7 +298,7 @@ func (r *TaskRunner) RestoreState() (string, error) {
 		}
 
 		if pre06ScriptCheck(snap.Version, r.task.Driver, r.task.Services) {
-			restartReason = "upgrading pre-0.6 script checks"
+			restartReason = pre06ScriptCheckReason
 		}
 
 		if err := r.registerServices(d, handle); err != nil {
@@ -320,11 +320,16 @@ func (r *TaskRunner) RestoreState() (string, error) {
 	return restartReason, nil
 }
 
+// ver06 is used for checking for pre-0.6 script checks
 var ver06 = version.Must(version.NewVersion("0.6.0dev"))
 
+// pre06ScriptCheckReason is the restart reason given when a pre-0.6 script
+// check is found on an exec/java task.
+const pre06ScriptCheckReason = "upgrading pre-0.6 script checks"
+
 // pre06ScriptCheck returns true if version is prior to 0.6.0dev.
 func pre06ScriptCheck(ver, driver string, services []*structs.Service) bool {
-	if driver != "exec" && driver != "java" {
+	if driver != "exec" && driver != "java" && driver != "mock_driver" {
 		// Only exec and java are affected
 		return false
 	}
@@ -352,6 +357,7 @@ func (r *TaskRunner) SaveState() error {
 	r.persistLock.Lock()
 	defer r.persistLock.Unlock()
 
+	r.logger.Printf("[XXX] task_runner: %q", r.config.Version)
 	snap := taskRunnerState{
 		Task:               r.task,
 		Version:            r.config.Version,
diff --git a/client/task_runner_test.go b/client/task_runner_test.go
index ede8cb164..b373d6db8 100644
--- a/client/task_runner_test.go
+++ b/client/task_runner_test.go
@@ -369,7 +369,7 @@ func TestTaskRunner_SaveRestoreState(t *testing.T) {
 	tr2 := NewTaskRunner(ctx.tr.logger, ctx.tr.config, ctx.upd.Update,
 		ctx.tr.taskDir, ctx.tr.alloc, task2, ctx.tr.vaultClient, ctx.tr.consul)
 	tr2.restartTracker = noRestartsTracker()
-	if err := tr2.RestoreState(); err != nil {
+	if _, err := tr2.RestoreState(); err != nil {
 		t.Fatalf("err: %v", err)
 	}
 	go tr2.Run()