git » chasquid » commit 3584441

diff --git a/internal/sts/sts_test.go b/internal/sts/sts_test.go
index 80e24ba..7ff4570 100644
--- a/internal/sts/sts_test.go
+++ b/internal/sts/sts_test.go
@@ -252,6 +252,8 @@ func expvarMustEq(t *testing.T, name string, v *expvar.Int, expected int64) {
 
 func TestCacheBasics(t *testing.T) {
 	dir := testlib.MustTempDir(t)
+	defer testlib.RemoveIfOk(t, dir)
+
 	c, err := NewCache(dir)
 	if err != nil {
 		t.Fatal(err)
@@ -308,15 +310,13 @@ func TestCacheBasics(t *testing.T) {
 	expvarMustEq(t, "cacheFetches", cacheFetches, 4)
 	expvarMustEq(t, "cacheHits", cacheHits, 1)
 	expvarMustEq(t, "cacheFailedFetch", cacheFailedFetch, 1)
-
-	if !t.Failed() {
-		os.RemoveAll(dir)
-	}
 }
 
 // Test how the cache behaves when the files are corrupt.
 func TestCacheBadData(t *testing.T) {
 	dir := testlib.MustTempDir(t)
+	defer testlib.RemoveIfOk(t, dir)
+
 	c, err := NewCache(dir)
 	if err != nil {
 		t.Fatal(err)
@@ -379,10 +379,6 @@ func TestCacheBadData(t *testing.T) {
 
 	expvarMustEq(t, "cacheUnmarshalErrors", cacheUnmarshalErrors, 1)
 	expvarMustEq(t, "cacheInvalid", cacheInvalid, 1)
-
-	if !t.Failed() {
-		os.RemoveAll(dir)
-	}
 }
 
 func (c *PolicyCache) mustFetch(ctx context.Context, t *testing.T, d string) *Policy {
@@ -408,6 +404,8 @@ func mustRewriteAndChtime(t *testing.T, fname, content string) {
 
 func TestCacheRefresh(t *testing.T) {
 	dir := testlib.MustTempDir(t)
+	defer testlib.RemoveIfOk(t, dir)
+
 	c, err := NewCache(dir)
 	if err != nil {
 		t.Fatal(err)
@@ -452,10 +450,6 @@ func TestCacheRefresh(t *testing.T) {
 	if p.MaxAge != 200*time.Second {
 		t.Fatalf("policy.MaxAge is %v, expected 200s", p.MaxAge)
 	}
-
-	if !t.Failed() {
-		os.RemoveAll(dir)
-	}
 }
 
 func TestCacheSlashSafe(t *testing.T) {
diff --git a/internal/userdb/userdb_test.go b/internal/userdb/userdb_test.go
index 37f2c64..f4b176a 100644
--- a/internal/userdb/userdb_test.go
+++ b/internal/userdb/userdb_test.go
@@ -5,12 +5,19 @@ import (
 	"io/ioutil"
 	"os"
 	"reflect"
+	"strings"
 	"testing"
 )
 
 // Remove the file if the test was successful. Used in defer statements, to
 // leave files around for inspection when the tests failed.
 func removeIfSuccessful(t *testing.T, fname string) {
+	// Safeguard, to make sure we only remove test files.
+	// This should help prevent accidental deletions.
+	if !strings.Contains(fname, "userdb_test") {
+		panic("invalid/dangerous directory")
+	}
+
 	if !t.Failed() {
 		os.Remove(fname)
 	}