From 6baf6a1f8f460691888e08222032ff67f7ea2d18 Mon Sep 17 00:00:00 2001
From: Seth Hoenig <seth.a.hoenig@gmail.com>
Date: Wed, 17 Aug 2022 15:22:26 -0500
Subject: [PATCH] cleanup: first pass at fixing command package warnings

This PR is the first of several for cleaning up warnings, and refactoring
bits of code in the command package. First pass is over acl_ files and
gets some helpers in place.
---
 command/acl_bootstrap_test.go             |  67 ++++++-------
 command/acl_policy_apply_test.go          |  31 +++---
 command/acl_policy_delete_test.go         |  19 ++--
 command/acl_policy_info_test.go           |  19 ++--
 command/acl_policy_list_test.go           |  28 +++---
 command/acl_token_create_test.go          |  17 ++--
 command/acl_token_delete_test.go          |  19 ++--
 command/acl_token_info_test.go            |  32 +++---
 command/acl_token_list_test.go            |  29 +++---
 command/acl_token_self_test.go            |  28 ++----
 command/acl_token_update_test.go          |  17 ++--
 command/alloc_status.go                   |   2 +-
 command/alloc_status_test.go              | 114 ++++++----------------
 command/helpers_test.go                   |   2 +-
 command/meta_test.go                      |  10 +-
 command/{util_test.go => testing_test.go} |  40 ++++++--
 command/ui_test.go                        |  20 ++--
 go.mod                                    |   2 +-
 go.sum                                    |   4 +-
 19 files changed, 205 insertions(+), 295 deletions(-)
 rename command/{util_test.go => testing_test.go} (77%)

