From 8b7f66499fced5522345b3df90b9bdad01115d9b Mon Sep 17 00:00:00 2001
From: Mahmood Ali <mahmood@hashicorp.com>
Date: Fri, 22 Feb 2019 14:35:21 -0500
Subject: [PATCH] address review comments

---
 client/allocrunner/alloc_runner_test.go       | 10 ++++----
 .../taskrunner/task_runner_test.go            | 24 ++++++++-----------
 2 files changed, 16 insertions(+), 18 deletions(-)

diff --git a/client/allocrunner/alloc_runner_test.go b/client/allocrunner/alloc_runner_test.go
index 11e01f718..139786b44 100644
--- a/client/allocrunner/alloc_runner_test.go
+++ b/client/allocrunner/alloc_runner_test.go
@@ -639,7 +639,7 @@ func TestAllocRunner_SimpleRun(t *testing.T) {
 	ar, err := NewAllocRunner(conf)
 	require.NoError(t, err)
 	go ar.Run()
-	defer ar.Destroy()
+	defer destroy(ar)
 
 	// Wait for alloc to be running
 	testutil.WaitForResult(func() (bool, error) {
@@ -662,6 +662,8 @@ func TestAllocRunner_SimpleRun(t *testing.T) {
 
 }
 
+// TestAllocRunner_MoveAllocDir asserts that a rescheduled
+// allocation copies ephemeral disk content from previous alloc run
 func TestAllocRunner_MoveAllocDir(t *testing.T) {
 	t.Parallel()
 
@@ -672,7 +674,7 @@ func TestAllocRunner_MoveAllocDir(t *testing.T) {
 	ar, err := NewAllocRunner(conf)
 	require.NoError(t, err)
 	ar.Run()
-	defer ar.Destroy()
+	defer destroy(ar)
 
 	require.Equal(t, structs.AllocClientStatusComplete, ar.AllocState().ClientStatus)
 
@@ -700,7 +702,7 @@ func TestAllocRunner_MoveAllocDir(t *testing.T) {
 	require.NoError(t, err)
 
 	ar2.Run()
-	defer ar2.Destroy()
+	defer destroy(ar2)
 
 	require.Equal(t, structs.AllocClientStatusComplete, ar2.AllocState().ClientStatus)
 
@@ -752,7 +754,7 @@ func TestAllocRunner_HandlesArtifactFailure(t *testing.T) {
 	ar, err := NewAllocRunner(conf)
 	require.NoError(t, err)
 	go ar.Run()
-	defer ar.Destroy()
+	defer destroy(ar)
 
 	testutil.WaitForResult(func() (bool, error) {
 		state := ar.AllocState()
diff --git a/client/allocrunner/taskrunner/task_runner_test.go b/client/allocrunner/taskrunner/task_runner_test.go
index ca56921e1..a0b92e3a2 100644
--- a/client/allocrunner/taskrunner/task_runner_test.go
+++ b/client/allocrunner/taskrunner/task_runner_test.go
@@ -1296,6 +1296,8 @@ func TestTaskRunner_Template_Artifact(t *testing.T) {
 	require.NoErrorf(t, err, "%v not rendered", f2)
 }
 
+// TestTaskRunner_Template_NewVaultToken asserts that a new vault token is
+// created when rendering template and that it is revoked on alloc completion
 func TestTaskRunner_Template_NewVaultToken(t *testing.T) {
 	t.Parallel()
 
@@ -1321,10 +1323,7 @@ func TestTaskRunner_Template_NewVaultToken(t *testing.T) {
 	// Wait for a Vault token
 	var token string
 	testutil.WaitForResult(func() (bool, error) {
-		tr.vaultTokenLock.Lock()
-		defer tr.vaultTokenLock.Unlock()
-
-		token = tr.vaultToken
+		token = tr.getVaultToken()
 
 		if token == "" {
 			return false, fmt.Errorf("No Vault token")
@@ -1344,10 +1343,7 @@ func TestTaskRunner_Template_NewVaultToken(t *testing.T) {
 
 	var token2 string
 	testutil.WaitForResult(func() (bool, error) {
-		tr.vaultTokenLock.Lock()
-		defer tr.vaultTokenLock.Unlock()
-
-		token2 = tr.vaultToken
+		token2 = tr.getVaultToken()
 
 		if token2 == "" {
 			return false, fmt.Errorf("No Vault token")
@@ -1379,15 +1375,15 @@ func TestTaskRunner_Template_NewVaultToken(t *testing.T) {
 
 }
 
+// TestTaskRunner_VaultManager_Restart asserts that the alloc is restarted when the alloc
+// derived vault token expires, when task is configured with Restart change mode
 func TestTaskRunner_VaultManager_Restart(t *testing.T) {
 	t.Parallel()
 
 	alloc := mock.BatchAlloc()
 	task := alloc.Job.TaskGroups[0].Tasks[0]
-	task.Driver = "mock_driver"
 	task.Config = map[string]interface{}{
-		"exit_code": "0",
-		"run_for":   "10s",
+		"run_for": "10s",
 	}
 	task.Vault = &structs.Vault{
 		Policies:   []string{"default"},
@@ -1446,15 +1442,15 @@ func TestTaskRunner_VaultManager_Restart(t *testing.T) {
 	})
 }
 
+// TestTaskRunner_VaultManager_Signal asserts that the alloc is signalled when the alloc
+// derived vault token expires, when task is configured with signal change mode
 func TestTaskRunner_VaultManager_Signal(t *testing.T) {
 	t.Parallel()
 
 	alloc := mock.BatchAlloc()
 	task := alloc.Job.TaskGroups[0].Tasks[0]
-	task.Driver = "mock_driver"
 	task.Config = map[string]interface{}{
-		"exit_code": "0",
-		"run_for":   "10s",
+		"run_for": "10s",
 	}
 	task.Vault = &structs.Vault{
 		Policies:     []string{"default"},