diff --git a/command/acl_bootstrap_test.go b/command/acl_bootstrap_test.go
index c91588b23..094a00a04 100644
--- a/command/acl_bootstrap_test.go
+++ b/command/acl_bootstrap_test.go
@@ -1,7 +1,6 @@
 package command
 
 import (
-	"io/ioutil"
 	"os"
 	"testing"
 
@@ -9,13 +8,11 @@ import (
 	"github.com/hashicorp/nomad/command/agent"
 	"github.com/hashicorp/nomad/nomad/mock"
 	"github.com/mitchellh/cli"
-	"github.com/stretchr/testify/assert"
-	"github.com/stretchr/testify/require"
+	"github.com/shoenig/test/must"
 )
 
 func TestACLBootstrapCommand(t *testing.T) {
 	ci.Parallel(t)
-	assert := assert.New(t)
 
 	// create a acl-enabled server without bootstrapping the token
 	config := func(c *agent.Config) {
@@ -24,61 +21,59 @@ func TestACLBootstrapCommand(t *testing.T) {
 	}
 
 	srv, _, url := testServer(t, true, config)
-	defer srv.Shutdown()
+	defer stopTestAgent(srv)
 
-	assert.Nil(srv.RootToken)
+	must.Nil(t, srv.RootToken)
 
 	ui := cli.NewMockUi()
 	cmd := &ACLBootstrapCommand{Meta: Meta{Ui: ui, flagAddress: url}}
 
 	code := cmd.Run([]string{"-address=" + url})
-	assert.Equal(0, code)
+	must.Zero(t, code)
 
 	out := ui.OutputWriter.String()
-	assert.Contains(out, "Secret ID")
+	must.StrContains(t, out, "Secret ID")
 }
 
 // If a bootstrap token has already been created, attempts to create more should
 // fail.
 func TestACLBootstrapCommand_ExistingBootstrapToken(t *testing.T) {
 	ci.Parallel(t)
-	assert := assert.New(t)
 
 	config := func(c *agent.Config) {
 		c.ACL.Enabled = true
 	}
 
 	srv, _, url := testServer(t, true, config)
-	defer srv.Shutdown()
+	defer stopTestAgent(srv)
 
-	assert.NotNil(srv.RootToken)
+	must.NotNil(t, srv.RootToken)
 
 	ui := cli.NewMockUi()
 	cmd := &ACLBootstrapCommand{Meta: Meta{Ui: ui, flagAddress: url}}
 
 	code := cmd.Run([]string{"-address=" + url})
-	assert.Equal(1, code)
+	must.One(t, code)
 
 	out := ui.OutputWriter.String()
-	assert.NotContains(out, "Secret ID")
+	must.StrNotContains(t, out, "Secret ID")
 }
 
 // Attempting to bootstrap a token on a non-ACL enabled server should fail.
 func TestACLBootstrapCommand_NonACLServer(t *testing.T) {
 	ci.Parallel(t)
-	assert := assert.New(t)
 
 	srv, _, url := testServer(t, true, nil)
-	defer srv.Shutdown()
+	defer stopTestAgent(srv)
 
 	ui := cli.NewMockUi()
 	cmd := &ACLBootstrapCommand{Meta: Meta{Ui: ui, flagAddress: url}}
 
 	code := cmd.Run([]string{"-address=" + url})
-	assert.Equal(1, code)
+	must.One(t, code)
 
 	out := ui.OutputWriter.String()
-	assert.NotContains(out, "Secret ID")
+	must.StrNotContains(t, out, "Secret ID")
 }
 
 // Attempting to bootstrap the server with an operator provided token in a file should
@@ -95,27 +90,26 @@ func TestACLBootstrapCommand_WithOperatorFileBootstrapToken(t *testing.T) {
 	mockToken := mock.ACLToken()
 
 	// Create temp file
-	f, err := ioutil.TempFile("", "nomad-token.token")
-	assert.Nil(t, err)
-	defer os.Remove(f.Name())
+	file, rm := getTempFile(t, "nomad-token.token")
+	t.Cleanup(rm)
 
 	// Write the token to the file
-	err = ioutil.WriteFile(f.Name(), []byte(mockToken.SecretID), 0700)
-	assert.Nil(t, err)
+	err := os.WriteFile(file, []byte(mockToken.SecretID), 0700)
+	must.NoError(t, err)
 
 	srv, _, url := testServer(t, true, config)
-	defer srv.Shutdown()
+	defer stopTestAgent(srv)
 
-	require.Nil(t, srv.RootToken)
+	must.Nil(t, srv.RootToken)
 
 	ui := cli.NewMockUi()
 	cmd := &ACLBootstrapCommand{Meta: Meta{Ui: ui, flagAddress: url}}
 
-	code := cmd.Run([]string{"-address=" + url, f.Name()})
-	assert.Equal(t, 0, code)
+	code := cmd.Run([]string{"-address=" + url, file})
+	must.Zero(t, code)
 
 	out := ui.OutputWriter.String()
-	assert.Contains(t, out, mockToken.SecretID)
+	must.StrContains(t, out, mockToken.SecretID)
 }
 
 // Attempting to bootstrap the server with an invalid operator provided token in a file should
@@ -133,25 +127,24 @@ func TestACLBootstrapCommand_WithBadOperatorFileBootstrapToken(t *testing.T) {
 	invalidToken := "invalid-token"
 
 	// Create temp file
-	f, err := ioutil.TempFile("", "nomad-token.token")
-	assert.Nil(t, err)
-	defer os.Remove(f.Name())
+	file, cleanup := getTempFile(t, "nomad-token.token")
+	t.Cleanup(cleanup)
 
 	// Write the token to the file
-	err = ioutil.WriteFile(f.Name(), []byte(invalidToken), 0700)
-	assert.Nil(t, err)
+	err := os.WriteFile(file, []byte(invalidToken), 0700)
+	must.NoError(t, err)
 
 	srv, _, url := testServer(t, true, config)
-	defer srv.Shutdown()
+	defer stopTestAgent(srv)
 
-	assert.Nil(t, srv.RootToken)
+	must.Nil(t, srv.RootToken)
 
 	ui := cli.NewMockUi()
 	cmd := &ACLBootstrapCommand{Meta: Meta{Ui: ui, flagAddress: url}}
 
-	code := cmd.Run([]string{"-address=" + url, f.Name()})
-	assert.Equal(t, 1, code)
+	code := cmd.Run([]string{"-address=" + url, file})
+	must.One(t, code)
 
 	out := ui.OutputWriter.String()
-	assert.NotContains(t, out, invalidToken)
+	must.StrNotContains(t, out, invalidToken)
 }
diff --git a/command/acl_policy_apply_test.go b/command/acl_policy_apply_test.go
index 076d0a551..d69fe9d60 100644
--- a/command/acl_policy_apply_test.go
+++ b/command/acl_policy_apply_test.go
@@ -1,31 +1,29 @@
 package command
 
 import (
-	"io/ioutil"
 	"os"
-	"strings"
 	"testing"
 
 	"github.com/hashicorp/nomad/ci"
 	"github.com/hashicorp/nomad/command/agent"
 	"github.com/hashicorp/nomad/nomad/mock"
 	"github.com/mitchellh/cli"
-	"github.com/stretchr/testify/assert"
+	"github.com/shoenig/test/must"
 )
 
 func TestACLPolicyApplyCommand(t *testing.T) {
 	ci.Parallel(t)
-	assert := assert.New(t)
+
 	config := func(c *agent.Config) {
 		c.ACL.Enabled = true
 	}
 
 	srv, _, url := testServer(t, true, config)
-	defer srv.Shutdown()
+	defer stopTestAgent(srv)
 
 	// Bootstrap an initial ACL token
 	token := srv.RootToken
-	assert.NotNil(token, "failed to bootstrap ACL token")
+	must.NotNil(t, token)
 
 	ui := cli.NewMockUi()
 	cmd := &ACLPolicyApplyCommand{Meta: Meta{Ui: ui, flagAddress: url}}
@@ -34,25 +32,22 @@ func TestACLPolicyApplyCommand(t *testing.T) {
 	policy := mock.ACLPolicy()
 
 	// Get a file
-	f, err := ioutil.TempFile("", "nomad-test")
-	assert.Nil(err)
-	defer os.Remove(f.Name())
+	file, rm := getTempFile(t, "nomad-test")
+	t.Cleanup(rm)
 
 	// Write the policy to the file
-	err = ioutil.WriteFile(f.Name(), []byte(policy.Rules), 0700)
-	assert.Nil(err)
+	err := os.WriteFile(file, []byte(policy.Rules), 0700)
+	must.NoError(t, err)
 
 	// Attempt to apply a policy without a valid management token
-	code := cmd.Run([]string{"-address=" + url, "-token=foo", "test-policy", f.Name()})
-	assert.Equal(1, code)
+	code := cmd.Run([]string{"-address=" + url, "-token=foo", "test-policy", file})
+	must.One(t, code)
 
 	// Apply a policy with a valid management token
-	code = cmd.Run([]string{"-address=" + url, "-token=" + token.SecretID, "test-policy", f.Name()})
-	assert.Equal(0, code)
+	code = cmd.Run([]string{"-address=" + url, "-token=" + token.SecretID, "test-policy", file})
+	must.Zero(t, code)
 
 	// Check the output
 	out := ui.OutputWriter.String()
-	if !strings.Contains(out, "Successfully wrote") {
-		t.Fatalf("bad: %v", out)
-	}
+	must.StrContains(t, out, "Successfully wrote")
 }
diff --git a/command/acl_policy_delete_test.go b/command/acl_policy_delete_test.go
index 2ca293827..dc9f915d8 100644
--- a/command/acl_policy_delete_test.go
+++ b/command/acl_policy_delete_test.go
@@ -2,7 +2,6 @@ package command
 
 import (
 	"fmt"
-	"strings"
 	"testing"
 
 	"github.com/hashicorp/nomad/acl"
@@ -11,23 +10,23 @@ import (
 	"github.com/hashicorp/nomad/nomad/mock"
 	"github.com/hashicorp/nomad/nomad/structs"
 	"github.com/mitchellh/cli"
-	"github.com/stretchr/testify/assert"
+	"github.com/shoenig/test/must"
 )
 
 func TestACLPolicyDeleteCommand(t *testing.T) {
 	ci.Parallel(t)
-	assert := assert.New(t)
 	config := func(c *agent.Config) {
 		c.ACL.Enabled = true
 	}
 
 	srv, _, url := testServer(t, true, config)
+	defer stopTestAgent(srv)
+
 	state := srv.Agent.Server().State()
-	defer srv.Shutdown()
 
 	// Bootstrap an initial ACL token
 	token := srv.RootToken
-	assert.NotNil(token, "failed to bootstrap ACL token")
+	must.NotNil(t, token)
 
 	// Create a test ACLPolicy
 	policy := &structs.ACLPolicy{
@@ -35,7 +34,7 @@ func TestACLPolicyDeleteCommand(t *testing.T) {
 		Rules: acl.PolicyWrite,
 	}
 	policy.SetHash()
-	assert.Nil(state.UpsertACLPolicies(structs.MsgTypeTestSetup, 1000, []*structs.ACLPolicy{policy}))
+	must.NoError(t, state.UpsertACLPolicies(structs.MsgTypeTestSetup, 1000, []*structs.ACLPolicy{policy}))
 
 	ui := cli.NewMockUi()
 	cmd := &ACLPolicyDeleteCommand{Meta: Meta{Ui: ui, flagAddress: url}}
@@ -43,15 +42,13 @@ func TestACLPolicyDeleteCommand(t *testing.T) {
 	// Delete the policy without a valid token fails
 	invalidToken := mock.ACLToken()
 	code := cmd.Run([]string{"-address=" + url, "-token=" + invalidToken.SecretID, policy.Name})
-	assert.Equal(1, code)
+	must.One(t, code)
 
 	// Delete the policy with a valid management token
 	code = cmd.Run([]string{"-address=" + url, "-token=" + token.SecretID, policy.Name})
-	assert.Equal(0, code)
+	must.Zero(t, code)
 
 	// Check the output
 	out := ui.OutputWriter.String()
-	if !strings.Contains(out, fmt.Sprintf("Successfully deleted %s policy", policy.Name)) {
-		t.Fatalf("bad: %v", out)
-	}
+	must.StrContains(t, out, fmt.Sprintf("Successfully deleted %s policy", policy.Name))
 }
diff --git a/command/acl_policy_info_test.go b/command/acl_policy_info_test.go
index 828b4022b..8d54dd8a3 100644
--- a/command/acl_policy_info_test.go
+++ b/command/acl_policy_info_test.go
@@ -1,7 +1,6 @@
 package command
 
 import (
-	"strings"
 	"testing"
 
 	"github.com/hashicorp/nomad/ci"
@@ -9,23 +8,23 @@ import (
 	"github.com/hashicorp/nomad/nomad/mock"
 	"github.com/hashicorp/nomad/nomad/structs"
 	"github.com/mitchellh/cli"
-	"github.com/stretchr/testify/assert"
+	"github.com/shoenig/test/must"
 )
 
 func TestACLPolicyInfoCommand(t *testing.T) {
 	ci.Parallel(t)
-	assert := assert.New(t)
+
 	config := func(c *agent.Config) {
 		c.ACL.Enabled = true
 	}
 
 	srv, _, url := testServer(t, true, config)
 	state := srv.Agent.Server().State()
-	defer srv.Shutdown()
+	defer stopTestAgent(srv)
 
 	// Bootstrap an initial ACL token
 	token := srv.RootToken
-	assert.NotNil(token, "failed to bootstrap ACL token")
+	must.NotNil(t, token)
 
 	// Create a test ACLPolicy
 	policy := &structs.ACLPolicy{
@@ -33,7 +32,7 @@ func TestACLPolicyInfoCommand(t *testing.T) {
 		Rules: "node { policy = \"read\" }",
 	}
 	policy.SetHash()
-	assert.Nil(state.UpsertACLPolicies(structs.MsgTypeTestSetup, 1000, []*structs.ACLPolicy{policy}))
+	must.NoError(t, state.UpsertACLPolicies(structs.MsgTypeTestSetup, 1000, []*structs.ACLPolicy{policy}))
 
 	ui := cli.NewMockUi()
 	cmd := &ACLPolicyInfoCommand{Meta: Meta{Ui: ui, flagAddress: url}}
@@ -41,15 +40,13 @@ func TestACLPolicyInfoCommand(t *testing.T) {
 	// Attempt to apply a policy without a valid management token
 	invalidToken := mock.ACLToken()
 	code := cmd.Run([]string{"-address=" + url, "-token=" + invalidToken.SecretID, policy.Name})
-	assert.Equal(1, code)
+	must.One(t, code)
 
 	// Apply a policy with a valid management token
 	code = cmd.Run([]string{"-address=" + url, "-token=" + token.SecretID, policy.Name})
-	assert.Equal(0, code)
+	must.Zero(t, code)
 
 	// Check the output
 	out := ui.OutputWriter.String()
-	if !strings.Contains(out, policy.Name) {
-		t.Fatalf("bad: %v", out)
-	}
+	must.StrContains(t, out, policy.Name)
 }
diff --git a/command/acl_policy_list_test.go b/command/acl_policy_list_test.go
index ce3f2bcf7..e72749d2e 100644
--- a/command/acl_policy_list_test.go
+++ b/command/acl_policy_list_test.go
@@ -1,7 +1,6 @@
 package command
 
 import (
-	"strings"
 	"testing"
 
 	"github.com/hashicorp/nomad/acl"
@@ -10,23 +9,23 @@ import (
 	"github.com/hashicorp/nomad/nomad/mock"
 	"github.com/hashicorp/nomad/nomad/structs"
 	"github.com/mitchellh/cli"
-	"github.com/stretchr/testify/assert"
+	"github.com/shoenig/test/must"
 )
 
 func TestACLPolicyListCommand(t *testing.T) {
 	ci.Parallel(t)
-	assert := assert.New(t)
+
 	config := func(c *agent.Config) {
 		c.ACL.Enabled = true
 	}
 
 	srv, _, url := testServer(t, true, config)
 	state := srv.Agent.Server().State()
-	defer srv.Shutdown()
+	defer stopTestAgent(srv)
 
 	// Bootstrap an initial ACL token
 	token := srv.RootToken
-	assert.NotNil(token, "failed to bootstrap ACL token")
+	must.NotNil(t, token)
 
 	// Create a test ACLPolicy
 	policy := &structs.ACLPolicy{
@@ -34,7 +33,7 @@ func TestACLPolicyListCommand(t *testing.T) {
 		Rules: acl.PolicyWrite,
 	}
 	policy.SetHash()
-	assert.Nil(state.UpsertACLPolicies(structs.MsgTypeTestSetup, 1000, []*structs.ACLPolicy{policy}))
+	must.NoError(t, state.UpsertACLPolicies(structs.MsgTypeTestSetup, 1000, []*structs.ACLPolicy{policy}))
 
 	ui := cli.NewMockUi()
 	cmd := &ACLPolicyListCommand{Meta: Meta{Ui: ui, flagAddress: url}}
@@ -42,25 +41,20 @@ func TestACLPolicyListCommand(t *testing.T) {
 	// Attempt to list policies without a valid management token
 	invalidToken := mock.ACLToken()
 	code := cmd.Run([]string{"-address=" + url, "-token=" + invalidToken.SecretID})
-	assert.Equal(1, code)
+	must.One(t, code)
 
 	// Apply a policy with a valid management token
 	code = cmd.Run([]string{"-address=" + url, "-token=" + token.SecretID})
-	assert.Equal(0, code)
+	must.Zero(t, code)
 
 	// Check the output
 	out := ui.OutputWriter.String()
-	if !strings.Contains(out, policy.Name) {
-		t.Fatalf("bad: %v", out)
-	}
+	must.StrContains(t, out, policy.Name)
 
 	// List json
-	if code := cmd.Run([]string{"-address=" + url, "-token=" + token.SecretID, "-json"}); code != 0 {
-		t.Fatalf("expected exit 0, got: %d; %v", code, ui.ErrorWriter.String())
-	}
+	must.Zero(t, cmd.Run([]string{"-address=" + url, "-token=" + token.SecretID, "-json"}))
+
 	out = ui.OutputWriter.String()
-	if !strings.Contains(out, "CreateIndex") {
-		t.Fatalf("expected json output, got: %s", out)
-	}
+	must.StrContains(t, out, "CreateIndex")
 	ui.OutputWriter.Reset()
 }
diff --git a/command/acl_token_create_test.go b/command/acl_token_create_test.go
index e24e4c507..dd6507bd4 100644
--- a/command/acl_token_create_test.go
+++ b/command/acl_token_create_test.go
@@ -1,43 +1,40 @@
 package command
 
 import (
-	"strings"
 	"testing"
 
 	"github.com/hashicorp/nomad/ci"
 	"github.com/hashicorp/nomad/command/agent"
 	"github.com/mitchellh/cli"
-	"github.com/stretchr/testify/assert"
+	"github.com/shoenig/test/must"
 )
 
 func TestACLTokenCreateCommand(t *testing.T) {
 	ci.Parallel(t)
-	assert := assert.New(t)
+
 	config := func(c *agent.Config) {
 		c.ACL.Enabled = true
 	}
 
 	srv, _, url := testServer(t, true, config)
-	defer srv.Shutdown()
+	defer stopTestAgent(srv)
 
 	// Bootstrap an initial ACL token
 	token := srv.RootToken
-	assert.NotNil(token, "failed to bootstrap ACL token")
+	must.NotNil(t, token)
 
 	ui := cli.NewMockUi()
 	cmd := &ACLTokenCreateCommand{Meta: Meta{Ui: ui, flagAddress: url}}
 
 	// Request to create a new token without providing a valid management token
 	code := cmd.Run([]string{"-address=" + url, "-token=foo", "-policy=foo", "-type=client"})
-	assert.Equal(1, code)
+	must.One(t, code)
 
 	// Request to create a new token with a valid management token
 	code = cmd.Run([]string{"-address=" + url, "-token=" + token.SecretID, "-policy=foo", "-type=client"})
-	assert.Equal(0, code)
+	must.Zero(t, code)
 
 	// Check the output
 	out := ui.OutputWriter.String()
-	if !strings.Contains(out, "[foo]") {
-		t.Fatalf("bad: %v", out)
-	}
+	must.StrContains(t, out, "[foo]")
 }
diff --git a/command/acl_token_delete_test.go b/command/acl_token_delete_test.go
index 8da29208d..4c1a85e1a 100644
--- a/command/acl_token_delete_test.go
+++ b/command/acl_token_delete_test.go
@@ -2,7 +2,6 @@ package command
 
 import (
 	"fmt"
-	"strings"
 	"testing"
 
 	"github.com/hashicorp/nomad/acl"
@@ -11,22 +10,22 @@ import (
 	"github.com/hashicorp/nomad/nomad/mock"
 	"github.com/hashicorp/nomad/nomad/structs"
 	"github.com/mitchellh/cli"
-	"github.com/stretchr/testify/assert"
+	"github.com/shoenig/test/must"
 )
 
 func TestACLTokenDeleteCommand_ViaEnvVariable(t *testing.T) {
 	ci.Parallel(t)
-	assert := assert.New(t)
+
 	config := func(c *agent.Config) {
 		c.ACL.Enabled = true
 	}
 
 	srv, _, url := testServer(t, true, config)
-	defer srv.Shutdown()
+	defer stopTestAgent(srv)
 
 	// Bootstrap an initial ACL token
 	token := srv.RootToken
-	assert.NotNil(token, "failed to bootstrap ACL token")
+	must.NotNil(t, token)
 
 	ui := cli.NewMockUi()
 	cmd := &ACLTokenDeleteCommand{Meta: Meta{Ui: ui, flagAddress: url}}
@@ -36,21 +35,19 @@ func TestACLTokenDeleteCommand_ViaEnvVariable(t *testing.T) {
 	mockToken := mock.ACLToken()
 	mockToken.Policies = []string{acl.PolicyWrite}
 	mockToken.SetHash()
-	assert.Nil(state.UpsertACLTokens(structs.MsgTypeTestSetup, 1000, []*structs.ACLToken{mockToken}))
+	must.NoError(t, state.UpsertACLTokens(structs.MsgTypeTestSetup, 1000, []*structs.ACLToken{mockToken}))
 
 	// Attempt to delete a token without providing a valid token with delete
 	// permissions
 	code := cmd.Run([]string{"-address=" + url, "-token=foo", mockToken.AccessorID})
-	assert.Equal(1, code)
+	must.One(t, code)
 
 	// Delete a token using a valid management token set via an environment
 	// variable
 	code = cmd.Run([]string{"-address=" + url, "-token=" + token.SecretID, mockToken.AccessorID})
-	assert.Equal(0, code)
+	must.Zero(t, code)
 
 	// Check the output
 	out := ui.OutputWriter.String()
-	if !strings.Contains(out, fmt.Sprintf("Token %s successfully deleted", mockToken.AccessorID)) {
-		t.Fatalf("bad: %v", out)
-	}
+	must.StrContains(t, out, fmt.Sprintf("Token %s successfully deleted", mockToken.AccessorID))
 }
diff --git a/command/acl_token_info_test.go b/command/acl_token_info_test.go
index 23a6e15b9..3725393a3 100644
--- a/command/acl_token_info_test.go
+++ b/command/acl_token_info_test.go
@@ -1,35 +1,29 @@
 package command
 
 import (
-	"os"
-	"strings"
 	"testing"
 
 	"github.com/hashicorp/nomad/acl"
-	"github.com/hashicorp/nomad/ci"
 	"github.com/hashicorp/nomad/command/agent"
 	"github.com/hashicorp/nomad/nomad/mock"
 	"github.com/hashicorp/nomad/nomad/structs"
 	"github.com/mitchellh/cli"
-	"github.com/stretchr/testify/assert"
+	"github.com/shoenig/test/must"
 )
 
 func TestACLTokenInfoCommand_ViaEnvVar(t *testing.T) {
-	ci.Parallel(t)
-	defer os.Setenv("NOMAD_TOKEN", os.Getenv("NOMAD_TOKEN"))
-
-	assert := assert.New(t)
 	config := func(c *agent.Config) {
 		c.ACL.Enabled = true
 	}
 
 	srv, _, url := testServer(t, true, config)
-	defer srv.Shutdown()
+	defer stopTestAgent(srv)
+
 	state := srv.Agent.Server().State()
 
 	// Bootstrap an initial ACL token
 	token := srv.RootToken
-	assert.NotNil(token, "failed to bootstrap ACL token")
+	must.NotNil(t, token)
 
 	ui := cli.NewMockUi()
 	cmd := &ACLTokenInfoCommand{Meta: Meta{Ui: ui, flagAddress: url}}
@@ -38,28 +32,26 @@ func TestACLTokenInfoCommand_ViaEnvVar(t *testing.T) {
 	mockToken := mock.ACLToken()
 	mockToken.Policies = []string{acl.PolicyWrite}
 	mockToken.SetHash()
-	assert.Nil(state.UpsertACLTokens(structs.MsgTypeTestSetup, 1000, []*structs.ACLToken{mockToken}))
+	must.NoError(t, state.UpsertACLTokens(structs.MsgTypeTestSetup, 1000, []*structs.ACLToken{mockToken}))
 
 	// Attempt to fetch info on a token without providing a valid management
 	// token
 	invalidToken := mock.ACLToken()
-	os.Setenv("NOMAD_TOKEN", invalidToken.SecretID)
+	t.Setenv("NOMAD_TOKEN", invalidToken.SecretID)
 	code := cmd.Run([]string{"-address=" + url, mockToken.AccessorID})
-	assert.Equal(1, code)
+	must.One(t, code)
 
 	// Fetch info on a token with a valid management token
-	os.Setenv("NOMAD_TOKEN", token.SecretID)
+	t.Setenv("NOMAD_TOKEN", token.SecretID)
 	code = cmd.Run([]string{"-address=" + url, mockToken.AccessorID})
-	assert.Equal(0, code)
+	must.Zero(t, code)
 
 	// Fetch info on a token with a valid management token via a CLI option
-	os.Setenv("NOMAD_TOKEN", "")
+	t.Setenv("NOMAD_TOKEN", "")
 	code = cmd.Run([]string{"-address=" + url, "-token=" + token.SecretID, mockToken.AccessorID})
-	assert.Equal(0, code)
+	must.Zero(t, code)
 
 	// Check the output
 	out := ui.OutputWriter.String()
-	if !strings.Contains(out, mockToken.AccessorID) {
-		t.Fatalf("bad: %v", out)
-	}
+	must.StrContains(t, out, mockToken.AccessorID)
 }
diff --git a/command/acl_token_list_test.go b/command/acl_token_list_test.go
index 59622838b..8ea91674d 100644
--- a/command/acl_token_list_test.go
+++ b/command/acl_token_list_test.go
@@ -1,7 +1,6 @@
 package command
 
 import (
-	"strings"
 	"testing"
 
 	"github.com/hashicorp/nomad/acl"
@@ -10,29 +9,30 @@ import (
 	"github.com/hashicorp/nomad/nomad/mock"
 	"github.com/hashicorp/nomad/nomad/structs"
 	"github.com/mitchellh/cli"
-	"github.com/stretchr/testify/assert"
+	"github.com/shoenig/test/must"
 )
 
 func TestACLTokenListCommand(t *testing.T) {
 	ci.Parallel(t)
-	assert := assert.New(t)
+
 	config := func(c *agent.Config) {
 		c.ACL.Enabled = true
 	}
 
 	srv, _, url := testServer(t, true, config)
+	defer stopTestAgent(srv)
+
 	state := srv.Agent.Server().State()
-	defer srv.Shutdown()
 
 	// Bootstrap an initial ACL token
 	token := srv.RootToken
-	assert.NotNil(token, "failed to bootstrap ACL token")
+	must.NotNil(t, token)
 
 	// Create a valid token
 	mockToken := mock.ACLToken()
 	mockToken.Policies = []string{acl.PolicyWrite}
 	mockToken.SetHash()
-	assert.Nil(state.UpsertACLTokens(structs.MsgTypeTestSetup, 1000, []*structs.ACLToken{mockToken}))
+	must.NoError(t, state.UpsertACLTokens(structs.MsgTypeTestSetup, 1000, []*structs.ACLToken{mockToken}))
 
 	ui := cli.NewMockUi()
 	cmd := &ACLTokenListCommand{Meta: Meta{Ui: ui, flagAddress: url}}
@@ -40,25 +40,20 @@ func TestACLTokenListCommand(t *testing.T) {
 	// Attempt to list tokens without a valid management token
 	invalidToken := mock.ACLToken()
 	code := cmd.Run([]string{"-address=" + url, "-token=" + invalidToken.SecretID})
-	assert.Equal(1, code)
+	must.One(t, code)
 
 	// Apply a token with a valid management token
 	code = cmd.Run([]string{"-address=" + url, "-token=" + token.SecretID})
-	assert.Equal(0, code)
+	must.Zero(t, code)
 
 	// Check the output
 	out := ui.OutputWriter.String()
-	if !strings.Contains(out, mockToken.Name) {
-		t.Fatalf("bad: %v", out)
-	}
+	must.StrContains(t, out, mockToken.Name)
 
 	// List json
-	if code := cmd.Run([]string{"-address=" + url, "-token=" + token.SecretID, "-json"}); code != 0 {
-		t.Fatalf("expected exit 0, got: %d; %v", code, ui.ErrorWriter.String())
-	}
+	must.Zero(t, cmd.Run([]string{"-address=" + url, "-token=" + token.SecretID, "-json"}))
+
 	out = ui.OutputWriter.String()
-	if !strings.Contains(out, "CreateIndex") {
-		t.Fatalf("expected json output, got: %s", out)
-	}
+	must.StrContains(t, out, "CreateIndex")
 	ui.OutputWriter.Reset()
 }
diff --git a/command/acl_token_self_test.go b/command/acl_token_self_test.go
index d907cd781..e8eb49439 100644
--- a/command/acl_token_self_test.go
+++ b/command/acl_token_self_test.go
@@ -1,35 +1,29 @@
 package command
 
 import (
-	"os"
-	"strings"
 	"testing"
 
 	"github.com/hashicorp/nomad/acl"
-	"github.com/hashicorp/nomad/ci"
 	"github.com/hashicorp/nomad/command/agent"
 	"github.com/hashicorp/nomad/nomad/mock"
 	"github.com/hashicorp/nomad/nomad/structs"
 	"github.com/mitchellh/cli"
-	"github.com/stretchr/testify/assert"
+	"github.com/shoenig/test/must"
 )
 
 func TestACLTokenSelfCommand_ViaEnvVar(t *testing.T) {
-	ci.Parallel(t)
-	defer os.Setenv("NOMAD_TOKEN", os.Getenv("NOMAD_TOKEN"))
-
-	assert := assert.New(t)
 	config := func(c *agent.Config) {
 		c.ACL.Enabled = true
 	}
 
 	srv, _, url := testServer(t, true, config)
-	defer srv.Shutdown()
+	defer stopTestAgent(srv)
+
 	state := srv.Agent.Server().State()
 
 	// Bootstrap an initial ACL token
 	token := srv.RootToken
-	assert.NotNil(token, "failed to bootstrap ACL token")
+	must.NotNil(t, token)
 
 	ui := cli.NewMockUi()
 	cmd := &ACLTokenSelfCommand{Meta: Meta{Ui: ui, flagAddress: url}}
@@ -38,23 +32,21 @@ func TestACLTokenSelfCommand_ViaEnvVar(t *testing.T) {
 	mockToken := mock.ACLToken()
 	mockToken.Policies = []string{acl.PolicyWrite}
 	mockToken.SetHash()
-	assert.Nil(state.UpsertACLTokens(structs.MsgTypeTestSetup, 1000, []*structs.ACLToken{mockToken}))
+	must.NoError(t, state.UpsertACLTokens(structs.MsgTypeTestSetup, 1000, []*structs.ACLToken{mockToken}))
 
 	// Attempt to fetch info on a token without providing a valid management
 	// token
 	invalidToken := mock.ACLToken()
-	os.Setenv("NOMAD_TOKEN", invalidToken.SecretID)
+	t.Setenv("NOMAD_TOKEN", invalidToken.SecretID)
 	code := cmd.Run([]string{"-address=" + url})
-	assert.Equal(1, code)
+	must.One(t, code)
 
 	// Fetch info on a token with a valid token
-	os.Setenv("NOMAD_TOKEN", mockToken.SecretID)
+	t.Setenv("NOMAD_TOKEN", mockToken.SecretID)
 	code = cmd.Run([]string{"-address=" + url})
-	assert.Equal(0, code)
+	must.Zero(t, code)
 
 	// Check the output
 	out := ui.OutputWriter.String()
-	if !strings.Contains(out, mockToken.AccessorID) {
-		t.Fatalf("bad: %v", out)
-	}
+	must.StrContains(t, out, mockToken.AccessorID)
 }
diff --git a/command/acl_token_update_test.go b/command/acl_token_update_test.go
index e98002217..7ad70c8d7 100644
--- a/command/acl_token_update_test.go
+++ b/command/acl_token_update_test.go
@@ -9,23 +9,22 @@ import (
 	"github.com/hashicorp/nomad/nomad/mock"
 	"github.com/hashicorp/nomad/nomad/structs"
 	"github.com/mitchellh/cli"
-	"github.com/stretchr/testify/assert"
+	"github.com/shoenig/test/must"
 )
 
 func TestACLTokenUpdateCommand(t *testing.T) {
 	ci.Parallel(t)
 
-	assert := assert.New(t)
 	config := func(c *agent.Config) {
 		c.ACL.Enabled = true
 	}
 
 	srv, _, url := testServer(t, true, config)
-	defer srv.Shutdown()
+	defer stopTestAgent(srv)
 
 	// Bootstrap an initial ACL token
 	token := srv.RootToken
-	assert.NotNil(token, "failed to bootstrap ACL token")
+	must.NotNil(t, token)
 
 	ui := cli.NewMockUi()
 	cmd := &ACLTokenUpdateCommand{Meta: Meta{Ui: ui, flagAddress: url}}
@@ -35,19 +34,19 @@ func TestACLTokenUpdateCommand(t *testing.T) {
 	mockToken := mock.ACLToken()
 	mockToken.Policies = []string{acl.PolicyWrite}
 	mockToken.SetHash()
-	assert.Nil(state.UpsertACLTokens(structs.MsgTypeTestSetup, 1000, []*structs.ACLToken{mockToken}))
+	must.NoError(t, state.UpsertACLTokens(structs.MsgTypeTestSetup, 1000, []*structs.ACLToken{mockToken}))
 
 	// Request to update a new token without providing a valid management token
 	invalidToken := mock.ACLToken()
 	code := cmd.Run([]string{"--token=" + invalidToken.SecretID, "-address=" + url, "-name=bar", mockToken.AccessorID})
-	assert.Equal(1, code)
+	must.One(t, code)
 
 	// Request to update a new token with a valid management token
 	code = cmd.Run([]string{"--token=" + token.SecretID, "-address=" + url, "-name=bar", mockToken.AccessorID})
-	assert.Equal(0, code)
+	must.Zero(t, code)
 
 	// Check the output
 	out := ui.OutputWriter.String()
-	assert.Contains(out, mockToken.AccessorID)
-	assert.Contains(out, "bar")
+	must.StrContains(t, out, mockToken.AccessorID)
+	must.StrContains(t, out, "bar")
 }
diff --git a/command/alloc_status.go b/command/alloc_status.go
index 900d5fd91..dccb70388 100644
--- a/command/alloc_status.go
+++ b/command/alloc_status.go
@@ -8,7 +8,7 @@ import (
 	"strings"
 	"time"
 
-	humanize "github.com/dustin/go-humanize"
+	"github.com/dustin/go-humanize"
 	"github.com/posener/complete"
 
 	"github.com/hashicorp/nomad/api"
diff --git a/command/alloc_status_test.go b/command/alloc_status_test.go
index 8f80003c4..eb9bd69fe 100644
--- a/command/alloc_status_test.go
+++ b/command/alloc_status_test.go
@@ -12,10 +12,9 @@ import (
 	"github.com/hashicorp/nomad/helper/uuid"
 	"github.com/hashicorp/nomad/nomad/mock"
 	"github.com/hashicorp/nomad/nomad/structs"
-	"github.com/hashicorp/nomad/testutil"
 	"github.com/mitchellh/cli"
 	"github.com/posener/complete"
-	"github.com/stretchr/testify/assert"
+	"github.com/shoenig/test/must"
 	"github.com/stretchr/testify/require"
 )
 
@@ -27,7 +26,7 @@ func TestAllocStatusCommand_Implements(t *testing.T) {
 func TestAllocStatusCommand_Fails(t *testing.T) {
 	ci.Parallel(t)
 	srv, _, url := testServer(t, false, nil)
-	defer srv.Shutdown()
+	defer stopTestAgent(srv)
 
 	ui := cli.NewMockUi()
 	cmd := &AllocStatusCommand{Meta: Meta{Ui: ui}}
@@ -89,23 +88,9 @@ func TestAllocStatusCommand_Fails(t *testing.T) {
 func TestAllocStatusCommand_LifecycleInfo(t *testing.T) {
 	ci.Parallel(t)
 	srv, client, url := testServer(t, true, nil)
-	defer srv.Shutdown()
+	defer stopTestAgent(srv)
 
-	// Wait for a node to be ready
-	testutil.WaitForResult(func() (bool, error) {
-		nodes, _, err := client.Nodes().List(nil)
-		if err != nil {
-			return false, err
-		}
-		for _, node := range nodes {
-			if node.Status == structs.NodeStatusReady {
-				return true, nil
-			}
-		}
-		return false, fmt.Errorf("no ready nodes")
-	}, func(err error) {
-		require.NoError(t, err)
-	})
+	waitForNodes(t, client)
 
 	ui := cli.NewMockUi()
 	cmd := &AllocStatusCommand{Meta: Meta{Ui: ui}}
@@ -152,24 +137,9 @@ func TestAllocStatusCommand_LifecycleInfo(t *testing.T) {
 func TestAllocStatusCommand_Run(t *testing.T) {
 	ci.Parallel(t)
 	srv, client, url := testServer(t, true, nil)
-	defer srv.Shutdown()
+	defer stopTestAgent(srv)
 
-	// Wait for a node to be ready
-	testutil.WaitForResult(func() (bool, error) {
-		nodes, _, err := client.Nodes().List(nil)
-		if err != nil {
-			return false, err
-		}
-		for _, node := range nodes {
-			if _, ok := node.Drivers["mock_driver"]; ok &&
-				node.Status == structs.NodeStatusReady {
-				return true, nil
-			}
-		}
-		return false, fmt.Errorf("no ready nodes")
-	}, func(err error) {
-		t.Fatalf("err: %v", err)
-	})
+	waitForNodes(t, client)
 
 	ui := cli.NewMockUi()
 	cmd := &AllocStatusCommand{Meta: Meta{Ui: ui}}
@@ -249,28 +219,13 @@ func TestAllocStatusCommand_Run(t *testing.T) {
 func TestAllocStatusCommand_RescheduleInfo(t *testing.T) {
 	ci.Parallel(t)
 	srv, client, url := testServer(t, true, nil)
-	defer srv.Shutdown()
+	defer stopTestAgent(srv)
 
-	// Wait for a node to be ready
-	testutil.WaitForResult(func() (bool, error) {
-		nodes, _, err := client.Nodes().List(nil)
-		if err != nil {
-			return false, err
-		}
-		for _, node := range nodes {
-			if node.Status == structs.NodeStatusReady {
-				return true, nil
-			}
-		}
-		return false, fmt.Errorf("no ready nodes")
-	}, func(err error) {
-		t.Fatalf("err: %v", err)
-	})
+	waitForNodes(t, client)
 
 	ui := cli.NewMockUi()
 	cmd := &AllocStatusCommand{Meta: Meta{Ui: ui}}
 	// Test reschedule attempt info
-	require := require.New(t)
 	state := srv.Agent.Server().State()
 	a := mock.Alloc()
 	a.Metrics = &structs.AllocMetric{}
@@ -285,41 +240,27 @@ func TestAllocStatusCommand_RescheduleInfo(t *testing.T) {
 			},
 		},
 	}
-	require.Nil(state.UpsertAllocs(structs.MsgTypeTestSetup, 1000, []*structs.Allocation{a}))
+	require.NoError(t, state.UpsertAllocs(structs.MsgTypeTestSetup, 1000, []*structs.Allocation{a}))
 
 	if code := cmd.Run([]string{"-address=" + url, a.ID}); code != 0 {
 		t.Fatalf("expected exit 0, got: %d", code)
 	}
 	out := ui.OutputWriter.String()
-	require.Contains(out, "Replacement Alloc ID")
-	require.Regexp(regexp.MustCompile(".*Reschedule Attempts\\s*=\\s*1/2"), out)
+	require.Contains(t, out, "Replacement Alloc ID")
+	require.Regexp(t, regexp.MustCompile(".*Reschedule Attempts\\s*=\\s*1/2"), out)
 }
 
 func TestAllocStatusCommand_ScoreMetrics(t *testing.T) {
 	ci.Parallel(t)
 	srv, client, url := testServer(t, true, nil)
-	defer srv.Shutdown()
+	defer stopTestAgent(srv)
 
-	// Wait for a node to be ready
-	testutil.WaitForResult(func() (bool, error) {
-		nodes, _, err := client.Nodes().List(nil)
-		if err != nil {
-			return false, err
-		}
-		for _, node := range nodes {
-			if node.Status == structs.NodeStatusReady {
-				return true, nil
-			}
-		}
-		return false, fmt.Errorf("no ready nodes")
-	}, func(err error) {
-		t.Fatalf("err: %v", err)
-	})
+	waitForNodes(t, client)
 
 	ui := cli.NewMockUi()
 	cmd := &AllocStatusCommand{Meta: Meta{Ui: ui}}
+
 	// Test node metrics
-	require := require.New(t)
 	state := srv.Agent.Server().State()
 	a := mock.Alloc()
 	mockNode1 := mock.Node()
@@ -342,27 +283,26 @@ func TestAllocStatusCommand_ScoreMetrics(t *testing.T) {
 			},
 		},
 	}
-	require.Nil(state.UpsertAllocs(structs.MsgTypeTestSetup, 1000, []*structs.Allocation{a}))
+	require.NoError(t, state.UpsertAllocs(structs.MsgTypeTestSetup, 1000, []*structs.Allocation{a}))
 
 	if code := cmd.Run([]string{"-address=" + url, "-verbose", a.ID}); code != 0 {
 		t.Fatalf("expected exit 0, got: %d", code)
 	}
 	out := ui.OutputWriter.String()
-	require.Contains(out, "Placement Metrics")
-	require.Contains(out, mockNode1.ID)
-	require.Contains(out, mockNode2.ID)
+	require.Contains(t, out, "Placement Metrics")
+	require.Contains(t, out, mockNode1.ID)
+	require.Contains(t, out, mockNode2.ID)
 
 	// assert we sort headers alphabetically
-	require.Contains(out, "binpack  node-affinity")
-	require.Contains(out, "final score")
+	require.Contains(t, out, "binpack  node-affinity")
+	require.Contains(t, out, "final score")
 }
 
 func TestAllocStatusCommand_AutocompleteArgs(t *testing.T) {
 	ci.Parallel(t)
-	assert := assert.New(t)
 
 	srv, _, url := testServer(t, true, nil)
-	defer srv.Shutdown()
+	defer stopTestAgent(srv)
 
 	ui := cli.NewMockUi()
 	cmd := &AllocStatusCommand{Meta: Meta{Ui: ui, flagAddress: url}}
@@ -370,15 +310,15 @@ func TestAllocStatusCommand_AutocompleteArgs(t *testing.T) {
 	// Create a fake alloc
 	state := srv.Agent.Server().State()
 	a := mock.Alloc()
-	assert.Nil(state.UpsertAllocs(structs.MsgTypeTestSetup, 1000, []*structs.Allocation{a}))
+	must.NoError(t, state.UpsertAllocs(structs.MsgTypeTestSetup, 1000, []*structs.Allocation{a}))
 
 	prefix := a.ID[:5]
 	args := complete.Args{Last: prefix}
 	predictor := cmd.AutocompleteArgs()
 
 	res := predictor.Predict(args)
-	assert.Equal(1, len(res))
-	assert.Equal(a.ID, res[0])
+	must.Len(t, 1, res)
+	must.Eq(t, a.ID, res[0])
 }
 
 func TestAllocStatusCommand_HostVolumes(t *testing.T) {
@@ -397,7 +337,8 @@ func TestAllocStatusCommand_HostVolumes(t *testing.T) {
 			},
 		}
 	})
-	defer srv.Shutdown()
+	defer stopTestAgent(srv)
+
 	state := srv.Agent.Server().State()
 
 	// Upsert the job and alloc
@@ -448,7 +389,8 @@ func TestAllocStatusCommand_HostVolumes(t *testing.T) {
 func TestAllocStatusCommand_CSIVolumes(t *testing.T) {
 	ci.Parallel(t)
 	srv, _, url := testServer(t, true, nil)
-	defer srv.Shutdown()
+	defer stopTestAgent(srv)
+
 	state := srv.Agent.Server().State()
 
 	// Upsert the node, plugin, and volume
diff --git a/command/helpers_test.go b/command/helpers_test.go
index f1e3b7436..187afee22 100644
--- a/command/helpers_test.go
+++ b/command/helpers_test.go
@@ -347,7 +347,7 @@ job "example" {
 }
 `
 
-	setEnv(t, "NOMAD_VAR_var4", "from-envvar")
+	t.Setenv("NOMAD_VAR_var4", "from-envvar")
 
 	cliArgs := []string{`var2=from-cli`}
 	fileVars := `var3 = "from-varfile"`
diff --git a/command/meta_test.go b/command/meta_test.go
index 27724faaa..1bb4d076b 100644
--- a/command/meta_test.go
+++ b/command/meta_test.go
@@ -92,7 +92,7 @@ func TestMeta_Colorize(t *testing.T) {
 		{
 			Name: "disable colors via env var",
 			SetupFn: func(t *testing.T, m *Meta) {
-				setEnv(t, EnvNomadCLINoColor, "1")
+				t.Setenv(EnvNomadCLINoColor, "1")
 				m.SetupUi([]string{})
 			},
 			ExpectColor: false,
@@ -107,7 +107,7 @@ func TestMeta_Colorize(t *testing.T) {
 		{
 			Name: "force colors via env var",
 			SetupFn: func(t *testing.T, m *Meta) {
-				setEnv(t, EnvNomadCLIForceColor, "1")
+				t.Setenv(EnvNomadCLIForceColor, "1")
 				m.SetupUi([]string{})
 			},
 			ExpectColor: true,
@@ -122,7 +122,7 @@ func TestMeta_Colorize(t *testing.T) {
 		{
 			Name: "no color take predecence over force color via env var",
 			SetupFn: func(t *testing.T, m *Meta) {
-				setEnv(t, EnvNomadCLINoColor, "1")
+				t.Setenv(EnvNomadCLINoColor, "1")
 				m.SetupUi([]string{"-force-color"})
 			},
 			ExpectColor: false,
@@ -141,8 +141,8 @@ func TestMeta_Colorize(t *testing.T) {
 			os.Stdout = tty
 
 			// Make sure color related environment variables are clean.
-			setEnv(t, EnvNomadCLIForceColor, "")
-			setEnv(t, EnvNomadCLINoColor, "")
+			t.Setenv(EnvNomadCLIForceColor, "")
+			t.Setenv(EnvNomadCLINoColor, "")
 
 			// Run test case.
 			m := &Meta{}
diff --git a/command/util_test.go b/command/testing_test.go
similarity index 77%
rename from command/util_test.go
rename to command/testing_test.go
index 03efe6bb0..59283c46f 100644
--- a/command/util_test.go
+++ b/command/testing_test.go
@@ -1,13 +1,16 @@
 package command
 
 import (
+	"fmt"
 	"os"
 	"testing"
 
 	"github.com/hashicorp/nomad/api"
 	"github.com/hashicorp/nomad/command/agent"
 	"github.com/hashicorp/nomad/helper/pointer"
+	"github.com/hashicorp/nomad/nomad/structs"
 	"github.com/hashicorp/nomad/testutil"
+	"github.com/shoenig/test/must"
 )
 
 func testServer(t *testing.T, runClient bool, cb func(*agent.Config)) (*agent.TestAgent, *api.Client, string) {
@@ -108,16 +111,33 @@ func testMultiRegionJob(jobID, region, datacenter string) *api.Job {
 	return job
 }
 
-// setEnv wraps os.Setenv(key, value) and restores the environment variable to initial value in test cleanup
-func setEnv(t *testing.T, key, value string) {
-	initial, ok := os.LookupEnv(key)
-	os.Setenv(key, value)
-
-	t.Cleanup(func() {
-		if ok {
-			os.Setenv(key, initial)
-		} else {
-			os.Unsetenv(key)
+func waitForNodes(t *testing.T, client *api.Client) {
+	testutil.WaitForResult(func() (bool, error) {
+		nodes, _, err := client.Nodes().List(nil)
+		if err != nil {
+			return false, err
 		}
+		for _, node := range nodes {
+			if _, ok := node.Drivers["mock_driver"]; ok &&
+				node.Status == structs.NodeStatusReady {
+				return true, nil
+			}
+		}
+		return false, fmt.Errorf("no ready nodes")
+	}, func(err error) {
+		must.NoError(t, err)
 	})
 }
+
+func stopTestAgent(a *agent.TestAgent) {
+	_ = a.Shutdown()
+}
+
+func getTempFile(t *testing.T, name string) (string, func()) {
+	f, err := os.CreateTemp("", name)
+	must.NoError(t, err)
+	must.NoError(t, f.Close())
+	return f.Name(), func() {
+		_ = os.Remove(f.Name())
+	}
+}
diff --git a/command/ui_test.go b/command/ui_test.go
index 8b0d049ae..621131703 100644
--- a/command/ui_test.go
+++ b/command/ui_test.go
@@ -43,38 +43,38 @@ func TestCommand_Ui(t *testing.T) {
 		{
 			Name: "set namespace via env var",
 			SetupFn: func(t *testing.T) {
-				setEnv(t, "NOMAD_NAMESPACE", "dev")
+				t.Setenv("NOMAD_NAMESPACE", "dev")
 			},
 			ExpectedURL: "http://127.0.0.1:4646?namespace=dev",
 		},
 		{
 			Name: "set region via env var",
 			SetupFn: func(t *testing.T) {
-				setEnv(t, "NOMAD_REGION", "earth")
+				t.Setenv("NOMAD_REGION", "earth")
 			},
 			ExpectedURL: "http://127.0.0.1:4646?region=earth",
 		},
 		{
 			Name: "set region and namespace via env var",
 			SetupFn: func(t *testing.T) {
-				setEnv(t, "NOMAD_REGION", "earth")
-				setEnv(t, "NOMAD_NAMESPACE", "dev")
+				t.Setenv("NOMAD_REGION", "earth")
+				t.Setenv("NOMAD_NAMESPACE", "dev")
 			},
 			ExpectedURL: "http://127.0.0.1:4646?namespace=dev&region=earth",
 		},
 		{
 			Name: "set region and namespace via env var",
 			SetupFn: func(t *testing.T) {
-				setEnv(t, "NOMAD_REGION", "earth")
-				setEnv(t, "NOMAD_NAMESPACE", "dev")
+				t.Setenv("NOMAD_REGION", "earth")
+				t.Setenv("NOMAD_NAMESPACE", "dev")
 			},
 			ExpectedURL: "http://127.0.0.1:4646?namespace=dev&region=earth",
 		},
 		{
 			Name: "flags have higher precedence",
 			SetupFn: func(t *testing.T) {
-				setEnv(t, "NOMAD_REGION", "earth")
-				setEnv(t, "NOMAD_NAMESPACE", "dev")
+				t.Setenv("NOMAD_REGION", "earth")
+				t.Setenv("NOMAD_NAMESPACE", "dev")
 			},
 			Args: []string{
 				"-region=mars",
@@ -87,8 +87,8 @@ func TestCommand_Ui(t *testing.T) {
 	for _, tc := range cases {
 		t.Run(tc.Name, func(t *testing.T) {
 			// Make sure environment variables are clean.
-			setEnv(t, "NOMAD_NAMESPACE", "")
-			setEnv(t, "NOMAD_REGION", "")
+			t.Setenv("NOMAD_NAMESPACE", "")
+			t.Setenv("NOMAD_REGION", "")
 
 			// Setup fake CLI UI and test case
 			ui := cli.NewMockUi()
diff --git a/go.mod b/go.mod
index 11e6c6f7c..574e2f88f 100644
--- a/go.mod
+++ b/go.mod
@@ -109,7 +109,7 @@ require (
 	github.com/ryanuber/go-glob v1.0.0
 	github.com/sean-/seed v0.0.0-20170313163322-e2103e2c3529
 	github.com/shirou/gopsutil/v3 v3.21.12
-	github.com/shoenig/test v0.3.0
+	github.com/shoenig/test v0.3.1
 	github.com/skratchdot/open-golang v0.0.0-20160302144031-75fb7ed4208c
 	github.com/stretchr/testify v1.8.0
 	github.com/syndtr/gocapability v0.0.0-20200815063812-42c35b437635
diff --git a/go.sum b/go.sum
index c3a1dbfeb..b4e20a417 100644
--- a/go.sum
+++ b/go.sum
@@ -1199,8 +1199,8 @@ github.com/shirou/gopsutil v0.0.0-20181107111621-48177ef5f880/go.mod h1:5b4v6he4
 github.com/shirou/gopsutil/v3 v3.21.12 h1:VoGxEW2hpmz0Vt3wUvHIl9fquzYLNpVpgNNB7pGJimA=
 github.com/shirou/gopsutil/v3 v3.21.12/go.mod h1:BToYZVTlSVlfazpDDYFnsVZLaoRG+g8ufT6fPQLdJzA=
 github.com/shirou/w32 v0.0.0-20160930032740-bb4de0191aa4/go.mod h1:qsXQc7+bwAM3Q1u/4XEfrquwF8Lw7D7y5cD8CuHnfIc=
-github.com/shoenig/test v0.3.0 h1:H6tfSvgLrPHRR5NH9S40+lOfoyeH2PbswBr4twgn9Po=
-github.com/shoenig/test v0.3.0/go.mod h1:xYtyGBC5Q3kzCNyJg/SjgNpfAa2kvmgA0i5+lQso8x0=
+github.com/shoenig/test v0.3.1 h1:dhGZztS6nQuvJ0o0RtUiQHaEO4hhArh/WmWwik3Ols0=
+github.com/shoenig/test v0.3.1/go.mod h1:xYtyGBC5Q3kzCNyJg/SjgNpfAa2kvmgA0i5+lQso8x0=
 github.com/shopspring/decimal v1.2.0 h1:abSATXmQEYyShuxI4/vyW3tV1MrKAJzCZ/0zLUXYbsQ=
 github.com/shopspring/decimal v1.2.0/go.mod h1:DKyhrW/HYNuLGql+MJL6WCR6knT2jwCFRcu2hWCYk4o=
 github.com/shurcooL/sanitized_anchor_name v1.0.0/go.mod h1:1NzhyTcUVG4SuEtjjoZeVRXNmyL/1OwPU0+IJeTBvfc=